Semi-Supervised Diffusion Model for Brain Age Prediction

TL;DR

Brain age prediction in neurodegenerative contexts often suffers from data quality and rapid disease progression. The authors introduce a semi-supervised diffusion model built as a diffusion autoencoder conditioned on a semantic latent, with an age predictor mapping $\mathbf{z}_{\text{sem}}$ to age, achieving a test correlation of $r=0.83$ ($p<0.01$) and $R^2=0.65$ with $MAE=5$ years on clinical-grade data. Importantly, brain-PADs derived from this model significantly correlate with ALS survival ($r=0.24$, $p<0.05$), suggesting the approach captures aging-related neuroanatomy linked to outcomes. The method is competitive with non-generative baselines and demonstrates that diffusion-based representations can be robust to data quality, with potential clinical utility in prognosis and trial design.

Abstract

Brain age prediction models have succeeded in predicting clinical outcomes in neurodegenerative diseases, but can struggle with tasks involving faster progressing diseases and low quality data. To enhance their performance, we employ a semi-supervised diffusion model, obtaining a 0.83(p<0.01) correlation between chronological and predicted age on low quality T1w MR images. This was competitive with state-of-the-art non-generative methods. Furthermore, the predictions produced by our model were significantly associated with survival length (r=0.24, p<0.05) in Amyotrophic Lateral Sclerosis. Thus, our approach demonstrates the value of diffusion-based architectures for the task of brain age prediction.

Figures (4)

• Figure 1: The Model architecture. First the image, $\mathbf{x}_{0}$ goes through the forward process $q(\mathbf{x}_{1:T}|\mathbf{x}_{0})$, to produce $\mathbf{x}_{T}$. The reverse process, $p_{\theta}(\mathbf{x_{0:T}}|\mathbf{z}_{\text{sem}})$ then reconstructs the image conditioned on the semantic latent, produced by $s_{\phi}(\mathbf{x}_{0})$. Finally, the age predictor, $f_{\psi}(\mathbf{z}_{\text{sem}})$ predicts an individual's age, $\hat{\mathbf{y}}$.
• Figure 2: A Boxplot of the distribution of brain-PADs produced by different models on our ALS dataset. Our model produces more variation within the brain-PADs of the ALS patients. Two-Sample Kolmogrov Smirnov tests (Table \ref{['Table2']}) demonstrated that our brain-PADs were significantly different from brain-PADs from the other three approaches .
• Figure 3: Age Interpolations. In the first row, the image in the first column is that of a 25-year-old, we then linearly interpolate between them and a 75-year-old (final column). The second row is the same but reversed, from age 92 to age 20. The difference in data quality, across the two rows, exemplifies the robustness of the model to varied data quality.
• Figure 4: Image Reconstructions. The images on the left columns are the original image, and the images on the right columns are their reconstructed version.