Trajectory Flow Matching with Applications to Clinical Time Series Modelling
Authors: Xi Zhang
Yuan Pu
Yuki Kawamura
Andrew Loza
Yoshua Bengio
Dennis Shung
Alexander Tong
Venue: NeurIPS 2024 spotlight (2024)
Keywords: Flow matching
stochastic differential equations
ODE
SDE
uncertainty
time series
EHR
Primary area: machine_learning_for_healthcare
TL;DR: Flow Matching for continuous stochastic modelling of time series data with applications to clinical time series

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 four clinical time series datasets both in terms of absolute performance and uncertainty prediction, a crucial parameter in this setting.

Trajectory Flow Matching with Applications to Clinical Time Series Modeling
Xi Zhang 2-
Yuan Pu - Yuki Kawamura'
Andrew Loza
Yoshua Bengio 2,5,6
Dennis L. Shung?
Alexander Tong? 2.5
McGill University, Mila
Quebee Al Institute.
Yale School of Medicine
4 School of Clinical Medicine, University of Cambridge, Université de Montréal, "CIFAR Fellow
Abstract
Modeling stochastic and irregularly sampled time series is a challenging problem found in 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 Neural SDE in simulation-free manner, bypassing backpropagation through the dynamics. TFM leverages the flow matching technique from generative modeling to model time series. this work we first establish necessary conditions for TFM to learn time series data. Next, we present reparameterization trick which improves training stability. Finally, we adapt TFM the clinical time series setting, demonstrating improved performance on four clinical time series datasets both terms of absolute performance and uncertainty prediction, crucial parameter in this setting.
Introduction
Real world problems often involve systems that evolve continuously over time, yet these systems are usually noisy and irregularly sampled. In addition, covariates, leading complex patterns such as intersecting trajectories. For instance, in the context of clinical trajectories in healthcare, patients' vital sign evolution can follow drastically different, crossing paths even if the initial measurements are similar, due to the influence the covariates such medication intervention and underlying health conditions, These covariates can be time-varying static, and often sparse.
time series often relate to other
Differential equation-based dynamical models are proficient at learning continuous variables with- out imputations [Chen et al.) 2018 Rubanova et al. 2019 Kidger et al., 20216). Nevertheless, systems governed by ordinary differential equations (ODEs) or stochastic differential equations (SDEs) are unable to accommodate intersecting trajectories, and thus requires modifications such as augmentation or modelling higher-order derivatives [Dupont et al. 2019). While ODEs model deterministic systems, SDEs contain diffusion term and can better represent the inherent uncertainty and fluctuations present many real world systems. However, fitting stochastic equations real life
"Joint first authorship
I Joint senior authorship. Correspondence to alexander all.aaneer.tongdmiiiil. tow.eeerr..tttmmmaaaaaaaaaaaaeee Code available at: httpp::// thub. wtt.//oo//////e http///hhaaggggaatttoyllllleeeea
38th Conference on Neural Information Processing Systems (NcurIPS 2024).

---

Uncertainty Flow
vo(t,Dr)
MA
predicted uncertainty ob
Lo (t, x))
Ita Fit 2: 2: Sa
Ug (t, xtt)
Figure 1: Trajectory Flow Matching trains both an estimator of the next timepoint (t, and an estimation of the uncertainty x+)) Using the conditional flow matching framework, these can be used to predict the instantaneous velocity ut 2x) and future observations. Both flows are conditioned on past data 111-1,111 and conditional variables c.
data challenging because they have thus far required time-consuming backpropagation through an SDE integration.
In the domain of generative models, diffusion models [Ho ct al. 20200 Nichol and Dhariwal 2021
Song et al. 2021] and more recently flow matching models [Lipman et al.J [023] Albergo al. 2023 Li et al. 2020] have had enormous success by training dynamical models is in simulation-free framework. The simulation-free framework facilitates the training of much larger models with significantly improved speed and stability. In this work generalize simulation-free training for fitting stochastic differential equations to time-series data, to learn population trajectories while preserving individual characteristics with conditionals. We present this method as Trajectory Flow Matching. We demonstrate that our method outperforms current state the art time series modelling architecture including RNN, ODE based and flow matching methods. We empirically demonstrate the utility of our method in clinical applications where hemodynamic trajectories are critical for ongoing dynamic monitoring and care. We applied our method to the following longitudinal electronic health record datasets: medical intensive care unit (MICU) data of patients with sepsis, ICU patients at risk for cardiae arrest, Emergency Department (ED) data of patients with acute gastrointestinal bleeding, and MICU data of patients with acute gastrointestinal bleeding.
Our main contributions are:
We prove the conditions under which continuous time dynamics can be trained simulation- free using matching techniques.
We extend the approach to irregularly sampled trajectories with a time predictive loss and to estimate uncertainty using uncertainty prediction loss.
We empirically demonstrate that our approach reduces the error by 15-83% when applied to the real world clinical data modelling.
2 Preliminaries
2.1
Notation
We consider the setting of a distribution trajectories over Idd denoted X := {x², 2?.......... where each z' is vector of T' datapoints i.e. z' := o2.......... x/} with associated times t := {ti, Es t Let -ft-,,t-1] denote vector of the last h observed time points. We denote (Lipschitz smooth) time dependent vector field conditioned arbitrary conditions e Re vt, Jz, -1 c) - dar di ((0, 1]. RRd, t RC) - d with flow (v(v) which induces the time- dependent density p. = dc(v) # (po) for any density Po Rd R4 with J&a Po 1. We also consider the coupling (20. 21) which operates the product space of marginal distributions Po. P1-
......")

---

2.2 Neural Stochastic Differential Equations
A stochastic differential equation (SDE) can be expressed in terms of a smooth drift f [0[0,]] x x
R and diffusion g [0, Rd
+ +R in the Ito sense as:
dx. =f g gdW
where We : [0, R the d-dimensional Wiener process. A density po(zo) evolved according to an SDE induces collection of marginal distributions pc(2;) viewed function p: [0, T) Rd -
R.. In a Newral SDE [Li et al., 2020. Kidger et al. 2021a the drift and diffusion terms are parameterized with neural networks fo(t, vx) and go(t, 21).
dx. fo(t, z4 zt)dt (t,
(I)
where the goal is to select 0 to enforce øT ~ Xtrue for some distributional notion of similarity such the Wasserstein distance [Kidger et al., 2021b or Kullback-Leibler divergence (Li al. 2020]. However, these objectives are simulation-base solver, which suffers from severe speed and stability issues. While some issues such as memory and numerical truncation can be ameliorated using the adjoint state method and advanced numerical solvers [Kidger et al. 2021b). optimization of Neural SDEs is still a significant issue.
requiring backpropagation through an SDE
We note that in the special case of zero-diffusion (i.e. go, (1)) 0) this reduces neural ordinary differential equation (Neural ODE) [Chen et al. 2018). which is easier to optimize than SDEs, but still presents challenges to scalability.
2.3 Matching algorithms
Matching algorithms are a simulation-free class of training algorithms which are able to bypass backpropagation through the solver during training by constructing the marginal distribution as mixture of tractable conditional probability paths.
The marginal density induced by eq. [ evolves according to the Fokker-Plank equation (FPE):
(pafa) 9² Ap:
Oep:
where Ap; V (Vpt) denotes the Laplacian of and gradients are taken with respect to Ty.
Matching algorithms first construct a factorization of pr into conditional densities pr (zz(z) such that Pt (x())) [pr 2)] and where Pe (3t (z) is generated by an SDE dz, (x(xxl=)dt + oz (x,
Given this construction it can be shown that the minimizer of
Lranch (0) (x (t,) (x]])==)] + X² (l,0((,
a. (xx
satisfies the FPE of the marginal pt. This is especially useful the generative modeling setting where 40 is samplable noise (e.g. (0, 1)) and q is the data distribution. Then we can define = (æo, xi) as tuple noise and data with ((=) == (x0)) 6 qi (x1). This makes eq, Blotimiie model which will draw new samples according to the data distribution gi (x1) using
30
403 -
fo(t, x. dt + a,),,
with the integration computed numerically using any off-the-shelf SDE solver. While this is guaran- guaranteed preserve the coupling of go and qı
teed to preserve the distribution over time,
(if given).
Paired bridge matching In generative modeling random pairings Liu et al. 2023c Albergo and
Vanden-Eijnden 2023 Albergo et al. 2023) or optimal transport [Tog et al. 2024 Pooladian et al. 2023] pairings are constructed for the conditional distribution ((2). However, in some problems
would like to match pairs points as the case image-to-image translation [Isola et al. 2017 Liu
al. 2023a Somnath al. 2023 In this case, training data comes pairs (xn, In this case
set (() := (Tv, F1) to be samples from these known pairs, and optimize eq. While empirically,
these models perform well, there are guarantees that the coupling will be preserved outside of the
special case when data comes from the (entropic) optimal transport coupling x (40. q1) and defined
as:
/ [d(zo, (1) x1)² dr (oo 21) + KL(T]lg0 @ q1),
Th (go. 91) = arg min (n)))
(5)
3

---

Algorithm General Trajectory Flow Matching
Input: Trajectories X while Training do
{x² x² ..... x²) noise o, initial networks vg and de.
x' (()), 1o ~~ u-,1T,T - 1}, t t~u(11,1)
H4 - (1 - t) + t
I2 NiW;; 02t(1 - -t)
(TFM(0) - v0 (K: (1,2,)
-
(,) - 1 (k: (,3+) LTFM 1
0 - Update(0, (fL) VoLa, (0))
return Ug,
where U (qo: (1) is the of admissible transport plans (i.e. joint distributions over 20 and ø whose marginals are equal to qo and ②1) as shown in (Shi et al., 2023 for some regularization parameter e E R 20-
Trajectory Flow Matching
We now describe our simulation-fre free method to learn SDEs on time-series data using trajectory flow matching as summarized in Alg. couplings are preserved. We first set out a general algorithm for flow matching on vector fields in $3.1 then present a numerical reparameterization which we find stabilizes training in $3.2 a next observation prediction for irregularly sampled time series in $33 and finally present how to learn the noise in $3.4
the case of time series we need to ensure that trajectory
3.1 Preserving Couplings
In this section, we assume access to fully observed and evenly spaced trajectories X' (x), ",,,,..,,.""" .T") with x' := (zj, .2....,,, .T) for clarity and notational simplicty. We note that our method is easily extensible to the more general setting of irregularly sampled trajectories. In this simplified case we let
= (I1, (T)
(6)
q(2) U(x)
(7)
(=(=2) =- NV([ 1)= - + (1 [(])= (M L)(l (()))
We (x/=) = 2TAT 2t
[L] -6
where ((X') is the uniform empirical distribution over X, [11 [-] are the ceiling and floor functions, and .(.-) .) is the multivariate normal distribution. This is a valid regression in the sense that function minimized with Alg. [ will return stochastic process that will match the observed marginal distributions over time as shown in the following lemma.
Lemma 3.1. The SDE dr, ue (z|2)dt + dd where u, is defined in eq. 9 Oenerrates Pe(x|2) in eq. g with initial condition Po == 82, where 8 is the Dirac delta function.
however, while useful, this is still insufficient for time series modeling, as it does not ensure coupling preservation. For intuition why this is an issue see Figure 2
In TFM we ensure that the couplings are preserved for history lengths f 0. i.c. (ETT T -ht .T) = <(IT -: 2T). We first establish a method to ensure that these couplings are preserved allowing us to use simulation-fre flow matching training for the time-series modeling task. Specifically, as long as the model takes as input (IT- Pr-h+1s... 2T)
in predicting the flow from T T 1, then there exists a function fo -k.T) such that the coupling is preserved.
Proposition 3.2 (Coupling Preservation). Under mild regulatory criteria on wx Pt, and q. if
(ur c) uu (xxlc))
and 2, 22)(2,, Pe (x)22 and un (2(22) are defined in eqs. 6 then II(x)* 11 (1

---

Where IT())* represents the coupling of a model which attains minimal loss according to eq. and II" (21r)) is the coupling of the data distribution. Intuitively, as long as no two paths cross given conditionals then the coupling is preserved. In prior work e = 0, and the coupling is only preserved in special cases such as eq.
We next enumerate three assumptions under which the coupling is guaranteed to be preserved the optima. We note that these are
(A1) When c = zo and there exists T X' + X' such that T(x)) = Ii iff II* (Co, X1). We note that this is equivalent asserting the existence of Monge map T for the coupling II*.
(A2) There exist no two trajectories x², xt such that xt = xi. for consecutive observations and
y 0.
(A3) Trajectories are associated with unique conditional vectors c independent of t.
Even in cases when (A1)-(A3) may not hold exactly, TFM is a useful model and can often still learn useful models of the data. In some sense uniqueness up to some history length is enough as it shows TFM as powerful as discrete-time autoregressive models. Proofs and further examples are available in A.
3.2 Target prediction reparameterization
While flow matching generally predicts the flow, there is a target predicting equivalent namely given (t 2 := ==" (t, and (x(2)) := [t)-t J, : 1 + (1 f)xro then it is easy to show that the target predicting loss is equivalent time- weighted flow- now-matching loss. Specifically let the target predicting loss be
=ft)_x which is equivalent to Zo when
[o]-
Curga (0) ( (t, - -x 1
(10)
then it is easy to show that Proposition 3.3. There exists a scaling function c(t) : 124 + R such that Liarge (0) = Cmutch (0).
3.3 Irregularly sampled trajectories
We next consider irregularly sampled time series of the form x² = ((x1, (a², 12), (arp t(r)) with ti < tr < ... < th with tneat denoting the next timepoint observed after time t. In this case, when combined with the target predicting reparameterization in $3.2 we can predict the time till next observation. We therefore parameterize an auxiliary model hg (t, 2z [0, x Rd - [0, which predicts the next observation time. This useful numerically, but also, perhaps more importantly, useful in a clinical setting, where the spacing between measurements can be as informative as the measurements themselves [Allam al. 2021]. ho is trained to predict the time till the next observation with the time predictive loss:
(w)
(t) (Enext t)
(11)
ET
where mexx is the time of the next measurement. This can be used in conjunction with the Thext predictor to calculate the flow time I. as
(1,11 Ft. ho IzI)
no(t, 24)
(12)
which can be used for inference on new trajectories.
3.4
Uncertainty prediction
Finally, we consider uncertainty prediction. till now we have defined conditional probability paths using a fixed noise parameter a. However, this does not have to be fixed. Instead, we consider learned (t 2x) which can be learned iteratively with the loss:
Luncertaincy (0, x)) = Σ (t. - || (t, - Foexil
(13)
tET
which learns to predict the error in the estimate of This loss can be interpreted as training an epistemic uncertainty predictor which is similar to that proposed in direct epistemic uncertainty prediction (DEUP) [ahlou et al. 2023).
5

---

peliinl
poikted
Figure 2: 1D harmonic oscillator overfitingg experiment results. Left: TFM-ODE (ours) with memory
3. Middle: TFM-ODE (ours) without memory. Right: Aligned FM [Liu et al. 2023a Somnath al. 2023
Experimental Results
In this section we empirically evaluate the performance of the trajectory flow matching objective in terms of time series modeling еπтоr, but also uncertainty quantification. We also evaluate and simulation-free methods including both stochastic and deterministic
variety of simulation-base methods. Stochastic methods are in general more difficult to fit, but can be used better model uncertainty and variance. Further experimental details can be found in S Experiments were run on a computing cluster with heterogenous cluster NVIDIA RTX8000, V100, A40, and A100 GPUs for approximately 24,000 GPU hours. Individual training runs require approximately one gpu day.
Baselines In addition different ablations of trajectory flow matching. we also evaluate Neu- raIODE [Chen et al. 2018). Ne [Li et al. 2020, Kidger et al. 2021b Kidger 2022], Latent NeuralODE [Rubanova
2019], and an aligned flow matching method (Aligned FM) [Liu et al.. 2023a, Somnath et al. 2023 where the couplings are sampled according to the ground truth coupling during training.
Metries We primarily make use of two metrics. The average mean-sques sean--suared-erro (Mean MSE) over left out time series to measure the time series modeling error defined as
MSE(
Σ <E[2.T]
|
x1112
(14)
where å and are the predicted and true trajectories respectively. We also use the maximum mean discrepancy with a radial basis function kernel (RBF MMD) which measures how well the distribution over next observation is modelled by comparing the predicted distribution to the distribution next states in the ground truth trajectory. Specifically we compute:
RBF-MMD (0, T
Σ MMD(Å,, A.) 12/2,11
(15)
where å = 8t 21-1, A samples from the model prediction at time t.
14-1, and & := f fo xx) as + go (ottxx))WW is a set of
4.1
Exploring coupling preservation with 1D harmonic oscillators
We begin by evaluating how trajectory flow matching performs in a simple one dimensional setting of harmonic oscillators. We show that the canonical conditional flow and bridge matching [Liu
al. 2023c Albergo and Vanden-Eijnden. 2023). specifically aligned approaches (Somnath et al. Liu et al. 2023a] are unable to preserve the coupling even in a simple one dimensional setting. However, augmented with our trajectory flow matching approach, and specifically using (A2), which includes information on previous observations, the model is able to fit the harmonic oscillator dataset well.
2023
The harmonic oscillator dataset consists of one-dimensional oscillatory trajectories from damped harmonic oscillator, with each trajectory distinguished by a unique damping coefficient Specifically we sample trajectories x from:
Iz Ii-1 + G--(tt t-11); Zo
(16)

---

where u is the velocity of the oscillator updated by
(- m U1-1 21-1 (t, -t-1); vo
c
t 1 +
(17)
with 42 0.1 for i 2, 99, spring constant k: 1, and mass m 1.
As increases, the trajectories evolve from underdamped scenarios with prolonged oscillations to critically and overdamped states where the oscillator quickly stabilizes. This leads to intersecting trajectories due to frequency and phase differences, despite their shared starting point. We perform overfitting experiments on three trajectories generated by varying c.
As shown in Figure 2 models without history information are unable to distinguish between the three crossing trajectories that share the same starting point, resulting in overlapping predictions. In contrast, TFM-ODE that incorporates three previous observations able to fit the crossing trajectories with high accuracy, with the predicted trajectories almost completely overlapping the ground truth. This is because the dataset with satisfies (A2) with h = 4 (TFM-ODE), but not h = 0 (TFM-ODE no memory and Aligned FM).
4.2 Experiments on clinical datasets
Next we compared the performance of TFM and TFM-ODE with the current SDE and ODE haselines, respectively, for modeling real-world patient trajectories formed with heart rate and mean arterial blood pressure measurements within the first 24 hours of admission across four different datasets. These are clinical measurements that are taken most frequently and used to evaluate the hemodynamic status of patients, key indicator of disease severity. Additionally, we evaluated our models against flow matching on these datasets, each with distinct characteristics, to assess their ability generalize across different distributions. A full description of the datasets are available in Appendix B.2 with the publicly available datasets used under The PhysioNet Credentialed Health Data License Version 1.5.0 and the EHR dataset with local institutional IRB approval:
ICU Sepsis: a subset of the eICU Collaborative Research Database v2.0 PPollrd et al. 2019 of patients admitted with sepsis as the primary diagnosis.
ICU Cardiac Arrest: a subset of the eICU Collaborative Research Database v2.0 Pollard et af.
2019] of patients risk for cardiac arrest.
ICU GIB: subset of the Medical Information Mart for Intensive Care III [Johnson et al. 2016] of patients with gastrointestinal bleeding as the primary diagnosis.
ED GIB: patients presenting with signs and symptoms of acute gastrointestinal bleeding to the emergency department of a large tertiary care academic health system.
4.2.1 Prediction accuracy and precision: TFM and TFM-ODE
TFM-ODE yields more accurate trajectory prediction Across the four datasets TFM-ODE outperformed the bascline models by 15% to 20%, as seen in table [ We noticed that TFM has a similar performance as TFM-ODE. In case TFM outperformed the non-stochastic TFM-ODE, as seen in the ICU GIB dataset. For ICU sepsis, the performance improvement from the baseline is the most significant, around 83%. This coincides with the ICU sepsis dataset having the most amount of measurement per trajectory. The improvement is seen both TFM and TFM-ODE, possibly indicating they are able to learn better given more data, resulting in more precise flow. Not formally measured, we noted that given the same time constraint, FM based models were significantly faster and often finished training before the time limit.
TFM yields better uncertainty prediction Though TFM- ODE had lower test MSE for half of the times, TFM yielded better uncertainty prediction overall, as seen in table Notably, TFM also had less variance in the uncertainty prediction than TFM-ODE. plausible explanation in this case is a sacrifice in bias that subsequently decreases the variance for the stochastic implementation, reflecting the bias-variance trade off. Sampled graphs of TFM can be seen in figure 3 It is notable that the model is able to detect the measurement uncertainty at certain timepoints, matching the increase in amplitude oscillation in patient trajectories.
4.2.2 Trajectory Variance Distribution Comparison
TFM trajectories accurately match the noise distribution in the data TFM able to match the noise distribution in addition to the overall trajectory shape, which is useful in settings where
7

---

Table Mean + Std. deviation MSE (x10-3 by models and datasets. Split into deterministic (top)
and stochastic models (bottom). Top performing model for each setting and dataset in bold.
ICU Sepsis
ICU Cardiac Arrest
ICU GIB
GIB
NeuralODE
4.776 + 0.000
6.153 + 0.000
3.170 + 0.000
10.859 0,000
FM baseline ODE
4,671 + 0.791
10.207:: 1.076
118,439 = 17.947
11.923
1.123
LatentODE: RNN
61.806 46.573
386.190 + 558.140
422.886 431.954
980.228++033..93
TFM-ODE (ours)
0.793 + 0.017
2.762 + 0.021
2.673 + 0.069
8.245 0.495
NeuralSDE
4.747 + 0.000
3.250 0.024
3.186 0.000
10.850 0.043
TFM (ours)
0.796 1 0.026
2.755 1 0.015
2.596 1 0.079
8.613 0.260
Table Uncertainty test MSE loss for TFM-ODE and TFM with two different ICU datasets.
ICU sepsis
ICU Cardiac Arrest
ICU GIB
TFM-ODE
1.039 0.1645
0,970 1426
0,9843 + 0.2233 0.605 + 0.0137
TFM
0.724 + 0.0072
0.636 + 0.0024
0.25
MAP
-0.25
0.-
20
40
60
10
20
30
40
30
10
15
20
25
0.0
002
10
70
Figure Three samples from predicted trajectory and uncertainty on ICU GIB set. Top: Predicted (orange) and the ground truth (blue) mean arterial pressure (MAP). Bottom: The absolute value the uncertainty predicted by TFM.
data has high stochasticity. We compared our models to NeuralODE and NeuralSDE in matching the variance in neighboring data points, seen in table We verify that between the baseline NeuralSDE and NeuralODE, NeuralSDE has a lower MMD and is better able to match data points. We find in ICU GIB and ED GIB datasets, TFM outperforms both in matching the variance in data. Notably, the performance pattern is reversed for the MMD metrics and mean MSE metrics with respect to TFM and TFM-ODE where better MSE leads to worse MMD and vice versa. As such, this further confirms the
trade-off for both TFM and TFM-ODE implementation.
Table 3: Data variance MMD for by models and datasets. Split into deterministic models (top) and stochastic models (bottom). Top performing model for each setting and dataset in bold.
ICU Sepsis
ICU Cardiac Arrest
ICU GIB
ED GIB
NeuralODE
1.988
0.000
2.246 + 0.000
2.090 + 0.000
2.192 + 0.000
TFM-ODE (ours)
1.172
0.017
1.295 + 0.006
1.087 + 0.02
1.063 + 0.031
NeuralSDE
1.212
0.000
3.261 + 0.020
1.332 + 0.000
1.465 + 0.122
TFM (ours)
1.199
0.006
0.993 + 0.003
0.844
0.013
0.717
0.016
4.2.3 Ablation Study
We performed ablation studies on TFM and TFM-ODE to attribute importance of various model components contributing to the performance, as seen in table We examined three aspects of the model, two of which were part of our main contributions: uncertainty prediction and memory. We also ablate the model hidden dimension width to infer its potential in scaling effect.
TFM and TFM-ODE performance scales with model size In contrast to Neural DE based models, TFM and TFM-ODE exhibit scaling effect, in which the model performance becomes better with a larger hidden dimension. This has been observed previous flow matching models in image generation Tong et al. 2024). This may pave the way for further improvements from larger models.

---

Table 4: Mean MSE (x10-)) by ablated versions of TFM, TFM-ODE, and datasets.
Uncertainty Prediction
Memory
Hidden Size
ICU Sepsis
ICU Arrest
Cardiac
ICU GIB
ED GIB
TFM-ODE
256
0.793 1 0.017
2.762 1 0.017
2.673 1 0.069
8.245 1 0,495
256
1.170 + 0.014
2.759 + 0.015
3.097 + 0.054
8.659 = 0.429
256
1,555 + 0.122
3,242 + 0.050
2.981 + 0.161
6,381 + 0.451
1,936 = 0.262
3.244= 0.025
4.003 = 0.347
11.253+ 4.597
TFM
V
256
0.796 = 0.026
2.596 + 0.079
2.762 + 0.021
8.613 + 0.260
256
0.816 = 0.031
2.778 + 0.021
2.754 + 0.095
8,600 = 0,389
64
1.965 1 0.289
3,271 + 0.031
4.037 1 0.314
7.549 1 0.737
Uncertainty improves performance of trajectory prediction For TFM and TFM-ODE, the flow network used to learn the uncertainty dz, is separate from the flow network learning 2,. The loss function of the network learning 2; is independent of uncertainty flow network. Therefore, it was unexpected that taking away the uncertainty prediction would result in increased MSE test loss for learning 24, This implies further process in the synergistic effects between flow and an, flow.
Trajectory memory may improve performance in high frequeney measurement settings We conditioned the model based on sliding window of trajectory history to disentangle data points that otherwise look indistinguishable to FM models. This improved the interpolation performance in the ICU Sepsis and ICU GIB dataset. Notably, this modification did not improve performance the ED GIB dataset, which could be due shorter trajectories for patients and lower measurement frequency in the defined time period. This may also be explained by the decreased severity of disease in the ED compared the ICU. Adding memory as a condition may be more suitable for patients whose clinical trajectories have higher frequency of measurements.
5 Related Work
Continuous-timg neural network architectures have outperformed traditional RNN methods in mod- eling irregularly sampled clinical time series to optimize interpolation and extrapolation. Neural ODE with latent representations of trajectories [Rubanova et al., 2019] outperformed RNN-based approaches [Lipton et al. 2016 Che et al. 2018 Cao et al.) [018 Rajkomar et al.) [018] for inter- polation while providing explicit uncertainty estimates about latent states. More recently, Neural SDEs appear to outperform LSTM [Hochreiter and Schmidhuber 1997]. Neural ODE [Chen et al.
2018 [De Brouwer et al. 2019 Dupont 2019 Lechner and Hasani 2020], and attention-based
[Shukla and Marlin, 202 Lee et al. 2022] approaches in interpolation performance while natively handling uncertainty using drift and diffusion terms [Oh et al. 2024].
Discrete-time approaches offer an alternative to our continuous-time model model transformers utilize a discrete-ting time representation with sequential processing [Gao et al., 2024 Nie et al. 2023
Woo et al. 2024 Ansari et al. [2024] Dong et al.) [2024] Garza and Mergenthaler Canseco [023]] Das
et al. 2024 Liu et al. [2024] Kuvshinova et al. 2024 models for traditional time series modeling.
Adaptations to baseline transformer includes structuring observations into text with finetuning
(Zhag al. [023 Zhou et al. 2023), without finetuning [Xue and Salim 2024 Gruver et al. 2023)
or using autoregressive model vision transformers to model unevenly spaced time series data by
converting time series into images (Li et al.) 2023].
Continuous-time systems are of great interest for learning causal representations using assumptions by using observations directly modify the system state (De Brouwer al. [022]] Jia and Benson 2019]. Variations include intervention modeling with separate ODEs for interventions and outcome
processes Gwak et al. 2020], using liquid time-constant networks (Hasani et 2021 Vorbach et al. 2021], or modeling treatment effects with either one [Bellot and van der Schaar 2021] or multiple interventions (Seedat et 2022). The importance of accounting for external interventions a particular challenge in clinical data, where external interventions (change in environment due to treatment decisions or clinical context such as ED or ICU) are common clinical data trajectories.
Conclusion
In this work we present Trajectory Flow Matching, simulation-free training algorithm for neural differential equation models. We show when trajectory flow matching valid theoretically, then demonstrate its usefulness empirically in clinical setting. The ability to model the underlying

---

continuous physiologic processes during critical illness using irregular, sparsely sampled, and noisy data has the potential for broad impacts in care settings such as the emergency department or ICU, These models could be used to improve clinical decision making, inform monitoring strategies, and optimize resource allocation by identifying which patients are likely to deteriorate or recover. These use cases will require thorough prospective validation and calibration for specific clinical outcomes, for example using the likelihood of a patient crossing a specific heart rate or blood pressure threshold for decisions on level of care (ICU versus inpatient floors) or specific interventions such as transfusions. In these applications, will be important to assess and control for bias that may be present due to which patient subpopulations are present in training data.
Limitations Limitations of the method includes the selective utility of integrating memory in clinical settings with high measurement frequency and no current capacity for estimating causal representations, though this will be an important future research direction. Potential harms include the following: erroneous predictions that either results in delayed care or overutilization of the health system. Accurate trajectory predictions have the potential to inform clinical decision-making regarding the appropriate level of care, leading more timely and appropriate interventions.
Future work We hope to extend our method to cover other types of time series that have periodicity in the components, potentially incorporating Fourier transform [Li al. [02T] and thysics-Inspiree Neural Networks (PINN). Since interpretability is an important factor for clinical reliability, we are
developing methods to further elucidate key components affecting the prediction. As well, we hope to incorporate functional flow matching for fully continuous setting [Kerrigan et al.) 2024).
7 Broader Impact
Our work extends flow matching into the domain of time series modeling, demonstrating a specific instance of clinical time series prediction. In contrast the large transformer-basee models, our method has fewer in parameters and less training time needed. Notably, it scales well with parameters. As well, our parameterization on Stochastic Differential Equations (SDE) allow faster training time than traditional SDE integration.
Accurate timeseries modeling in healthcare has the potential for significant benefits, but also in- troduces risks. Benefits that could be derived from more accurate prediction of clinical courses include improved treatment decisions, resource allocation, as well as more informative discussions of prognosis with patients or family members. Risks may come from inaccuracies in predictions which could lead to harms by biasing decision making of clinical teams. In the general case of false negative prediction (prediction of trajectories with falsely favorable outcomes) this may lead to undertreatment and in the case of false positive prediction (prediction of trajectories with incorrect detrimental outcomes) or overtreating patients. These inaccuracies may also propagate biases in training data.
To move towards broad impact in the clinical domain, this work will require validation and bias esti- mates. Furthermore, models deployed in domains with high-stakes prediction require interpretability. which can help identify biases, miscalibration, discordance with domain knowledge, as well as build trust with teams using predictions from the model. At this time, flow-based methods have limited tools for interpretability, and we recognize this as gap in need of future work.
Acknowledgements
The authors would like to thank Mathieu Blanchette for useful comments on early versions of this manuscript. We are also grateful to the anonymous reviewers for suggesting numerous improvements.
The authors acknowledge funding from the National Institutes of Health, UNIQUE, CIFAR, NSERC, Intel, and Samsung. The research was enabled in part by computational resources provided by the Dig- ital Research Alliance of Canada (https: /tttss:/aallannees Yale School of Medicine and NVIDIA.
ca), Mila battps://mile quebec),
References
M. S. Albergo and F. Vanden-Eijnden. Building normalizing flows with stochastic interpolants. In The Eleventh International Conference on Learning Representations, 2023. URL. https //openreview.ro htt.:e/eoorrm/ed1117eeooooo
10

---

M. S. Albergo, N. M. Boffi, and E. Vanden-Eijnden. Stochastic interpolants: A unifying framework for flows and diffusions. CoRR. abs/2303.08797, 2023. /ntpps//:///oi org 4 arXiv 2303.08797
A. Allam, S. Feuerriegel, M. Rebhan, and M. Krauthammer. Analyzing patient trajectories with artificial intelligence. J Med Internet Res, 23(12)::29812, Dec 2021. ISSN 1438-8871. doi:
10.2196/29812 12. URL wwwphhttp/////tt http:://www. jmir org/ (trg/2221/12/e29811
A. F. Ansari, L. Stella, C. Turkmen, X. Zhang, P. Mercado, H. Shen, O. Shchur, S. S. Rangapuram, S. Pineda Arango, S. Kapoor, J. Zschiegner, D. C. Maddix, H. Wang, M. W. Mahoney, K. Torkkola, A. Gordon Wilson, M. Bohlke-SShhniide and Y. Wang. Chronos: Learning the language of time series. arXiv preprint arXiv:2403.0755 2024.
A. Bellot and M. van der Schaar. Policy analysis using synthetic controls in continuous-timg time.. In M. Meila and T. Zhang, editors, Proceedings of the 38th International Conference on Machine Learning, volume 139 of Proceedings of Machine Learning Research, pages 759-768. PMLR, 18-24 Jul 2021. URL
mlr
http:///v/eeeelttaaaaaae html
W. Cao, D. Wang, Li, H. Zhou, L. Li, and Y. Li. Brits: Bidirectional recurrent imputation for time series. In S. Bengio, H. Wallach, H. Larochelle, K. Grauman, N. Cesa-Bianchi, and R. Garnett, editors, Advances in Neural Information Processing Systems, volume 31. Curran As- sociates, Inc., 2018 URL http//tttts///rreeeedigg hww.riinisniinnnsnneeeeee. cctsscccpppppppe www.se/pppeeeee/..///// htle:9/eeee.c5Beccccc..bbbbb1b1. Pappeer. pdf
Z. Che, S. Purushotham, K. Cho, D. Sontag, and Y. Liu. Recurrent neural networks for multivariate time series with missing values. Scientific Reports, 8(1), 2018. doi: 10. 10038//41598-01--22277
R. T. Q. Chen, Y. Rubanova, J. Bettencourt, and D. K. Duvenaud. Neural ordinary differential equations. In S. Bengio, H. Wallach, H. Larochelle, K. Grauman, N. Cesa-Bianchi, and R. Gar- nett, editors, Advances in Neural Information Processing Systems, volume 31. Curran Asso- ciates, Inc., 2018. URL htttps::/pprcceeednggg htwprooceeeinngn.neurrii..eee
www.ieuu...... ipa. wwcpaper_.aeeeeeeeeeeeeee pdf
T. Chen and C. Guestrin. Xgboost: A scalable tree boosting system. In Proceedings of the 22nd acm sigkdd international conference on knowledge discovery and data mining, pages 785-794, 2016.
M. M. Churpek, T. C. Yuen, S. Y. Park, D. O. Meltzer, J. B. Hall, and D. P. Edelson. Derivation of cardiac arrest prediction model using ward vital signs. Critical Care Medicine, 2012.
A. Das, W. Kong, R. Sen, and Y. Zhou. A decoder-only foundation model for time-series forecasting. In Forty-first International Conference on Machine Learning. 2024. URL /tRp: net/foruee jttp:/orrrnn-jj2iiiaaaaas.s
E. De Brouwer, J. Simm, A. Arany, and Y. Moreau. Gru-ode-bayes: Continuous modeling of sporadically-osserved time series. In H. Wallach, H. Larochelle, A. Beygelzimer, F. d'Alché- Buc, E. Fox, and R. Garnett, editors, Advances in Neural Information Processing Systems, vol- ume 32. Curran Associates, Inc., 2019, URL https: htttps://ppocceediinnge www.eeedddnnnnnnnensnee..ee.m.e. neurips. cctpppp.cccppppp.ppppppppp pdf
E. De Brouwer. J. Gonzalez, and S. Hyland. Predicting the impact of treatments over time with uncertainty aware neural differential equations. In G. Camps- Valls, F. J. R. Ruiz, and I. Valera, editors, Proceedings The 25th International Conference Artificial Intelligence and Statistics, volume 151 of Proceedings of Machine Learning Research, pages 4705-4722, PMLR, 28-30 Mar 2022. URL hntpss//pproceedingg mll. gress/v151/der http://teeebrouue22aahhthta
J. Dong. H. Wu, Y. Wang, Y.-Z. Qiu, L. Zhang. J. Wang, and M. Long. Timesiam: A pre-training framework for siamese sime-series modeling. In Forty-first International Conference on Machine Learning, 2024. URL https:://ppperrereee..u www.:nneeffooooe http:p/wwwwwweeeeeeeooooooo//aaa
E. Dupont, A. Doucet, and Y. W. Teh. Augmented neural odes. In H. Wal- lach, H. Larochelle, A. Beygelzimer, F. d'Alché-Bue, E. Fox, and R. Garnett, edi- tors, Advances in Neural Information Processiny Systems, volume 32. Curran Associates, Inc., 2019. URL CbD@@alldeee11000e11111111111--2 .pff
www.rroeeeeiiinnneeeeeeeeereeeee nwurips. htappp.cc/ppaee wwwpcc//ppppeeeeee
fttppppppr_fflss___eepppppppppee files/fieee/eeeeeeee Shappr_fflees//ppr//201111 /aper//0119/fiee/

---

S. Gao, T. Koker, O. Queen, T. Hartvigsen, T. Tsiligkaridis, and M. Zitnik. Units: Building a unified time series model. arXiv, 2024. UtLhttsss//arxi org. wwtp:/tgppp/pp/o/oooooo 201.//2223..0.011..
A. Garza and M. Mergenthaler-Canseeo Timegpt-1, 2023.
N. Gruver, M. Finzi, S. Qiu, and A. G. Wilson. Large language models are zero-shot time series forecasters. In A. Oh, T. Naumann, A. Globerson, K. Saenko. M. Hardt, and S. Levine, editors, Advances in Neural Information Processing Systems, volume 36, pages files/paper/2023/ sociates, Inc., 2023. URL http:///rrrooooeeinngg neurips htcc/ppprrrfieee///eeeeeeea ConeereeCeeeeeccc.coo ftle///77ase22077772222222222112
Curran As-
D. Gwak, G. Sim, M. Poli, S. Massaroli, J. Choo, and E. Choi. Neural Ordinary Differential Equations for Intervention Modeling. arXiv e-prints, art. arXiv:2010.08304, Oct. 2020. doi: 10. 100488500arXxv.20000.8888
R. Hasani, M. Lechner, A. Amini, D. Rus, and R. Grosu, Liquid time-constant networks. Proceedings of the AAAI Conference on Artificial Intelligence, 350))7657-7666, May 2021. doi: 10. 1609/aia.... v3519.16936. URL https://ooj/ aaai. wrg.i.og//aaaaaaa..aaaaaaaa index http://wwww.ppp...aaaaaaaaaaae/e
J. Ho, A. Jain, and P. Abbeel. Denoising diffusion probabilistic models. Advances in Neural Information Processing Systems, 33:6840-68511 2020.
S. Hochreiter and J. Schmidhuber. Long short-term memory. Neural Computation, 175--773--7888... 1997. ISSN 1530-888X. doi: 10.1162//eco.11/ 19.7.9.8.11.... 1735.. URI. http::/dx...d..ee org. org/10 1162/neco 1997.99999 9.87.998.0 1735
Isola, Zhu, Zhou, and A. A. Efros. Image-to-image translation with conditional adversarial networks. CVPR, 2017.
Jia and A. Benson. Neural jump stochastic differential equations. In H. Wal- lach, H. Larochelle, A. Beygelzimer, F. d'Alché-Buc, E. Fox, and R. Garnett, edi- tors, Advances Neural Information Processing Systems, volume 32. Curran Associates. Inc., 2019. URL htttps:/ Paperr Paper. pdf
neurips. wttppp...ppppppppe wcwpaper........ files/ /ttpppeea.leeepp////////////1/11 paper/2019/file/
A. Johnson, J. Pollard, L. Shen, L.-w. H. Lehman, M. Feng, M. Ghassemi, B. Moody, P. Szolovits, L. Anthony Celi, and R. G. Mark. Mimic-iii, freely accessible critical care database. Scientific data, 3(1):1-9, 2016.
G. Kerrigan, G. Migliorini, and Smyth. Functional flow matching. In S. Dasgupta, Mandt, and Y. Li, editors, Proceedings of The 27th International Conference on Artificial Intelligence and Statistics, volume 238 Proceedings of Machine Learning Revearch, pages 3934-3942, PMLR, 02-04 May 2024. URL http:htts://// htttt::s//prrocedingg. mmr. wtt....resss//33//errrraaa222 html
Kidger. On neural differential equations, 2022. URL
P. Kidger, J. Foster, X. Li, H. Oberhauser, and T. Lyons. Neural sdes as infinie-ddimensional gans. In International conference on machine learning. PMLR, 2021a.
P. Kidger, J. Foster, X. C. Li, and T. Lyons. Efficient and accurate gradients for neural sdes. In M. Ranzato, A. Beygelzimer, Y. Dauphin, P. Liang, and J. W. Vaughan, editors, Advances in Neurul Information Processing Systems, volume 34, pages 18747-18761. Curran Associates, Inc., 2021b. URL 1ba.9956r7eeeeaaaaa0994440002222 Paper pdff
wwwpeeeeddiinnnnenineeeeeeeeeeee neurips. cc/ ftlpeeeeeeepppppppe.eeeeeeee/e/e Seapper_fileesppppr/22//////1/
K. Kuvshinova, O. Tsymboi, A. Kostromina, Simakov, and Kovtun. Towards foundation time series model: To synthesize not to synthesize? arXiv preprint arXiv:2403.002534 2024.
Lahlou, M. Jain, H. Nekoei, Butoi, P. Bertin, Rector-Brooks, M. Korablyov, and Y. Bengio. DEUP: Direct epistemic uncertainty prediction. Transactions on Machine Learning Research, 2023. ISSN 2835-8856. URL. htttae///pppaarreee..eee///ee/ed Certification.
Expert
12

---

M. Lechner and R. Hasani. Learning long-term dependencies in irregularly-sappl time series. arXiv preprint Stlv:2006.044188 2020.
Y. Lee, E. Jun, J. Choi, and H.-I. Suk. Multi-view integrative attention-based deep representation learning for irregular clinical time-series data. IEEE Journal of Biomedical and Health Informatics. (6(8):4270-42888) 2022. doi: 10.1 20.1099/BHL2222.31725544
X. Li, T.-K. L. Wong. R. T. Q. Chen, and D. K. Duvenaud. Scalable gradients and variational inference for stochastic differential equations. In C. Zhang, F. Ruiz, T. Bui, A. B. Dieng, and D. Liang, editors, Proceedings of The 2nd Symposium on Advances in Approximate Bayesian Inference, volume 118 of Proceedings of Machine Learning Research, pages 1 1-28. PMLR, 08 Dec 2020. URL httpsstt////rooeeediaa
mlr. press/v118/1i20a.
Z. Li, N. B. Kovachki, K. Azizzadenesheli, B. liu, K. Bhattacharya, A. Stuart, and A. Anandkumar. Fourier neural operator for parametric partial differential equations. In International Conference on Learning Representations, 2021. URL
Z. Li, Li, and X. Yan. Time series images: Vision transformer for irregularly sampled time series. In A. Oh, T. Naumann, A. Globerson, K. Saenko, M. Hardt, and S. Levine, editors, Advances in Neural Information Processing Systems, volume 36, pages 49187-49204. Curran As- sociates, Inc., 2023. URL//ttpss//prroceedddagg wwwrips. cc/pper f 21llss/paper/22233 file/
Cawaaaperrrooffffeeeeeeeee pdf
Y. Lipman, R. Q. Chen, H. Ben-Hamu, M. Nickel, and M. Le. Flow matching for generative modeling. In The Eleventh International Conference Learning Representations, 2023. URL https://openrevial
Z. C. Lipton, D. Kale, and R. Wetzel. Directly modeling missing data in sequences with rnss: Improved classification of clinical time series. F. Doshi-Velez, J. Fackler, D. Kale, B. Wallace, and J. Wiens, editors, Proceedings of the 1st Machine Learning for Healthcare Conference, volume 56 of Proceedings of Machine Learning Research, pages 253-270, Northeastern University, Boston, MA, USA, 18-19 Aug 2016. PMLR. URL Lipton16.1. html
mlr wwtpsemww........eeeeeeeeee press/
G.-H. Liu, A. Vahdat, D.-A. Huang, E. A. Theodorou, W. Nie, and A. Anandkumar. 12sb: image- to-image schrödinger bridge. Proceedings of the 40th International Conference on Machine Learning, ICML'23. JMLR.org. 2023a.
X. Liu, C. Gong, and Q. Liu. Flow straight and fast: Learning to generate and transfer data with rectified flow. The Eleventh International Conference on Learning Representations (ICLR), 2023b. URL https: //par. nsf http//bbl1o/1004455555
X. Liu, L. Wu, M. Ye, and qiang liu. Learning diffusion bridges on constrained domains. In The Eleventh International Conference on Learning Representations, 2023c. URL https:/7 openreview
Y. Liu, H. Zhang, C. Li, X. Huang, J. Wang, and M. Long. Timer: Transformers for Time Series Analysis at Scale. arXiv e-prints, art. arXiv:2402.022668 Feb. 2024. doi: 10.485500/rxivv2440.. 02368.
A. Q. Nichol and Dhariwal. Improved denoising diffusion probabilistic models. In M. Meila and T. Zhang, editors, Proceedings of the 38th International Conference Machine Learning, volume 139 of Proceedings of Machine Learning Research, pages 8162-8171, PMLR, 18-24 Jul 2021. URL httpp::/pproceedingg/ mlr. wtt..esssss///iccollless1 /ttp::///ppeeeiiiig......//////s
html
Y. Nic. N. H. Nguyen, Sinthong, and J. Kalagnanam. A time series is worth 64 words: Long- term forecasting with transformers. In The Eleventh International Conference Learning Representations, 2023. URL httl:httt///// httpttts//ppeereiiiee
Y. Oh, D. Lim, and S. Kim. Stable neural stochastic differential equations in analyzing irregular time series data. In The Twelfth International Conference on Learning Representutions, 2024. URL httpp:://ppenreviee wttp://ooooo..dd/iiiinieeeee
13

---

T. Pollard, A. Johnson. J. Raffa, L. A. Celi, O. Badawi, and R. Mark. eicu collaborative research database (version 2.0). 2019. httpttt:///oo..rgg///// http:///2aa.rrr..1..110001////1 13026/C2WM1R
A.-A. Pooladian, H. Ben-Hamu, C. Domingo-Enrich, B. Amos, Y. Lipman, and R. T. Q. Chen. Multisample flow matching: straightening flows with minibatch couplings. In Proceedings of the 40th International Conference on Machine Learning, ICML'23. JMLR.org, 2023.
A. Rajkomar, E. Oren, K. Chen, A. M. Dai, N. Hajaj. M. Hardt, P. Liu, X. Liu, J. Marcus, M. Sun, P. Sundberg, H. Yee, K. Zhang, Y. Zhang, G. Flores, G. E. Duggan, J. Irvine, Q. Le. K. Litsch, A. Mossin, J. Tansuwan, D. Wang, J. Wexler, J. Wilson, D. Ludwig, S. L. Volchenboum, K. Chou, M. Pearson, S. Madabushi, N. H. Shah, A. J. Butte, M. D. Howell, C. Cui, G. S. Corrado, and Dean. Scalable and accurate deep learning with electronic health records. upj Digital Medicine, 1(1):18, 2018. doi: 10. (038//411446-11 1--0099-1
Y. Rubanova, R. Т. Q. Chen, and D. K. Duvenaud. Latent ordinary differential equations for irreuulaly-sampledd time series. In H. Wallach, H. Larochelle, A. Beygelzimer, F. d'Alché- Buc, E. Fox, and R. Garnett, editors, Advances in Neural Information Processing Systems, vol- ume 32. Curran Associates, Inc., 2019. URL https::/procceeeeeee
.ew.iiee...p.... cc/paper_ pdf
N. Seedat, F. Imrie, A. Bellot, Z. Qian, and M. van der Schaar. Continuous-time modeling of counterfactual outcomes using neural controlled differential equations. In K. Chaudhuri, S. Jegelka, I.. Song, C. Szepesvari, G. Niu, and S. Sabato, editors, Proceedings of the 39th International Conference on Machine Learning, volume 162 of Proceedings of Machine Learning Research, pages 19497-19521. PMLR, 17-23 Jul 2022. URL https://proocedings.ct mlr mwt.s/ss.r..eeeeee12222 presa/v162/ seedat22b. html
Y. Shi, V. De Bortoli, A. Camphell, and A. Doucet. Diffusion schrödinger bridge matching. In A. Oh, T. Naumann, A. Globerson, K. Saenko, M. Hardt, and S. Levine, editors, Advances in Neural Information Processing Systems, volume 36, pages 62183 66223.. Curran Associates, Inc., 2023. URL htttps://procceeddiii neurips cc/paper...p......pcppp filefillll wwwes//aleeleiee paper/2023//tiee/ 212266ff777222202222442266222222
pdf
S. N. Shukla and B. Marlin. Multi-time attention networks for irregularly sampled time series. In International Conference on Learning Representations, 2021. URL httph:t////////o //tp::::///perevveee net/
forum7id=4c006lv02
V. R. Somnath, M. Pariset, Y.-P. Hsieh, M. R. Martinez. A. Krause, and C. Bunne. Aligned diffusion Schrödinger bridges. In R. J. Evans and I. Shpitser, editors, Proceedings of the Thirty-Ninth Conference on Uncertainty in Artificial Intelligence, volume 216 of Proceedings of Machine Learning Research, pages 1985-1995. PMLR, Jul-04 Aug 2023. WtRLhhttps//pprceeddii mlr. 2rrss//216/ssomah2223 htmm
Y. Song, J. Sohl-Dickstein, D. P. Kingma, A. Kumar, S. Ermon, and B. Poole. Score-based generative modeling through stochastic differential equations. In International Conference Learning Representations, 2021. URL http:tttp//:/oorrrriri
A. Tong, K. FATRAS, N. Malkin, G. Huguet, Y. Zhang, J. Rector-Brooks, G. Wolf, and Y. Ben- gio. Improving and generalizing flow-based generative models with minibatch optimal trans- port. Transactions on Machine Learning Research, 2024. ISSN 2835-8856. URL https wtwp://www.neetooumiieeeooooo
Expert Certification.
C. Vorbach, R. Hasani, A. Amini, M. Lechner, and D. Rus. Causal navigation by continuous-time neural networks. In M. Ranzato, A. Beygelzimer, Y. Dauphin, P. Liang, and J. W. Vaughan, editors, Advances in Neural Information Processing Systems, volume 34. pages 12425- 12440. Curran As- sociates, Inc., 2021. URL URU/attss::/pprooeddnaee htw..ip.cccpppppe.fffeee///epeep
neurips cc/paperppppp..cccpccco ccc/paper_filessee fiwes/ ffle6611022002210000550077777772 Paper. "paper.dff
G. Woo, C. Liu, A. Kumar, C. Xiong, S. Savarese, and D. Sahoo. Unified training of universal time series forecasting transformers. Forty-first International Conference on Machine Learning, 2024. URL htllhtttss//opeeeiiiie htt.:pawwwweee.eeeeeooaaaaateaan
14

---

H. Xue and F. D. Salim. PrompiCast: A New Comptt-Based Learning Paradigm for Time Series Forecasting IEEE Transactions on Knowledge & Data Engineering. 36(11):6851-
6864, Nov. 2024. ISSN 1558-2191. doi: 1011099TKKDE.22333422111
URL https
htt.://wwcccpperroocciiooeoooo/a 110.////0..11.EEEE.11.. 2023 3342137
Y. Zhang, K. Gong, K. Zhang, H. Li, Y. Qiao, W. Ouyang, and X. Yue. Meta transformer: unified framework for multimodal learning. arXiv preprint wrXiv:2307.108800 2023.
T. Zhou, P. Niu, wang, L. Sun, and R. Jin. One fits all: Power general time series analysis by pretrained Im. In A. Oh, T. Naumann, A. Globerson, K. Saenko, M. Hardt, and S. Levine, editors, Advances in Neural Information Processing Systems, volume 36, pages 43322-43333 33322-43355. Curran As- sociates, Inc., 2023. URL htrL//tttps:/prrceeeding www.diiiinnnnnn..n......m wwwpppppppppppppppppp. f wiles/paper/2000 F1117577777657700222224444622222
pdf
Zimmerman, A. A. Kramer, S. McNair, and M. Malila. Acute physiology and chronic health evaluation (apache) iv: hospital mortality assessment for today's 's critically ill patients. Critical care medicine, 1995)1127 2006.
15

---

A Proof of theorems
We first prove Lemma which shows TFM learns valid flows between distributions with the target prediction reparameterization trick.
Lemma A.1. If (e()) > 0, òdata is Lipschitz continuous for all x E IR and e [0, 11. LFM and LrPM are equal,
= (
Proof. This proof is simple extension of Lipman et al. 2023 Tong al. 2024 which proved LCFM and LFM are equal under similar constraint.
Given daota 4 lo. we have ux(x) = - where lo is the previous time in the time series, and t is the current time for inference. For the time series data, we are assuming it to be Lipschitz continuous there exist
such that for all x, R", If(x) f(y) S L||x yll-
Vo (o)))ttt)) (t) ux(x)]1 = (e.((s) (1 t)2 (t, x1 1|
(18)
1 [olLea((hh (1 (1- ( 2 +
(19)
( 2
1
Voleeg(s)
(20)
By bilinearity of the 2-norm and since æ1 is independent of 0. Next,
/ x1)
E (x) (1 - 1)2 (t.t x)11
//
= 1utt)
Finally,
/ (b),, I [=12((a())9(2ddd2
Ep,(.) (t),
(e(x)d
(e(x)
/ ab((t / =1p((=l2)((2)d)) dx
// (ab 22,, 21) (c(=)))(())ddd
(x =) (i) x)tt))),z))
Where we first substitute then change the order integration for the final equality. Since all times t the gradients of LFM and LTFM equal, VoLFM VoLTFM
by substitution.
(eadlssmdeyllllottt.)) 2 E, (x ([t] t)2
1 21 (t,
(21)
s|r] (t, x)) -
I1 r
Eg (esttt) -c ((.(()))
2=(2)2)112
(22)
[t]
[t]
[[t] t)z / [c] (t, x) x 11
(e.(())
(23)
16

---

Lemma 3.1
The SDE dz, ur(z|2)dt + dW,, where de is defined in eq. generates pa (2(22) in with initial condition Po := òn, where ó the Dirac delta function.
eq. g
Proof. For simplicity of notation we first show the case where [t] 1.
dx. - (z|=))) o'd,, 1 1-- dl dW,
(24)
which is equivalent to the d dimensional Brownian bridge which has marginal
(V1( - (1)o + Lx,, 02 ())
(25)
which completes the proof for [t]
Proposition 3.2 (Coupling Preservation) Under mild regulatory criteria on we(-)z), Pr. and q, if
(ea (F, (u(
and (z(22), and (z(22) are as defined in eqs. then II(u)* = (x1.r).
Proof. We prove the deterministic case with T 1. The extensions to stochastic and T are evident. The couplings are equal if the marginal vector field un (Z+]c) = u everywhere as the coupling is governed by the push forward flows Φ(20, c) = So ue (xxlcc)dd and ((z0,c,2) fo (u,(0,12,)) c). If
(.--ur......sogg)e)eedddd.eed... ( c) u. (x.
then @( c, 2) o ( c) for all zo and therefore the couplings of the optimal map are equivalent. We note that this requires exchange of integrals under the same conditions as Lemma A.1.
Next show how (A1)-(A3) satisfy Prop. 3.2.
(A1) When zo and there exists T' X X' such that T(m) = æi if and only if 11" (20, 21). We note that this is equivalent to asserting the existence of Monge map T for the coupling II*
In the two timepoint case, c 2o is sufficient as long as there aren't two trajectories that have the same zo but different 11s. Conditioning this way ensures the conditions of of Prop. 3.2 as the uniqueness property ensures the uniqueness uy (x)cc).
(A2) There exist two trajectories x², x² such that xi x for h + consecutive observations. In this case notice that this is simply multi-timepoint extension of A1 Tt-h 1:t-1, i.c. conditioned on a history of length h. If this is the case then the same reasoning as A1 applies.
(A3) Trajectories associated with unique conditional vectors independent of t. This satisfies Prop 3.2 by definition.
Proposition 3.3 There exists scaling function c(t) R__ R such that Lungel (0) = c(t) Cmunchi (0).
Proof. We start with the matching loss.
1 Besnts) us (x|=))11 (xo(() ([t) t)2 |1= (t, - x (26)
by substitution,
x1 (t, xx)
E, (estess - ((
(27)
[t] -
[t]
([t) - - |// (t, x m
1
(28)
completing the proof.
17

---

0.08
ICU Sepsis
ICU Cardiac Arrest
Density
0.06
ICU GIB
0.04
0.02
0.00
50
100
150
200
Complete Measurements per Trajectory
ICU Sepsis
0.020
ICU Cardiac Arrest
Density
0.015
ICU GIB
ED GIB
0.010
0.005
0.000
100
150
200
250
Heart Rate (BPM)
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.
B Experimental Details
B.1
1D Oscillators
The three oscillation trajectories correspond to c 0.25 (the red trajectory in Figure 2) = 2 (blue), and 3.75 (green). Before used as input, scaled to between 0 and 1 by dividing by 10.
B.2
Clinical Datasets
B.2.1
Clinical Data Characteristics
In order to accurately model the perturbations in the physiologic signals (mean arterial pressure and heart rate) of the underlying patient states, need leam beyond the general trend of the data.
While the physiologic measurements themselves reflect patient status and drive clinical decision making, the degree of variation holds information that goes beyond the snapshot at a single time point. Our approach models the data distribution and stochasticity rather than just fitting the average trajectory. Other time-varying such as treatment conditions and non-time-varyinn covariates such as underlying disease states may also hold information that may impact the underlying state generating the physiologic signals. Our approach also incorporates this information to inform the trajectory modeling.
The data distribution in the ICU datasets reflect its status as the most resource-intensir clinical setting with increased measurement frequency and data distribution shift towards more abnormal physiologic values (Figure 4. The ED dataset reflects its status the clinical setting focused on triaging patients, with sparser and physiologic measurements that fall in the normal range.
B.2.2 Clinical Data Preprocessing
For each clinical dataset, we modeled patient trajectories formed with heart rate and blood pressure measurements during the first 24 hours following admission. The timeline for each trajectory. originally minutes, was scaled to range between 0 and by dividing by 1440. Additionally, heart rate and blood pressure values were z-score normalized to standardize the data.
Intensive Care Unit Sepsis (ICU Sepsis) Dataset The eICU Collaborative Research Database v2.0 [Pollard et al. 2019 is a database including deidentified information collected from over 200,000 patients in multiple intensive care units (ICUs) in the United States from 2014 to 2015. The ICU Sepsis Dataset was created by subsetting the eICU Database for 3362 patients with sepsis as the
18

---

primary admission diagnosis (2689 patients in training set, 336 in validation set, and 337 in test set). The following data fields were extracted: patient sex, age, heart rate, mean arterial pressure. norepinephrine dose and infusion rate, and a validated ICU score (APACHE-IV), Each patient's complete pair measurements of heart rate and mean arterial pressure over time form one trajectory to be modeled.
Norepinephrine infusion rates were calculated by converting drug doses or infusion rates to pg/kg/min, and where drug doses were not explicitly available, the dose was inferred from the free text given in the drug name. Start and end times for norepinephrine infusion were calculated by dividing the dose by the infusion rate. Where there appeared to be multiple infusions at the same time, the maximum infusion rate was taken as the infusion rate. As conditional input to the models, the norepinephrine infusion doses are then scaled between and by dividing by the maximum norepinephrine value in the dataset.
The APACHE-IV score, a validated critical care risk score, predicts individual patient mortality risk Zimmerman et al., 2006]. In data preprocessing, we uses logistic regression of the score against binary hospital mortality data to generate a probability for each patient, serving as an additional input condition for models.
Intensive Care Unit Cardiac Arrest (ICU Cardiac Arrest) Dataset This dataset was extracted from the to eICU Collaborative Research Database v2.0 [Pollard et al., [019] described above to reflect ICU patients at risk for cardiac arrest. This dataset excludes patients who presented with myocardial infarction (MI) and includes variables used in the Cardiac Arrest Risk Triage (CART) score [Churpek et al.) 2012]: respiratory rate, heart rate, diastolic blood pressure, and age at the time of ICU admission. As an input to the model, the age was 7-score normalized. 51671 patients were included in the training set, with 6459 patients each in the validation and test sets.
Intensive Care Unit Acute Gastrointestinal Bleeding (ICU GIB) Dataset The Medical Infor- mation Mart for Intensive Care III (MIMIC-III critical care database contains data for over 40,000 patients in the Beth Israel Deaconess Medical Center from 2001 to 2012 requiring an ICU stay (Johnson et al. 2016]. We selected a cohort of 2602 ICU patients with the primary diagnosis of gastrointestinal bleeding to form the ICU GIB dataset, split into training set of 2082 patients, and validation set and test set of 260 patients each. We extracted the following variables: age, sex, heart rate, systolic blood pressure, diastolic blood pressure, usage of vasopressor, usage of blood product, usage of packed red blood cells, and liver disease. Since the vasopressor and blood product usage are encoded as a binary value and may not represent actual infusion amount that are most likely decaying, experimented with adding a Gaussian decay to them to use as conditional inputs. Likewise, trajectories to model consist of complete pairs of heart rate and mean arterial pressure (calculated from systolic blood pressure and diastolic blood pressure) measurements.
Emergency Department Acute Gastrointestinal Bleeding (ED GIB) Dataset This dataset reflects 3348 patients presenting with signs and symptoms of acute gastrointestinal bleeding to two hospital campuses in Yale New Haven Hospital between 2014 and 2018. The patients were split into training set, a validation set, and a test set of 2636, 352, and 360 patients. Variables extracted include patient sex, age, heart rate, mean arterial pressure, initial measurements of 24 lab tests, and 17 pre-existing medical conditions as determined by ICD-10 codes. Like ICU Sepsis data, the trajectoires consist of complete pairs of heart rate and mean arterial pressure measurements.
Age. initial lab test measurements (three labs omitted due to missing data), and pre-existing medical conditions were used to train an XGBoost model [Chen and Guestrin 2016) to predict the binary outcome variable indicating the need for hospital- based care. The resulting probabilities of requiring hospital- -based care (outcome of 1) for each patient were then calculated using the trained model and used conditional input conditional models in experiments this dataset.
Of note, the outcome variable was defined as 1 if patient (1) requires red blood cell transfusion, (2) requires urgent intervention (endoscopic, interventional radiologic, or surgical) to stop bleeding or (3) all-cause 30-day mortality. Labs and medical conditions included in this dataset are listed below. Labs in bold were excluded from the XGBoost risk score calculation due to missing data.
Labs: Sodium, Potassium, Chloride, Carbon Dioxide, Blood Urea Nitrogen, Creatinine, International Normalized Ratio, Partial Thromboplastin Time, White Blood Cell Count, Hemoglobin, Platelet Count, Hematocrit, Mean Corpuscular Volume, Mean Corpuscular Hemoglobin, Mean Corpuscular Hemoglobin Concentration, Red Cell Distribution Width,
19

---

ICU Sepsis
ICU GIB
Mean MSE Test
x -0-1
10-3
Model
TFM-ODE
TF TFM-ODE-Memory
10
3
0.0
0.2
0.4
0.0
0.2
0.4
sigma
sigma
Figure 5: Sigma mean MSE comparison
Red Blood Cell Count, Aspartate Aminotransferase, Alanine Aminotransferase, Alkaline Phosphatase, Total Bilirubin, Direct Bilirubin, Albumin, Lactate.
Previous Medical Histories: Charlson Comorbidity Index, Cerebrovascular Accident, Deep Vein Thrombosis, Pulmonary Embolism, Atrial Fibrillation, Upper Gastrointestinal Bleeding, Lower Gastrointestinal Bleeding, Unspecified Gastrointestinal Bleeding, Peptic Uleer Disease, Helicobacter Pylori Infection, Coronary Artery Disease, Heart Failure, Hypertension, Type 2 Diabetes Mellitus, Chronic Kidney Disease, Alcohol Use Disorder, Cirrhosis.
B.3 Training
B.3.1
1D Oscillators
Since the trajectories in this dataset are deterministic and regularly sampled, we deployed only TFM-ODE and applied solely the Lmatc loss (i.e, no uncertainty or time predictive loss), as these methods sufficiently address the structured nature of the data to generate proof-of-coneee
results.
The three models presented Figure 2 have hidden size 256, a of 0.1, trained under seed=0) with Adam optimizer with learning rate 10-3 3 for a maximum of 1000 epochs with early stopping (patience=3) monitoring validation loss.
B.3.2 Clinical Data
All the models for clinical data experiments are trained with Adam optimizer. maximum training time and epochs are to 48 hours and 300, with early stopping (patience=3) monitoring validation loss. All metrics reported were ran with 5 seeds (0,1,2,3,4 to ensure it is reproducible.
TFM, TFM-ODE, and ablations The TFM models were trained with learning rate x 10 and had a of 0. The complete models have hidden size 256 and memory of 3, while ablation study with hidden size of 64 and/or no memory was performed (Table The noise parameter for the SDE implementation was
0.1 for ablations without Luncertainty- The hyperparameters a and memory=3 for full models were selected through experiments with different values of and memory (Figure and and 6
FM The FM baseline models were trained with learning rate x 10-6, All models had a hidden size of 64 with of 0.1.
Latent Neural ODE The latent Neural ODE models were trained with a learning rate of 10-3 100 GRU units used for the encoder model and the number of latent dimensions was 2.
Baseline Neural SDE and Neural ODE Both baseline models were trained with learning rate 1 x 10 and had a hidden size of 64,
20

---

ICU Sepsis
ICU GIB
x 10
5
N
Model model_ty _type TFM-ODE TFM
Mean MSE Test
3
4 10
2x10-
64
3 x
256
10
5 memory
memory
Figure 6: Memory Mean MSE comparison
NeurIPS Paper Checklist
Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
The abstract and introduction clearly state the claim. The theoretical portions of the claim supported in sections 2 and of the main text and in appendix A; the claims on its performance are supported by section of the main text. We discuss the limitations our models context of the ablation study and in the related works section. See the next section for more details on the specifics on the discussion of limitations.
Justification:
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made in the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No NA answer to this question will not perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize other settings.
is fine to include aspirational goals as motivation as long as is clear that these goals are not attained by the paper.
2. Limitations
Question:
Does the paper discuss the limitations of the work performed by the authors?
Answer: [Yes]
Justification: We discuss the limits of our current model depending characteristics of the dataset. the ablation study, we discuss the variable effect memory had across datasets of different data distributions and measurement frequency framework. We also provide context in the related work section that describes the limitation of our current framework in not being able to provide causal explanations.
limitation to the current
Guidelines:
The answer NA means that the paper has no limitation while the answer No means that paper has limitations, but those are not discussed in the paper.
The authors are encouraged to create a separate "Limitations" section their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification. asymptotic approximations only holding locally). The authors
21

---

should reflect on how these assumptions might be violated in practice and what the implications would be.
The authors should reflect on the scope of the claims made, e.g., if the approach was only tested few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution is low images are taken in low lighting. Or speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness.
While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an impor- tant role in developing norms that preserve the integrity of the community, Reviewers will be specifically instructed to not penalize honesty concerning limitations.
3. Theory Assumptions and Proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
Answer: [Yes]
Justification: Proof of theorems are included in section 3 and appendix A.
Guidelines:
The answer NA means that the paper does include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated or referenced the statement of any theorems.
The proofs can either appear in the main paper or the supplemental material, but they appear in the supplemental material, the authors are encouraged provide short proof sketch to provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental Result Reproducibility
Question: Does the paper fully disclose all the information needed to reproduce the main ex- perimental results the paper to the extent that affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?
Answer: [Yes]
Justification: Reproducibility is accomplished by providing a full description of the model the main text addition to providing publically available code at /tttp::/// thub The experimental data are available publicly online (ICU Sepsis, ICU GIB) or upon reasonable request (ED GIB) subject to institutional approval and HIPAA regulations. Data preprocessing steps are extensively outlined in Appendix B.2.
Guidelines:
The answer NA means that the paper does not include experiments.
If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data provided or not.
22

---

If the contribution is dataset and/or model, the authors should describe the steps taken make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution a specific model and empirical evaluation, may be necessary to either make it possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility instructions for how to replicate the results, access to a hosted model (c.g., in the case of large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.
also be provided via detailed
While NeurIPS does not require releasing code, the conference does require all submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) If the contribution is primarily a new algorithm, the paper should make it clear how reproduce that algorithm. (b) If the contribution is primarily
model architecture, the paper should describe the architecture clearly and fully.
(c) If the contribution is a new model (e.g.. a large language model), then there should either be way to access this model for reproducing the results or way to reproduce the model (e.g., with an open-source dataset instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case authors welcome describe the particular way they provide for reproducibilityy In the case of closed-source models, it may be that access to the model is limited in some way (e.g., registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results.
5. Open access to data and code
Question: Does the paper provide open access to the data and code, with sufficient instruc- tions to faithfully reproduce the main experimental results, as described in supplemental material?
Answer: [Yes]
Justification: Anonymised code is available as supplementary material. The ICU Sepsis dataset is available publicly online via the eICU database and the ICU GIB dataset is available via the MIMIC-III database, whereas the ED GIB datasets are datasets which include identifiable information. A deidentified dataset may be made available upon reasonable request subject to guidelines set by the institutional review board and in accordance with HIPAA policy.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurIPS code and data submission guidelines (https://nips.cm cc/ sablic/guidese https///oootttbaiiiiooonoooooooo
for more details.
While encourage the release of code and data, we understand that this might not be possible, so "No" is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurlPS code and data submission guidelines (https
//nips ctips..c/public/guiii hwwpoo/lcgggiiiee/eeeooo http://udddee/oeeSubmmmm
wwwssssosliiioooeoooo. for more details.
The authors should provide instructions on data access and preparation, including how access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
23

---

Providing as much information as possible in supplemental material (appended to the paper) recommended, but including URLs to data and code permitted.
6. Experimental Setting/Details
Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Question:
Answer: [Yes]
Justification: We provide an overview with the significance of each clinical dataset in the body of the paper. In the Appendix we have written detailed descriptions of each dataset, including the methodology for identifying the cohorts of patients, input variables extracted, number of patient encounters, summary statistics, well as frequency distributions. We also detail the process of training and testing for methods presented in the paper, from baseline to our novel method. We have also provided anonymized link to all code used for the experiments.
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core of the paper to level of detail that is necessary to appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, or as supplemental material.
7. Experiment Statistical Significance
Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Answer: [Yes]
Justification: The paper includes standard error of the mean for all the results reported in numerical format. There are no figures with error bars for which the nature of error bar calculation would be relevant. There are no statistical tests that were perform which would make discussion of the validity of statistical tests relevant.
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by enror bars, confi- dence intervals, or statistical significance tests, least for the experiments that support the main claims of the paper.
The factors of variability that the error bars capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall
with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)
The assumptions made should be given (e.g.. Normally distributed errors).
It should be clear whether the error bar is the standard deviation the standard error of the mean.
is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out of range (e.g. negative еттог rates).
If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text.
8. Experiments Compute Resources
Question: For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time of execution) needed reproduce the experiments?
24

---

Answer: [Yes]
Justification: We provide the description of the compute resources utilized with available computing clusters detailed in the Experimental Results section.
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU or GPU, internal cluster, cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't 't make it into the paper).
9. Code Of Ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurlPS Code of Ethics http:///wwwpppc.pppp/ooaaaaacaaa
Answer: [Yes]
Justification: We have all reviewed the NeurIPS Code of Ethics and striven maintain and preserve anonymity.
Guidelines:
The answer NA means that the authors have reviewed the NeurIPS Code of Ethics.
If the authors answer No, they should explain the special circumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g.. if there is special consid- eration due to laws or regulations their jurisdiction).
10. Broader Impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [Yes]
Justification: In the conclusion section we discuss how improved modeling of Emergency Department and ICU physiology can lead the positive societal impact of improving disease risk prediction which can enable medical teams to make better clinical decisions, provide patients and family members with increased information on the likely course the illness, and improve resource allocation by identifying which patients may require costly or rare resources such a blood transfusions.
Guidelines:
The answer NA means that there is no societal impact of the work performed.
If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (c.g., deployment of technologics that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
The conference expects that many papers will be foundational research and tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology being used as intended but gives incorrect results, and harms following from (intentional unintentional) misuse of the technology.
25

---

If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition attacks. mechanisms for monitoring misuse, mechanisms to monitor how system learns from feedback over time, improving the efficiency and accessibility of MI.).
11. Safeguards
Does the paper deseribe safeguards that have been put in place for responsible release of data or models that have high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
Justification: The current study proposes predictive model for time-varying datasets and as such does not pose risks that are significantly higher in comparison to conventional models that fit data trends. There is generative component in model, and little risk that the results could be misused in way that misleads or proves otherwise detrimental to the broader public. We therefore believe this item does not apply for this submission. Guidelines:
The answer NA means that the paper poses such risks.
Released models that have high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make best faith effort.
12. Licenses for existing assets
Are the creators or original owners of assets (e.g.. code, data, models), used in the paper, properly credited and the license and terms of use explicitly mentioned and properly respected?
Question:
Answer: [Yes]
The publicly available datasets from MIMIC-III and eICU databases are properly credited, respected, mentioned and used under the PhysioNet Credentialed Health Data License Version 1.5.0. The ED GIB EHR data is owned by the authors.
Justification:
Guidelines:
The answer NA means that the paper does not use existing assets.
The authors should cite the original paper that produced the code package or dataset.
The authors should state which version of the asset is used and, if possible, include URL.
The name of the license (e.g., CC-BY 4.0) should be included for each asset.
For scraped data from a particular source (e.g.. website). the copyright and terms of service of that souree should be provided.
assets are released, the license, copyright information, and terms use in the package should be provided. For popular datasets, paperswi thcode. com com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
For existing datasets that are re- ee-pacaaedd both the original license and the license of the derived asset (if it has changed) should be provided.
If this information is not available online, the authors encouraged to reach out to the asset's 's creators.
13. New Assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
26

---

Answer:
[No]
Justification:
We will not be releasing new datasets.
Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details the dataset/coode/modll as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent was obtained from people whose asset is used.
At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.
14. Crowdsourcing and Research with Human Subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well as details about compensation (if any)?
Answer: [NA]
Justification: While the data involves patient data, does not include any experiments performed on humans but instead based the use of retrospectively collected healthcare records generated as part of routine clinical care. Publicly available datasets were used in accordance in the regulations set out by the hosting institution, whereas collection of data from patients in non-publicly available datasets were performed in accordance with institutional review board guidance and HIPAA data protection regulations. As the data collection does not involve experiments on human subjects or crowdsourcing experiments, this item does not apply to our current study.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Including this information in the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then as much detail as possible should be included in the main paper.
According to the NeurIPS Code of Ethics, workers involved in data collection, curation,
or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?
Answer: [Yes]
Justification: We have obtained institutional IRB consent for the use of the ED GIB EHR data under an approved protocol that we are happy to provide upon request. MIMIC-III and eICU databases are pre-approved, de-identified, publicly available clinical data sources.
Guidelines:
The answer NA means that the paper does not involve crowdsoureing nor research with human subjects.
Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere the NeurIPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applicable), such the institution conducting the review.
27


Reviews:
[{"review_id": "50HpVNDuui", "reviewer": "Reviewer_TMfk", "summary": "This paper presents Trajectory Flow Matching (TFM), a simulation-free training algorithm for neural differential equation models. This enables modeling of continuous physiologic processes using irregular, sparsely sampled, and noisy data, all with better scalability.\n\nThe authors provide theoretical proof that matching techniques can allow for simulation-free training, and empirical proof of the effectiveness of TFM in clinical settings. An ablation study is also provided to validate components of TFM.", "questions": "- Did you compare TFM-ODE with other models that are able to process irregular clinical time series?\n- What is your reasoning behind modeling specifically the heart rate + MAP pair ?", "limitations": "The authors have adequately addressed the technical limitations (memory not always suited depending on the dataset, no causal representations estimations) as well as societal impacts (false predictions).", "rating": 7, "confidence": 3, "soundness": 4, "presentation": 3, "contribution": 3, "strengths": "- Originality : Bridging the gap between flow matching and differential equation-based models is a novel idea that is carried out rigorously by the authors.\n- Quality : Proofs are provided for the theoretical results; experiments are performed on three different real-world datasets; uncertainty prediction is provided by design and helps clinical usability. Responses to the NeurIPS checklist are also extensive.\n- Clarity : The paper is easy to parse, as all ideas follow each other logically.\n- Significance : SDEs are notoriously expensive to train and this work promises to solve this issue while offering state-of-the-art performance, by a notable margin.", "weaknesses": "- The motivation for the paper is that SDEs are not scalable. While this is a known fact, the paper lacks a proper discussion of the impact of simulation-free training on the scalability of TFM.\n- The provided code lacks instructions to reproduce the results (empty README). \n- Explanability studies would have been welcome to help clinical usability.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their valuable feedback that helps us further improve our paper. We have added a README, and impact statement on simulation-free training scalability in our manuscript. \n\nWe agree having explainability is important for others who may apply our method. We thank the reviewer for bringing this point up as it is quite important in high-stake applications such as healthcare. As such, we have included this in our future work section. \n\nWe would like to address some of the reviewers questions below:\n\n> Did you compare TFM-ODE with other models that are able to process irregular clinical time series?\n\nYes, we compared TFM-ODE’s performance to NeuralODE [1], FM baseline ODE [2], and LatentODE RNN [3], as well as SDE-based models like TFM and NeuralSDE [4]. All these models are able to process irregularly sampled time series. (See Table 1)\n\n> What is your reasoning behind modeling specifically the heart rate + MAP pair ?\n\nWe chose the heart rate + MAP pair since they are the most obtainable and non-invasive metric from a patient that provide a clinically relevant representation of patient state trajectories in the management of clinical conditions where hemodynamic monitoring is important (e.g., sepsis and GI bleeding). We chose MAP since it reflects both cardiac output and peripheral vascular resistance; this is used for decisions regarding resuscitation and vasopressor treatment, which are indications for hospital-based treatment. Heart rate may be affected by either cardiac arrhythmias or clinical conditions related to but independent from cardiac output and peripheral vascular resistance. Therefore we used a combination of the two to reflect clinical trajectories.\n\nWe hope we have addressed the questions of the reviewer and we would like to thank the reviewer once more for helping us improve the clarity of our responses and propose future directions.\n\n[1] Chen, R. T. Q., Rubanova, Y., Bettencourt, J., & Duvenaud, D. (2018). Neural Ordinary Differential Equations (Version 5). arXiv. `https://doi.org/10.48550/ARXIV.1806.07366`\n\n[2] Lipman, Y., Chen, R. T. Q., Ben-Hamu, H., Nickel, M., & Le, M. (2022). Flow Matching for Generative Modeling (Version 2). arXiv. `https://doi.org/10.48550/ARXIV.2210.02747`\n\n[3] Rubanova, Y., Chen, R. T. Q., & Duvenaud, D. (2019). Latent ODEs for Irregularly-Sampled Time Series (Version 1). arXiv. `https://doi.org/10.48550/ARXIV.1907.03907`\n\n[4] Liu, X., Xiao, T., Si, S., Cao, Q., Kumar, S., & Hsieh, C.-J. (2019). Neural SDE: Stabilizing Neural ODE Networks with Stochastic Noise (Version 1). arXiv. `https://doi.org/10.48550/ARXIV.1906.02355`", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "rtun6X9b3Y", "reviewer": "Reviewer_s476", "summary": "This paper presents Trajectory Flow Matching, an extension of flow matching to time series. It can model irregular, sparsely sampled, and noisy time series. It trains in a simulation-free manner, bypassing backpropagation through the dynamics. The method is tested on ICU physiological time series, demonstration SoTA performance, and uncertainty prediction.", "questions": "1) It would be interesting to understand if this could be applied to periodic time series like ECG. Do you see ways to introduce this behaviour, and if so, could you call this out as a future research direction?", "limitations": "The authors clearly describe the limitations.", "rating": 10, "confidence": 4, "soundness": 4, "presentation": 4, "contribution": 4, "strengths": "1) A strong research contribution, extending flow matching to time series.\n2) Explicit modelling of uncertainty, handles noisy, irregular, and sparsely sampled data.\n3) Application to a domain with significant societal benefit.\n4) Clearly defined paper, easy to follow, and well-written.", "weaknesses": "None of particular note.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We sincerely thank the reviewer for their extremely positive response to our work. We are thrilled that you found our TFM to be a strong research contribution. Additionally, we are grateful for your suggestion regarding the future application of our model to periodic time series. We have some preliminary thoughts on this particular behavior that we would like to explore, such as modeling in the Fourier domain, are eager to explore them in future work, and have added this to the future work section of our manuscript.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "OzEpqHEwIn", "reviewer": "Reviewer_4EUb", "summary": "This paper proposes a trajectory flow matching method for time series data, by using flow matching at each time points. To preserve the coupling of time series, the vector fields are conditioned on history lengths (or even more general $c$). The method also provides ways for model stability, irregularly sampled trajectories and uncertainty prediction. After evaluating the model performance on simple harmonic oscillators, they further apply their method to three clinical time series datasets.", "questions": "The paper is clearly written. No major questions.", "limitations": "The same as weakness", "rating": 7, "confidence": 4, "soundness": 4, "presentation": 3, "contribution": 3, "strengths": "The proposed method provides a generative model for time series data, and it takes the benefits of FM, in terms of flexibility, scalability and stability, for time series data modeling.", "weaknesses": "The TFM is the model for discrete time series. It would potentially be more useful to model things in a continuous way.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their positive assessment of our paper. We appreciate the reviewer’s insight in the importance of continuous time series modeling. We are also interested in the fully continuous setting and would be interested in adapting ideas from functional flow matching [1] to this domain.\n\n[1] Kerrigan, G., Migliorini, G., & Smyth, P. (2023). Functional Flow Matching (Version 2). arXiv. `https://doi.org/10.48550/ARXIV.2305.17209`", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:
# Global Response

We thank the reviewers for their time evaluating our paper. We are grateful for the thoughtful comments, insights, and potential directions for future work. We have addressed each of the points raised and provided clarifications where necessary. Based on suggestions from the reviewers we have made the following improvements:

* Included README.md with instructions on running the code as well as other code quality improvements. (reviewer TMfk)
* Expanded impact statement of simulation-free training and scalability of TFM to include discussion of scalability relative to simulation-based methods. (reviewer TMfk)
* Expanded discussion of future directions to include modeling periodic time series, functional flow matching for continuous time series, and explainability. (reviewer s476, 4EUb, TMfk)

We look forward to any further feedback and discussions and appreciate the opportunity to improve our work based on your valuable input.

Meta-review: The paper presents a significant advancement by bridging flow matching and differential equation-based models, a novel and rigorous approach that is well-supported by theoretical proofs and extensive experiments across multiple real-world datasets. The method’s ability to provide uncertainty prediction and its potential to address the scalability issues of stochastic differential equations (SDEs) while delivering state-of-the-art performance highlight its originality and practical significance. The clarity of the paper, with its logically flowing ideas, and its application to a domain with notable societal benefits further underscore its strengths.

All reviewers are highly supportive of this paper. It stands out by combining substantial algorithmic and theoretical contributions with practical significance.
Decision: Accept (spotlight)