Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Probabilistic Temporal Prediction of Continuous Disease Trajectories and Treatment Effects Using Neural SDEs

Joshua Durso-Finley, Berardino Barile, Jean-Pierre Falet, Douglas L. Arnold, Nick Pawlowski, Tal Arbel

TL;DR

The paper tackles predicting continuous, individualized MS disability trajectories and treatment effects from baseline MRI data by introducing a probabilistic temporal causal model based on Neural Stochastic Differential Equations (NSDE). It encodes baseline MRI and clinical data into latent space, propagates trajectories in time conditioned on treatment via an NSDE, and decodes to EDSS while quantifying uncertainty; it also estimates longitudinal individual treatment effects (ITE) through counterfactual trajectories with pathwise uncertainty, using Neyman–Rubin causal reasoning. Evaluated on >3600 patients from six randomized trials across RRMS and SPMS, the NSDE approach outperformed LSTM and fixed-timepoint encoders, achieving an overall EDSS MSE of about 0.575 and enabling uncertainty-based selection of high-confidence subgroups with larger treatment effects. This framework advances image-based personalized medicine by delivering continuous trajectories with calibrated uncertainty and causal interpretability, with implications for patient care and drug development in chronic neurological diseases.

Abstract

Personalized medicine based on medical images, including predicting future individualized clinical disease progression and treatment response, would have an enormous impact on healthcare and drug development, particularly for diseases (e.g. multiple sclerosis (MS)) with long term, complex, heterogeneous evolutions and no cure. In this work, we present the first stochastic causal temporal framework to model the continuous temporal evolution of disease progression via Neural Stochastic Differential Equations (NSDE). The proposed causal inference model takes as input the patient's high dimensional images (MRI) and tabular data, and predicts both factual and counterfactual progression trajectories on different treatments in latent space. The NSDE permits the estimation of high-confidence personalized trajectories and treatment effects. Extensive experiments were performed on a large, multi-centre, proprietary dataset of patient 3D MRI and clinical data acquired during several randomized clinical trials for MS treatments. Our results present the first successful uncertainty-based causal Deep Learning (DL) model to: (a) accurately predict future patient MS disability evolution (e.g. EDSS) and treatment effects leveraging baseline MRI, and (b) permit the discovery of subgroups of patients for which the model has high confidence in their response to treatment even in clinical trials which did not reach their clinical endpoints.

Probabilistic Temporal Prediction of Continuous Disease Trajectories and Treatment Effects Using Neural SDEs

TL;DR

The paper tackles predicting continuous, individualized MS disability trajectories and treatment effects from baseline MRI data by introducing a probabilistic temporal causal model based on Neural Stochastic Differential Equations (NSDE). It encodes baseline MRI and clinical data into latent space, propagates trajectories in time conditioned on treatment via an NSDE, and decodes to EDSS while quantifying uncertainty; it also estimates longitudinal individual treatment effects (ITE) through counterfactual trajectories with pathwise uncertainty, using Neyman–Rubin causal reasoning. Evaluated on >3600 patients from six randomized trials across RRMS and SPMS, the NSDE approach outperformed LSTM and fixed-timepoint encoders, achieving an overall EDSS MSE of about 0.575 and enabling uncertainty-based selection of high-confidence subgroups with larger treatment effects. This framework advances image-based personalized medicine by delivering continuous trajectories with calibrated uncertainty and causal interpretability, with implications for patient care and drug development in chronic neurological diseases.

Abstract

Personalized medicine based on medical images, including predicting future individualized clinical disease progression and treatment response, would have an enormous impact on healthcare and drug development, particularly for diseases (e.g. multiple sclerosis (MS)) with long term, complex, heterogeneous evolutions and no cure. In this work, we present the first stochastic causal temporal framework to model the continuous temporal evolution of disease progression via Neural Stochastic Differential Equations (NSDE). The proposed causal inference model takes as input the patient's high dimensional images (MRI) and tabular data, and predicts both factual and counterfactual progression trajectories on different treatments in latent space. The NSDE permits the estimation of high-confidence personalized trajectories and treatment effects. Extensive experiments were performed on a large, multi-centre, proprietary dataset of patient 3D MRI and clinical data acquired during several randomized clinical trials for MS treatments. Our results present the first successful uncertainty-based causal Deep Learning (DL) model to: (a) accurately predict future patient MS disability evolution (e.g. EDSS) and treatment effects leveraging baseline MRI, and (b) permit the discovery of subgroups of patients for which the model has high confidence in their response to treatment even in clinical trials which did not reach their clinical endpoints.
Paper Structure (10 sections, 3 equations, 3 figures, 1 table)

This paper contains 10 sections, 3 equations, 3 figures, 1 table.

Table of Contents

  1. Introduction
  2. Method
  3. Learning Probabilistic Trajectories via Stochastic Neural ODE
  4. Temporal Counterfactual Trajectories and ITE
  5. Experiments and Results
  6. Dataset
  7. Experimental and Implementation Details
  8. Evaluating Predicted EDSS Trajectories and Uncertainties
  9. Counterfactual Predictions and Treatment Response
  10. Conclusions

Figures (3)

  • Figure 1: Overall framework: 3D MRI and tabular data (clinical, demographic, and subtypes) are encoded, concatenated and passed to the NSDE solver. During supervised Training (end-to-end), the model learns to represent the temporal trajectory in latent space conditioned on the treatment the patient received. The decoder maps the predicted longitudinal latent representations into outcomes (EDSS scores). During Inference, the model projects the embeddings onto the latent space forwarded in time and conditioned on the treatment (or control), permitting factual and counterfactual latent trajectories and uncertainties to be estimated. The decoder predicts future outcomes, and uncertainties, at any point along the continuous latent trajectories. ITE, and associated uncertainties, can be estimated by comparing probabilistic trajectories on and off treatments.
  • Figure 2: (a) MSE of factual EDSS predictions as a function of time. Evaluations are shown at 12-week intervals up to week 96. Results are shown for 7 active treatments and 2 placebo control groups (RRMS and SPMS). (b) Normalized Mean Squared Error (MSE) of factual EDSS predictions for patients kept based on model uncertainties. One can see a clear decrease in error when moving from right (100% of the patients) to left (30% of the patients with the most confident predictions). Shaded areas show the variance from different sets of outer fold aggregation.
  • Figure 3: Treatment response estimates (uplift) for three drugs (Ocrelizumab, DMF, Natalizumab). Comparison of predicted responders and non-responders at three different levels of uncertainties (30%, 50% and 100%-no uncertainty). Negative values indicate response through a reduction in average EDSS change on the drug as compared to placebo. Our model finds subgroups of responders for these drugs, known to be of moderate and high efficacy. The others drugs are considered to be low efficacy (Results in Supplemental Materials).