MambaControl: Anatomy Graph-Enhanced Mamba ControlNet with Fourier Refinement for Diffusion-Based Disease Trajectory Prediction
Hao Yang, Tao Tan, Shuai Tan, Weiqin Yang, Kunyan Cai, Calvin Chen, Yue Sun
TL;DR
The paper addresses the challenge of modeling disease progression in medical imaging by simultaneously capturing long-range temporal dynamics and maintaining anatomical fidelity. It introduces MambaControl, a framework that fuses a selective state-space model (Mamba) with diffusion-based trajectory generation and graph-guided anatomical priors, augmented by Fourier-enhanced spectral graphs and ControlNet-style conditioning. The method achieves state-of-the-art Alzheimer’s disease trajectory prediction on the ADNI3 dataset, delivering superior PSNR and SSIM scores and improved region-wise MAE with significantly fewer parameters than competing approaches. Ablation studies confirm the value of the spatial and Fourier graphs and the graph-guided integration, underscoring the approach’s potential for reliable, personalized prognosis and clinical decision support.
Abstract
Modelling disease progression in precision medicine requires capturing complex spatio-temporal dynamics while preserving anatomical integrity. Existing methods often struggle with longitudinal dependencies and structural consistency in progressive disorders. To address these limitations, we introduce MambaControl, a novel framework that integrates selective state-space modelling with diffusion processes for high-fidelity prediction of medical image trajectories. To better capture subtle structural changes over time while maintaining anatomical consistency, MambaControl combines Mamba-based long-range modelling with graph-guided anatomical control to more effectively represent anatomical correlations. Furthermore, we introduce Fourier-enhanced spectral graph representations to capture spatial coherence and multiscale detail, enabling MambaControl to achieve state-of-the-art performance in Alzheimer's disease prediction. Quantitative and regional evaluations demonstrate improved progression prediction quality and anatomical fidelity, highlighting its potential for personalised prognosis and clinical decision support.
