TADM: Temporally-Aware Diffusion Model for Neurodegenerative Progression on Brain MRI

TL;DR

This paper tackles the challenge of forecasting neurodegenerative progression in brain MRIs by introducing TADM, a temporally aware diffusion model that learns the distribution of intensity changes between baseline and follow-up scans. It conditions diffusion on the baseline image representation, the time interval via an age-gap parameter, and patient-specific data, and uses a Brain-Age Estimator to enforce temporal alignment through an auxiliary loss. The approach focuses on predicting residual changes rather than full images, yielding improved region-wise accuracy and overall similarity metrics on the OASIS-3 dataset, outperforming state-of-the-art diffusion-based methods. The method has potential clinical utility for predicting patient outcomes and guiding treatment planning, with future work aimed at extending to 3D data and other imaging modalities for broader applicability.

Abstract

Generating realistic images to accurately predict changes in the structure of brain MRI is a crucial tool for clinicians. Such applications help assess patients' outcomes and analyze how diseases progress at the individual level. However, existing methods for this task present some limitations. Some approaches attempt to model the distribution of MRI scans directly by conditioning the model on patients' ages, but they fail to explicitly capture the relationship between structural changes in the brain and time intervals, especially on age-unbalanced datasets. Other approaches simply rely on interpolation between scans, which limits their clinical application as they do not predict future MRIs. To address these challenges, we propose a Temporally-Aware Diffusion Model (TADM), which introduces a novel approach to accurately infer progression in brain MRIs. TADM learns the distribution of structural changes in terms of intensity differences between scans and combines the prediction of these changes with the initial baseline scans to generate future MRIs. Furthermore, during training, we propose to leverage a pre-trained Brain-Age Estimator (BAE) to refine the model's training process, enhancing its ability to produce accurate MRIs that match the expected age gap between baseline and generated scans. Our assessment, conducted on the OASIS-3 dataset, uses similarity metrics and region sizes computed by comparing predicted and real follow-up scans on 3 relevant brain regions. TADM achieves large improvements over existing approaches, with an average decrease of 24% in region size error and an improvement of 4% in similarity metrics. These evaluations demonstrate the improvement of our model in mimicking temporal brain neurodegenerative progression compared to existing methods. Our approach will benefit applications, such as predicting patient outcomes or improving treatments for patients.

Figures (2)

• Figure 1: TADM comprises three main parts (surrounded by coloured dashed boxes). red: The DDPM takes a residual image, calculated as the difference between two scans acquired at two the time points $T_a$ and $T_b$, on which random noise $\epsilon$ is applied. The DDPM is trained to denoise the residual image through several diffusion steps. The denoised residual image $\widehat{I}_{\Delta_{a,b}}$ and the scan $I_{T_a}$ are summed together to estimate the scan $\widehat{I}_{T_b}$ at time $T_{b}$. blue: Here, we encode the scan $I_{T_a}$ to extract a representation $z_a$ used to condition the DDPM, in conjunction with other patient-specific data. green: The estimated $\widehat{I}_{T_b}$ is provided to the encoder to extract the features $z_b$ that, together with the previously extracted features $z_a$, are provided as inputs to a BAE to predict the time interval $\widehat{\Delta}_{a,b}$. The padlock indicates a model with frozen parameters.
• Figure 2: Comparison of the temporal progression on a 66-year-old subject with AD, obtained by our approach against other state-of-the-art methods. We show predicted slice-MRIs on the left and the corresponding error with the subject's real brain MRI on the right.