Anatomically Guided Latent Diffusion for Brain MRI Progression Modeling

TL;DR

The paper tackles the challenge of modeling longitudinal brain MRI progression with anatomical fidelity, addressing limitations of prior multi-stage models. It introduces Anatomically Guided Latent Diffusion Model (AG-LDM), a two-stage latent diffusion framework that fuses baseline anatomy, follow-up latent, and clinical covariates at the input and employs WarpSeg for anatomical supervision. AG-LDM achieves state-of-the-art image quality and reduces regional volumetric errors by about 15-20% on ADNI, with up to 31.5x higher sensitivity to covariates and strong zero-shot generalization to OASIS-3, plus plausible counterfactual trajectories. The approach offers an efficient, anatomically grounded tool for reliable brain MRI progression modeling and could enable hypothesis testing and personalized prognosis in neurodegenerative diseases.

Abstract

Accurately modeling longitudinal brain MRI progression is crucial for understanding neurodegenerative diseases and predicting individualized structural changes. Existing state-of-the-art approaches, such as Brain Latent Progression (BrLP), often use multi-stage training pipelines with auxiliary conditioning modules but suffer from architectural complexity, suboptimal use of conditional clinical covariates, and limited guarantees of anatomical consistency. We propose Anatomically Guided Latent Diffusion Model (AG-LDM), a segmentation-guided framework that enforces anatomically consistent progression while substantially simplifying the training pipeline. AG-LDM conditions latent diffusion by directly fusing baseline anatomy, noisy follow-up states, and clinical covariates at the input level, a strategy that avoids auxiliary control networks by learning a unified, end-to-end model that represents both anatomy and progression. A lightweight 3D tissue segmentation model (WarpSeg) provides explicit anatomical supervision during both autoencoder fine-tuning and diffusion model training, ensuring consistent brain tissue boundaries and morphometric fidelity. Experiments on 31,713 ADNI longitudinal pairs and zero-shot evaluation on OASIS-3 demonstrate that AG-LDM matches or surpasses more complex diffusion models, achieving state-of-the-art image quality and 15-20\% reduction in volumetric errors in generated images. AG-LDM also exhibits markedly stronger utilization of temporal and clinical covariates (up to 31.5x higher sensitivity than BrLP) and generates biologically plausible counterfactual trajectories, accurately capturing hallmarks of Alzheimer's progression such as limbic atrophy and ventricular expansion. These results highlight AG-LDM as an efficient, anatomically grounded framework for reliable brain MRI progression modeling.

Figures (5)

• Figure 1: AG-LDM is a two-stage brain MRI generative model: Stage 1 first fine-tunes a Variational Autoencoder (VAE) to learn latent representations of MRIs, whereas Stage 2 trains a Latent Diffusion Model (LDM) that generates the latent of the future MRI conditioned on the latent of the baseline MRI and other clinical covariates (e.g., sex, age and diagnosis information at time A and time B). In both stages, segmentation guidance is used to ensure the reconstructed MRI has consistent GM and WM tissue segmentation as the ground-truth.
• Figure 2: Effect of segmentation-guided supervision: GM and WM segmentation of an MRI synthesized by AG-LDM Base, AG-LDM AE-Seg, and AG-LDM Full (conditioned on the same baseline MRI); Adding segmentation guidance at both stages produces segmentations closer to the ground truth.
• Figure 3: Ten-year brain progression synthesis (73$\rightarrow$83 years) for a randomly selected AD patient: Columns correspond to the ground truth follow-up scan and scans generated by AG-LDM, BrLP, and CounterSynth. For each method we show matched sagittal/coronal/axial slices with SynthSeg overlays (right image in each pair).
• Figure 4: Anatomical consistency measured by Dice coefficient for GM, WM, and CSF on the validation set. Each box represents 8 AG-LDM variants trained with different hyperparameter configurations.
• Figure 5: Comparing Real and Synthesized Trajectories: Top row: comparing observed volume change among 200 real CN subjects with trajectories synthesized by AG-LDM and BrLP. Bottom row: comparing observed volume change among 200 real subjects converting from CN to AD with trajectories counterfactually synthesized by AG-LDM and BrLP. Each tile shows the relative volume change ($\Delta V_{\text{rel}}$) in an ROI between the baseline and (real or synthesized) follow-up MRIs.