Covariate-Conditioned One-Step Diffusion for Multi-Horizon 3D Brain MRI Prediction in Aging and Neurodegeneration

Abstract

Predicting longitudinal brain changes from magnetic resonance imaging (MRI) can help characterize structural progression associated with aging and neurodegeneration. However, diffusion-based forecasting typically requires iterative denoising, limiting inference efficiency. We propose a covariate-conditioned one-step diffusion framework for multi-horizon prediction of 3D gray-matter density (GMD) maps. Rather than directly generating future anatomy, the model predicts longitudinal GMD changes through residual prediction in a single denoising step. The predicted changes are added to the baseline GMD map to reconstruct future anatomy. The model is conditioned on baseline anatomy, age, diagnosis, APOE4 allele count, and continuous follow-up time. Training combines mean squared error with a Pearson correlation loss to improve the spatial agreement between predicted and observed GMD changes. A unified model supports prediction horizons from 2 to 10 years. On the held-out ADNI test set, our method achieves a $\Delta$-Pearson of 0.478, mean squared error (MSE) of $5.176\times 10^{-3}$, and structural similarity index measure (SSIM) of 0.979, with $\Delta$-Pearson reaching 0.582 at 10 years. Without retraining, external evaluation on AIBL yields a $\Delta$-Pearson of 0.450, MSE of $4.016\times 10^{-3}$, and SSIM of 0.985. Ablation studies further demonstrate the benefits of residual prediction, covariate conditioning, and correlation-based supervision. Code: https://github.com/natenaile/OneStepBrainMRI.

arXiv comment5 pages, 2 figures, 4 tables. Submitted to IEEE ISBI 2027