- The paper introduces CLIMB, which combines a Mamba-based latent diffusion model with a Gaussian-aligned autoencoder to synthesize longitudinal brain MRI scans.
- It replaces conventional self-attention with state space modeling, improving computational efficiency and anatomical fidelity as shown by high SSIM and low MSE.
- The method conditions on multimodal patient data to enable personalized predictions, advancing applications in precision medicine and in silico clinical trials.
CLIMB: A Mamba-based Latent Diffusion Model for Controllable Longitudinal Brain Image Synthesis
Technical Overview
CLIMB addresses the challenge of synthesizing longitudinal brain MRI scans, a critical task for disease progression modeling and precision medicine. The method combines a latent diffusion framework with a Gaussian-aligned autoencoder (GATE) and leverages the Mamba selective state space model as an architectural backbone. Unlike existing LDMs, which predominantly use transformer-derived self-attention mechanisms for contextual feature modeling, CLIMB replaces this with state space modeling, resulting in significant computational efficiency improvements without degrading generation quality.
The framework is conditioned on multi-faceted patient data, including baseline MRI, demographic variables (age, gender), clinical markers (disease status, genetic factors), and structural brain volumes. This extensive conditioning supports controllability and personalization in generated outputs.
Autoencoder and Latent Space Alignment
Rather than adopting a VAE with inherent noise injection—which typically results in over-smooth reconstructions—CLIMB's GATE employs sliced CDF alignment to enforce Gaussianity in the latent codes without sampling stochasticity. This is achieved with a deterministic CDF matching strategy over multiple random projections of the latent space. The approach produces latent encodings that are structurally compatible with the latent diffusion process, ensuring high-fidelity reconstructions and improved perceptual similarity.
Multi-loss training combines reconstruction, sliced density alignment, perceptual, and adversarial objectives, further enhancing anatomical and visual realism.
Diffusion Process with State Space Modeling
The denoising network within the latent diffusion pipeline adopts a U-Net-like architecture with the critical substitution of Mamba state space modules for standard self-attention. This allows the model to efficiently capture both global and local temporal dependencies in high-dimensional latent representations, crucial for tracking anatomical changes over long periods.
Conditioning is performed via cross-attention to integrate projected patient data and baseline image features. A two-stage training pipeline is used: initial training on projected variables alone, followed by joint finetuning with image features, stabilizing convergence and promoting robust latent manifold adaptation.
Inference
During inference, CLIMB encodes the baseline MRI using GATE and predicts future conditional variables (via a pretrained IRLSTM) such as projected disease state and brain volumes. These are input to the Mamba-based denoising network, which samples a future latent vector. The decoder reconstructs the future MRI scan. Efficient sampling is achieved using a DDIM-based reverse process with only 25 steps and latent averaging over 10 samples per run.
Quantitative and Qualitative Results
CLIMB was rigorously evaluated on 6,306 MRI scans from 1,390 participants in the ADNI dataset. Extensive preprocessing and segmentation were used to extract relevant region-of-interest volumes.
Key Numerical Results
CLIMB achieves strong improvements over the SADM and BrLP baselines:
- SSIM: 0.9433 (highest among all compared methods)
- MSE: 2.01e-3 (lowest, indicating highly accurate reconstruction)
- PSNR: 27.82
- LPIPS: 0.0587 (best perceptual score)
- Inference time: 2.92 seconds per image
- Memory usage: 5699 MB
Ablation studies show that GATE outperforms traditional VAEs on all core metrics; Mamba-based state space blocks yield comparable or better fidelity with a substantial reduction in computation compared to self-attention; model scaling experiments further elucidate the memory-quality tradeoff.
Qualitative Analysis
Generated MRI sequences demonstrate anatomical consistency across temporal horizons, with substantially reduced structural artifacts compared to prior art. CLIMB's outputs show lower reconstruction errors in central brain regions and preserve long-range temporal consistency, as evidenced by direct visual comparisons and deviation heatmaps.
Predicted conditional variables (disease status, brain structure volumes) also show higher predictive accuracy, supporting more reliable downstream synthesis.
Implications and Future Perspectives
Practically, CLIMB's approach enables highly efficient and controllable longitudinal brain imaging simulation, a significant step toward in silico clinical trial design, patient-specific disease trajectory modeling, and augmentation of incomplete clinical imaging datasets. The ability to forecast future MRI scans given multi-modal, high-dimensional patient data with high anatomical fidelity has direct relevance for early disease detection, prognosis, and therapy planning.
Theoretically, the replacement of self-attention with Mamba-type SSMs in high-dimensional image generative modeling demonstrates that state space approaches can match or exceed transformer-based backbones in both efficiency and generation quality for temporal medical imaging. The decoupling of Gaussian alignment from stochastic VAE sampling with a sliced CDF strategy constitutes a notable methodological advance for latent generative models.
Regulatory and clinical integration will require further validation by domain experts, as high SSIM/MSE benchmarks do not guarantee clinical interpretability or diagnostic utility.
Anticipated future directions include extending CLIMB to multi-modal imaging, further scaling of SSM backbones, semi-supervised or self-supervised adaptation to new populations, and real-time interactive simulation tools for clinical use.
Conclusion
CLIMB provides an advanced and computationally efficient paradigm for controllable, longitudinal brain image synthesis by integrating a Gaussian-aligned autoencoder with a Mamba-based latent diffusion model. Demonstrating state-of-the-art quantitative and perceptual performance on a large-scale longitudinal neuroimaging dataset, the approach substantially advances the modeling of temporal anatomical change. The system's efficiency and flexibility position it as a promising tool for both research and future clinical applications, though real-world deployment will require further cross-disciplinary validation.