Technical Summary: AnaDiffusion

Problem Statement

High-fidelity 3D brain MRI generation is critical for medical imaging, simulation, and anatomical analysis, yet existing generative models face significant limitations. Current state-of-the-art approaches, including GANs and monolithic Latent Diffusion Models (LDMs), typically synthesize 3D volumes as a single, entangled unit. This "monolithic" strategy often overlooks regional anatomical structures, leading to localized failures such as blurred tissue interfaces, topological inconsistencies in the cortex, and unnatural asymmetries between hemispheres. Furthermore, methods that rely on segmentation masks for conditioning require dense, subject-specific annotations at inference time, limiting their practical utility. The core challenge is to generate globally coherent brain volumes that simultaneously preserve fine-grained, region-specific anatomical details and allow for controllable editing without requiring external segmentation maps during generation.

Methodology: AnaDiffusion

The authors propose AnaDiffusion, an anatomically compositional latent diffusion framework that decomposes the generation process into distinct, anatomically meaningful regions before assembling them into a whole. The framework operates through a "part-to-whole" strategy involving three main stages:

1. Anatomical Part Factorization

The brain is decomposed into three coarse anatomical parts based on neuroanatomical organization:

• Left Hemisphere ($Hemi_L$)
• Right Hemisphere ($Hemi_R$)
• Cerebellar-Brainstem Complex ($CB$)

Cerebrospinal fluid (CSF) is excluded from these parts as it is a fluid compartment rather than a distinct structure. During training, part volumes are extracted using SynthSeg-generated masks and cropped to fixed bounding boxes. Crucially, these masks are used only offline to construct training data; they are not provided as conditioning inputs to the diffusion models during inference.

2. Local Anatomical Priors (Part Models)

The framework trains separate latent diffusion models for each anatomical part to capture local structural priors:

• Hemisphere Canonicalization: To address bilateral symmetry while allowing for side-specific variations, a single hemisphere LDM is trained on both left hemispheres and right hemispheres that have been flipped into a left-oriented canonical frame. This model is conditioned on a side indicator ($y \in \{0, 1\}$) to distinguish left from right.
• Part-Specific Autoencoders: Each part (hemispheres and CB) has its own autoencoder and diffusion prior, allowing the model to devote generative capacity to the specific morphological complexities of each region.

3. Part-to-Whole Latent Scaffold and Refinement

The core innovation lies in the integration of these parts into a global volume without requiring subject-specific segmentation at inference:

• Assembly: Generated part assets are assembled in image space into a "scaffold" ($x_{scaf}$), omitting CSF and background.
• Latent Injection: This scaffold is re-encoded into the whole-brain latent space ($z_{scaf}^0$) using a whole-brain encoder.
• Refinement Training: The whole-brain denoiser is trained to resume reverse sampling from an injected latent state. At a specific timestep $t_{inj}$, the scaffold latent is injected into the global latent trajectory ($z_{inj}$) using a mask $m$:
  $$z_{inj} = (1 - m) \odot z_{tgt} + m \odot z_{scaf}$$
  The whole-brain denoiser then continues the denoising process to harmonize the injected parts with the surrounding context (e.g., synthesizing CSF and smoothing boundaries).
• Inference Modes:
  • Unconditional Generation: The model generates parts, assembles them, and refines the whole.
  • Controllable Part Editing: Specific part assets can be replaced with alternative generated parts. The modified set is reassembled, re-encoded, injected, and refined, allowing for localized edits while preserving the non-target anatomy.

Key Contributions

1. Compositional Framework: AnaDiffusion introduces a framework that factorizes 3D brain MRI synthesis into anatomical parts and a whole-brain latent integration process, moving beyond monolithic volume generation.
2. Part-to-Whole Refinement: The method proposes a latent refinement mechanism that integrates independently generated parts into a coherent whole-brain volume. This supports controllable part replacement while maintaining global consistency and avoiding the need for dense segmentation maps at inference.
3. Symmetry-Aware Hemisphere Generator: A hemisphere generator is designed with canonicalized weight sharing and side-indicator conditioning to encourage anatomically consistent left-right generation while improving data efficiency.
4. Inference-Time Control: The system enables localized part replacement and compositional refinement without requiring subject-specific dense segmentation maps during the inference phase.

Experimental Results

The model was evaluated on a subject-disjoint split of the Alzheimer's Disease Neuroimaging Initiative (ADNI) dataset (1,735 scans from 407 subjects).

• Distributional Fidelity (FID): AnaDiffusion achieved the lowest Fréchet Inception Distance (FID) across the whole brain, left and right hemispheres, cerebellar-brainstem complex, and seam regions compared to seven baselines (including VAE-GAN, HA-GAN, standard LDMs, ControlNet, and MorphLDM).
• Anatomical Plausibility (Cohen's $|d|$): The method achieved the best absolute Cohen's $d$ values for the cerebellum and the second-best for ventricles and brainstem, indicating superior alignment with real anatomical volume distributions.
• Controllable Editing: In localized editing experiments, AnaDiffusion demonstrated significantly higher target transfer (preservation of the inserted part) compared to segmentation-conditioned baselines, while maintaining high off-target preservation (stability of non-edited regions).
• Ablation Studies:
  • Removing latent injection significantly degraded performance, confirming the necessity of the scaffold for local fidelity.
  • Using separate models for left and right hemispheres showed a trade-off: better left-hemisphere FID but worse right-hemisphere and whole-brain FID, suggesting that the shared canonicalization strategy offers better sample efficiency and regularization.
  • The injection timing ($r_{inj}$) was critical; a moderate number of remaining denoising steps allowed for optimal seam harmonization without weakening the injected structure.

Significance and Claims

The paper claims that AnaDiffusion bridges the gap between local anatomical controllability and global brain integrity. By explicitly generating anatomical part assets and refining their assembled scaffold, the model transforms the output of a brain generative model from a static volume into a coherent sample with reusable and editable components.

The authors emphasize that this approach improves part-sensitive distributional quality and interface coherence over strong baselines. A key practical significance is the ability to perform controllable part editing and generate high-fidelity volumes without requiring subject-specific dense segmentation maps at inference, addressing a major bottleneck in clinical applications where such annotations are often unavailable. The work suggests that treating brain synthesis as a compositional problem, rather than a monolithic one, is essential for capturing the spatial heterogeneity and structural dependencies of the human brain.

Limitations Acknowledged: The authors note that the multi-stage design increases training and inference complexity. The current fixed anatomical factorization does not fully solve broader tissue-level calibration and assumes approximate correspondence with the MNI152 template, which may not hold for severe pathologies with mass effects or displaced boundaries. Future work is proposed to explore adaptive decompositions and region-dependent refinement schedules.