Table of Contents
Fetching ...

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.

MambaControl: Anatomy Graph-Enhanced Mamba ControlNet with Fourier Refinement for Diffusion-Based Disease Trajectory Prediction

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.
Paper Structure (17 sections, 12 equations, 3 figures, 1 table)

This paper contains 17 sections, 12 equations, 3 figures, 1 table.

Figures (3)

  • Figure 1: Comparison of ControlNet architectures for guiding diffusion models. (a) shows the baseline U-Net-based ControlNet. (b) shows a direct replacement using Mamba blocks. (c) shows our proposed MambaControl method, which utilizes separate Mamba pathways for control and diffusion, embedding graph-processed anatomical features from the control pathway to guide the main diffusion pathway generation.
  • Figure 2: Framework of our proposed MambaControl. It consists of two Mamba-based pathways: Diffusion and Control. The Control pathway incorporates a graph convolutional network in its downsampling phase to extract anatomical and spatial features. These features are then integrated into the Diffusion pathway's upsampling stage, enabling precise control over the generation process.
  • Figure 3: Qualitative evaluation of disease progression models. Each row represents a time point in longitudinal order (top to bottom). For each method, brain MRI predictions (left) are shown with corresponding residual maps (right) showing deviations from Ground Truth.