Source-linked AI summary
AnaDiffusion: Anatomically CompositionalLatent Diffusion for Controllable 3D Brain MRI Generation
Huiwen Han, Lulin Liu, Bangya Liu, Yuanhao Cai, Nuo Chen, Xiaoqing Wang, Ziqian Xie, Chenyu You, Shuiwang Ji, Degui Zhi, Zhiwen Fan
TL;DR
Existing 3D brain MRI generators often overlook fine-grained regional anatomy and limit local control. AnaDiffusion composes anatomically generated parts into a whole-brain latent diffusion process, outperforming compared baselines across evaluated synthesis measures while supporting localized replacement.
Problem
Existing 3D brain MRI generators synthesize volumes monolithically, providing limited preservation of fine-grained local anatomy and controllability.
Method
AnaDiffusion trains part-specific latent diffusion models and integrates their assembled anatomical scaffold into ongoing whole-brain denoising for coherent, controllable synthesis.
Results
AnaDiffusion improves evaluated synthesis metrics on subject-disjoint ADNI test subjects and enables localized part replacement with high target transfer and limited off-target drift.
Takeaways & Limitations
AnaDiffusion bridges local anatomical controllability and global brain integrity through explicit part assets and whole-brain refinement.
Takeaways & Limitations
The multi-stage design adds training and inference complexity, while fixed predefined regions may miss other biological organizations and pathology with displaced anatomical boundaries.
Abstract
from arXiv · showhide
3D brain MRI generation has made significant advances in medical imaging, simulation, and controllable anatomical analysis. However, existing generative models typically synthesize 3D volumes monolithically, often overlooking regional anatomical structures and limiting local controllability. To address these limitations, we introduce AnaDiffusion, an anatomically compositional latent diffusion framework that factorizes the generation process into distinct, anatomically meaningful regions, followed by part-to-whole assembly and global refinement. Our approach first trains part diffusion models to capture local structural priors. We then inject an assembled anatomical composite of the parts into the whole-brain latent representation and continue denoising. This mechanism enables the model to resolve global context while preserving the injected anatomy. As a result, AnaDiffusion produces both explicit part assets and a globally coherent volume, thereby enabling controllable part editing without requiring subject-specific dense segmentation maps at inference time while maintaining consistent part-to-whole brain structure. On the subject-disjoint ADNI test split, AnaDiffusion achieves the lowest FID across the whole brain, left and right hemispheres, cerebellar-brainstem complex, and seam regions. It also achieves the best cerebellar and second-best ventricular and brainstem absolute Cohen's d values among the evaluated methods. In localized editing experiments, paired MS-SSIM demonstrates high target transfer and off-target preservation, supporting controllable part replacement with minimal unintended anatomical alterations.
1 Introduction
AnaDiffusion addresses limitations of monolithic 3D brain MRI synthesis by decomposing generation into anatomical parts and integrating them into a refined, coherent whole-brain volume. The framework enables controllable part replacement while preserving non-edited regions and improving regional and global anatomical fidelity on held-out ADNI subjects.
- Method: Part-specific latent diffusion models generate local anatomical assets, which are assembled and injected into an ongoing whole-brain generative trajectory for global refinement.This coupling is designed to preserve high-fidelity local anatomy while maintaining global coherence.
- Motivation and contribution: AnaDiffusion introduces a compositional latent diffusion framework that models anatomical parts alongside whole-brain latent integration for controllable 3D brain MRI generation.The approach targets heterogeneous brain structures more effectively than monolithic whole-volume generation.
- Controllable editing: Part-to-whole latent refinement supports controllable part replacement while keeping non-edited brain regions stable and maintaining consistent part-to-whole anatomy.The model produces explicit part-level assets that can be replaced at inference time.
- Method: A shared hemisphere generator uses left-right canonicalization and side-indicator conditioning to model the brain’s bilateral symmetry.The design is motivated by the functional and structural bilateral symmetry of the human brain.
- Evaluation: On held-out subjects from a subject-disjoint ADNI test split, AnaDiffusion improves evaluated synthesis metrics, regional distributional fidelity, and segmentation-based anatomical alignment while enabling localized replacement without dense subject-specific segmentation maps.Localized editing shows high target transfer and limited off-target drift.
2 Related Work
Brain MRI generation has progressed from slice-wise or 2.5D methods toward full 3D synthesis, while recent approaches add structural or morphological priors. AnaDiffusion extends this direction toward explicit part-level controllability through part-specific priors and whole-brain latent integration.
- 2D-to-3D Brain MRI Generation: Recent work has shifted from 2D or 2.5D brain MRI generation toward full 3D synthesis with improved volumetric consistency.
- 2D-to-3D Brain MRI Generation: GAN-based models including α-GAN, CCE-GAN, and HA-GAN are widely explored for 3D brain MRI generation.
- Structural and Morphological Priors: Recent methods incorporate explicit structural priors by using segmentation maps as conditioning signals or auxiliary supervision for anatomically localized synthesis.
- Structural and Morphological Priors: Morphology-aware methods impose geometric constraints through deformation fields or surface-based representations, including template deformation and conditioned geodesic transformations.
- Gap in Existing Methods: These approaches regularize anatomical localization, tissue boundaries, or global morphology but do not explicitly decompose brain generation into part-specific components.
- AnaDiffusion: AnaDiffusion addresses this gap by combining part-specific generative priors with whole-brain latent inpainting for controllable editing without subject-specific dense segmentation maps at inference.
3 Methods
AnaDiffusion factorizes 3D brain MRI generation into anatomical part and whole-brain latent diffusion models, using assembled part scaffolds for whole-brain refinement. The framework supports unconditional compositional generation and localized part editing while preserving global anatomical context.
- Framework overview: AnaDiffusion trains separate part and whole-brain latent diffusion models, then uses an assemble-then-refine strategy that optimizes only the whole-brain denoiser.Frozen part models generate anatomical assets that are assembled in image space, re-encoded, and injected into the whole-brain latent trajectory.
- Anatomical decomposition: The anatomical decomposition contains three parts: left hemisphere, right hemisphere, and cerebellar-brainstem complex, while CSF, background, and uncertain boundaries are omitted from assembly.CSF is excluded because it is a fluid compartment with unstable, morphology-sensitive boundaries; a deterministic operator places parts in template-aligned locations.
- Anatomical decomposition: Hemisphere canonicalization flips right-hemisphere crops into a left-oriented frame and trains one hemisphere LDM with a side indicator for statistical sharing without enforcing exact symmetry.The method maps right crops using a flip operator before combining them with left crops for training.
- Scaffold injection: Part assets are composed in image space and re-encoded with the whole-brain encoder, avoiding direct pasting between incompatible part and whole-brain latent spaces.Generated regional self-masks define the scaffold support, which is forward-noised to the whole-brain trajectory before reverse denoising.
- Inference and editing: During inference, the framework refines globally contextualized crops or replaces selected part assets, then reassembles, reinjects, and refines them to preserve non-target anatomy and repair seams.The refinement budget r controls the preservation-harmonization tradeoff: smaller r favors inserted-part preservation, whereas moderately larger r supports interface harmonization and missing-context synthesis.
4 Data Preprocessing and Anatomical Factorization
The study uses preprocessed T1-weighted brain MRIs from ADNI and factorizes anatomy into three task-driven components: the left hemisphere, right hemisphere, and cerebellar-brainstem complex.
- Dataset and preprocessing: ADNI contributes 1,735 T1-weighted brain MRI scans from 407 subjects spanning cognitively normal, mild cognitive impairment, and Alzheimer’s disease groups.Subjects are 55–93 years old, with a mean age of 76.7 and 47.4% female subjects.
- Dataset and preprocessing: Each scan undergoes N4 bias-field correction, skull stripping, and affine registration to MNI152 space.
- Anatomical factorization: The anatomical factorization decomposes each brain into the left hemisphere (HemiL), right hemisphere (HemiR), and cerebellar-brainstem complex (CB).This coarse, task-driven decomposition follows widely used neuroanatomical organization in neuroimaging and neurobiology.
5 Experiments
Experiments show that AnaDiffusion improves whole-brain and regional generation quality while supporting controllable localized editing. Ablations further show benefits from latent injection and an intermediate injection time, with trade-offs for separate hemisphere models.
- Generation quality: AnaDiffusion achieves the lowest mean MedicalNet FID across the whole brain, both hemispheres, cerebellar–brainstem complex, and seam regions.The seam-region improvement is attributed to whole-brain refinement harmonizing independently generated parts.
- Morphometric evaluation: 0.169 [0.007, 0.482] is AnaDiffusion’s lowest cerebellar absolute Cohen’s |d| mismatch on the SynthSeg-derived metric.The reported metric evaluates structures tied to part-to-whole integration.
- Ablation studies: Disabling latent injection raises whole-brain mean FID from 36.16 [19.16, 59.61] to 52.27 [30.69, 79.63] and seam FID from 0.29 [0.17, 0.46] to 0.55 [0.35, 0.81].The comparison is between direct part-latent replacement (α = 1) and no injection (α = 0).
- Ablation studies: rinj = 10 yields the lowest mean FID across all five evaluated regions, balancing local part anchoring against global harmonization.The tested settings are rinj ∈ {7, 10, 15}, approximately corresponding to tinj ∈ {120, 180, 280}.
- Localized editing: AnaDiffusion achieves higher target transfer than Segm. cLDM for the left hemisphere (0.9427 vs. 0.7952), right hemisphere (0.9484 vs. 0.8343), and cerebellar–brainstem complex (0.9618 vs. 0.9151).Paired MS-SSIM evaluates donor-part transfer and preservation of non-edited anatomy using matched donor-recipient editing pairs.
6 Conclusion
AnaDiffusion is presented as a part-to-whole latent diffusion framework for controllable 3D brain MRI generation that connects local anatomical controllability with global brain integrity. Its limitations include added complexity, incomplete tissue-level calibration, and reliance on predefined anatomical regions, motivating more unified and adaptive compositional strategies.
- Conclusion: AnaDiffusion frames 3D brain MRI synthesis as part-to-whole generation, bridging local anatomical controllability and global brain integrity.The framework is motivated by distinct local anatomical distributions, structured interfaces, and region-specific control needs.
- Limitations: The multi-stage design introduces additional training and inference complexity compared with a monolithic LDM.This is identified as a remaining limitation despite AnaDiffusion’s advantages.
- Limitations: Fixed anatomy-driven factorization improves local part and seam behavior but does not fully solve broader tissue-level calibration.The current decomposition also focuses on predefined anatomical regions and may miss other biologically meaningful organizations, such as tissue classes.
- Future work: Future work will explore unified and adaptive compositional generation using anatomical or tissue-type experts, learned decompositions, and region-dependent refinement schedules.These directions aim to allocate modeling capacity according to local anatomical complexity and improve coordination between part-level generation and global structure.
Appendix · A Technical Appendices and Supplementary Material · A.1 Training and Inference
AnaDiffusion uses MONAI-based AutoencoderKL and diffusion UNet architectures implemented in PyTorch with CUDA and trained using AdamW. Training specifies region-dependent resolutions and autoencoder optimization settings, including learning rates, warm-starting, and adversarial supervision.
- A.1 Training and Inference: AnaDiffusion adopts AutoencoderKL and diffusion UNet architectures from the MONAI Generative Models library.The models are built on MONAI.
- A.1 Training and Inference: All models are implemented in PyTorch with CUDA and trained using the AdamW optimizer.
- A.1 Training and Inference: 128 × 128 × 128 is the final generation resolution for the whole brain.
- A.1 Training and Inference: 64 × 128 × 128 is the final generation resolution for each hemisphere.
- A.1 Training and Inference: 128 × 96 × 64 is the final generation resolution for the cerebellar-brainstem complex.
- A.1 Training and Inference: 1 × 10−4 and 5 × 10−5 are the autoencoder learning rates for the whole-brain and part models, respectively.Training is warm-started for 20 epochs before adversarial supervision with a patch-based discriminator; its learning rate is half the corresponding autoencoder rate.
A.2 Parts Factorization and Assembly Details
AnaDiffusion factorizes registered brain MRIs into three fixed anatomical parts using SynthSeg labels, deterministic template-aligned crops, and dilated template masks. CSF is excluded from part masks and synthesized during whole-brain refinement.
- Part Definition: Structural MRIs are registered to MNI152 space, then SynthSeg labels define three fixed parts: left hemisphere, right hemisphere, and cerebellar-brainstem complex.These parts provide the anatomical units for compositional generation.
- Part Definition: SynthSeg label sets assign anatomical regions to each part, while CSF label 24 is excluded from all part masks and synthesized during whole-brain refinement.The left hemisphere uses labels {2, 3, 4, 5, 10, 11, 12, 13, 17, 18, 26, 28}; the right uses {41, 42, 43, 44, 49, 50, 51, 52, 53, 54, 58, 60}; the cerebellar-brainstem complex uses {7, 8, 46, 47, 14, 15, 16}.
- Part Cropping: Within the 128 × 128 × 128 whole-brain grid, hemisphere crops measure 64 × 128 × 128, while the cerebellar-brainstem crop measures 128 × 96 × 64.Hemisphere crops split the volume along the brain midline; the cerebellar-brainstem crop comes from the inferior posterior registered brain.
- Part Masks: Template masks merge SynthSeg labels for the three parts and are dilated by r = 2 voxels, corresponding to 3 mm at 1.5 mm isotropic resolution.The dilation provides an anatomical margin and accommodates residual inter-subject variation after registration.