# Spatio-Temporal Conditional Diffusion Models for Forecasting Future Multiple Sclerosis Lesion Masks Conditioned on Treatments

arXiv: https://arxiv.org/abs/2508.07006 · HTML: https://mlsanity.com/paper/2508.07006/

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Placed in Medical imaging and segmentation, one of 14 directions in Diffusion models.
- Builds on 16 of its 26 references that are in the map.
- First posted on arXiv 9 Aug 2025.

## Placed in

- Diffusion models / [Medical imaging and segmentation](https://mlsanity.com/diffusion/medical-imaging-and-segmentation/)

## Builds on (most cited first, 16 of 16)

- [U-Net: Convolutional Networks for Biomedical Image Segmentation](https://mlsanity.com/paper/1505.04597/) (2015)
- [Diffusion Models Beat GANs on Image Synthesis](https://mlsanity.com/paper/2105.05233/) (2021)
- [Classifier-Free Diffusion Guidance](https://mlsanity.com/paper/2207.12598/) (2022)
- [Adding Conditional Control to Text-to-Image Diffusion Models](https://mlsanity.com/paper/2302.05543/) (2023)
- [Variational Diffusion Models](https://mlsanity.com/paper/2107.00630/) (2021)
- [Simple diffusion: End-to-end diffusion for high resolution images](https://mlsanity.com/paper/2301.11093/) (2023)
- [Efficient Diffusion Training via Min-SNR Weighting Strategy](https://mlsanity.com/paper/2303.09556/) (2023)
- [SynthSeg: Segmentation of brain MRI scans of any contrast and resolution without retraining](https://mlsanity.com/paper/2107.09559/) (2021)
- [Adaptive Latent Diffusion Model for 3D Medical Image to Image Translation: Multi-modal Magnetic Resonance Imaging Study](https://mlsanity.com/paper/2311.00265/) (2023)
- [SADM: Sequence-Aware Diffusion Model for Longitudinal Medical Image Generation](https://mlsanity.com/paper/2212.08228/) (2022)
- [WDM: 3D Wavelet Diffusion Models for High-Resolution Medical Image Synthesis](https://mlsanity.com/paper/2402.19043/) (2024)
- [3D MedDiffusion: A 3D Medical Latent Diffusion Model for Controllable and High-quality Medical Image Generation](https://mlsanity.com/paper/2412.13059/) (2024)
- [Enhancing Spatiotemporal Disease Progression Models via Latent Diffusion and Prior Knowledge](https://mlsanity.com/paper/2405.03328/) (2024)
- [Conditional Diffusion Models are Medical Image Classifiers that Provide Explainability and Uncertainty for Free](https://mlsanity.com/paper/2502.03687/) (2025)
- [Personalized Prediction of Future Lesion Activity and Treatment Effect in Multiple Sclerosis from Baseline MRI](https://mlsanity.com/paper/2204.01702/) (2022)
- [Improving Image-Based Precision Medicine with Uncertainty-Aware Causal Models](https://mlsanity.com/paper/2305.03829/) (2023)

## Abstract (arXiv)

Image-based personalized medicine has the potential to transform healthcare, particularly for diseases that exhibit heterogeneous progression such as Multiple Sclerosis (MS). In this work, we introduce the first treatment-aware spatio-temporal diffusion model that is able to generate future masks demonstrating lesion evolution in MS. Our voxel-space approach incorporates multi-modal patient data, including MRI and treatment information, to forecast new and enlarging T2 (NET2) lesion masks at a future time point. Extensive experiments on a multi-centre dataset of 2131 patient 3D MRIs from randomized clinical trials for relapsing-remitting MS demonstrate that our generative model is able to accurately predict NET2 lesion masks for patients across six different treatments. Moreover, we demonstrate our model has the potential for real-world clinical applications through downstream tasks such as future lesion count and location estimation, binary lesion activity classification, and generating counterfactual future NET2 masks for several treatments with different efficacies. This work highlights the potential of causal, image-based generative models as powerful tools for advancing data-driven prognostics in MS.