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

2025-08-09

AI Summary

This study addresses the challenge of predicting the spatial distribution of new and enlarging T2 (NET2) lesions in multiple sclerosis (MS) patients to inform personalized treatment decisions. We propose the first treatment-aware spatiotemporal diffusion model that jointly models 3D MRI sequences and structured treatment information, enabling voxel-level, causally interpretable lesion evolution generation. The model supports counterfactual reasoning and comparative evaluation across multiple therapeutic regimens, generating future NET2 lesion masks while performing clinically relevant tasks—including lesion counting, spatial localization, and activity classification. Evaluated on 2,131 multicenter clinical trial cases, our method significantly outperforms existing approaches and demonstrates robust generalization across six major disease-modifying therapies (DMTs). This work establishes a novel paradigm for imaging-driven MS prognosis modeling.

Technology Category

• Computer Vision: Diffusion Models for Vision
• Machine Learning: Multimodal Learning
• Data Mining & Knowledge Management: Mining of Spatial, Temporal or Spatio-Temporal Data

Application Category

• User Modeling, Personalization and Recommendation: Attacks and countermeasures in recommendation systems
• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking
• Web Mining and Content Analysis: Large pretrained models with web data

Abstract

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.

Problem

• Predict future MS lesion masks using treatment-aware diffusion models
• Incorporate multi-modal data to forecast new and enlarging lesions
• Enable clinical applications like lesion count and location estimation

Innovation

• Treatment-aware spatio-temporal diffusion model
• Multi-modal patient data integration
• Causal image-based generative forecasting