Trajectory Flow Matching with Applications to Clinical Time Series Modeling

TL;DR

This work introduces Trajectory Flow Matching (TFM), a simulation-free framework for training Neural SDEs to model stochastic, irregular clinical time series. By preserving couplings through conditional flow matching, a target-prediction reparameterization, and uncertainty modeling, TFM achieves accurate trajectory prediction and calibrated uncertainty on multiple ICU datasets, with a notable 15–83% reduction in error in some settings. The approach scales with model size and benefits from memory of recent history, while maintaining stability and speed advantages over backpropagation through SDE solvers. These results suggest practical potential for real-time clinical trajectory monitoring, provided thorough prospective validation, bias assessment, and interpretability enhancements. Key contributions include: (i) theoretical conditions ensuring simulation-free learning for continuous-time dynamics via flow matching, (ii) a reparameterization trick that stabilizes training, (iii) adaptation to irregular sampling with time-to-next-observation prediction, and (iv) demonstrated improvements in both trajectory accuracy and uncertainty prediction on four clinical datasets, highlighting the method’s relevance for clinical decision support.

Abstract

Modeling stochastic and irregularly sampled time series is a challenging problem found in a wide range of applications, especially in medicine. Neural stochastic differential equations (Neural SDEs) are an attractive modeling technique for this problem, which parameterize the drift and diffusion terms of an SDE with neural networks. However, current algorithms for training Neural SDEs require backpropagation through the SDE dynamics, greatly limiting their scalability and stability. To address this, we propose Trajectory Flow Matching (TFM), which trains a Neural SDE in a simulation-free manner, bypassing backpropagation through the dynamics. TFM leverages the flow matching technique from generative modeling to model time series. In this work we first establish necessary conditions for TFM to learn time series data. Next, we present a reparameterization trick which improves training stability. Finally, we adapt TFM to the clinical time series setting, demonstrating improved performance on three clinical time series datasets both in terms of absolute performance and uncertainty prediction.

Figures (6)

• Figure 1: Trajectory Flow Matching trains both an estimator of the next timepoint ($\hat{x}_\theta(t,x)$) and an estimation of the uncertainty ($\sigma_\theta(t, x_t)$). Using the conditional flow matching framework, these can be used to predict the instantaneous velocity $v_\theta(t, x_t)$ and future observations. Both flows are conditioned on past data $x_{[t-h, t-1]}$ and conditional variables $c$.
• Figure 2: 1D harmonic oscillator overfitting experiment results. Left: TFM-ODE (ours) with memory = 3. Middle: TFM-ODE (ours) without memory. Right: Aligned FM liu_sb_2023somnath_aligned_2023.
• Figure 3: Three samples from predicted trajectory and uncertainty on ICU GIB test set. Top: Predicted (orange) and the ground truth (blue) mean arterial pressure (MAP). Bottom: The absolute value of the uncertainty predicted by TFM.
• Figure 4: Left: Distribution of number of complete vital measurements per patient trajectory within the first 24 hours of admission in each clinical dataset. Right: Distribution of raw heart rate values in each clinical dataset.
• Figure 5: Sigma mean MSE comparison

Theorems & Definitions (8)

• Lemma 3.1
• Proposition 3.2: Coupling Preservation
• Proposition 3.3
• Lemma A.1
• proof
• proof
• proof
• proof