SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Table 2: Test errors in the irregular synthetic benchmarks, estimated over 5 seeds and 1000 test trajectories per seed, with standard deviation calculated across seeds.
Model	Complex dynamics eigenvalues	Real dynamics eigenvalues
MSE	OOD MSE	MSE	OOD MSE
LSTM	0.23 ± 0.001	0.589 ± 0.02	0.381 ± 0.002	2.354 ± 0.84
GRU-ODE-Bayes	0.182 ± 0.0004	0.361 ± 0.044	0.219 ± 0.0004	0.355 ± 0.005
CRU	0.233 ± 0.0054	0.584 ± 0.009	0.231 ± 0.001	0.541 ± 0.026
NESDE (ours)	0.176 ± 0.0001	0.178 ± 0.001	0.222 ± 0.0005	0.332 ± 0.005
As demonstrated below, the unique architecture of NESDE provides effective regularization and data efficiency (due to piecewise linearity), along with tunable expressiveness (neural updates with controlled frequency). Yet, it is important to note that the piecewise linear SDE operator does limit the expressiveness of the model (e.g., in comparison to other neural-ODE models). Further, NESDE is only optimal under a restrictive set of assumptions, as specified in Proposition 1.
5 Synthetic Data Experiments
In this section, we test three main aspects of NESDE: (1) prediction from partial and irregular observations, (2) robustness to out-of-distribution control (OOD), and (3) sample efficiency. We experiment with data of a simulated stochastic process, designed to mimic partially observable medical processes with indirect control.
The simulated data includes trajectories of a 1-dimensional signal  $Y$ , with noiseless measurements at random irregular times. The goal is to predict the future values of  $Y$  given its past observations. However,  $Y$  is mixed with a latent (unobservable) variable, and they follow linear dynamics with both decay and periodicity (i.e., complex dynamics eigenvalues). In addition, we observe a control signal that affects the latent variable (hence affects  $Y$ , but only indirectly through the dynamics). The control negatively correlated with the observations:  $u_{t} = b_{t} - 0.5\cdot Y_{t}$ .  $b_{t}\sim U[0,0.5]$  is a piecewise constant additive noise (changing 10 times per trajectory)
(a) Same control distribution (b) Out of distribution control
Figure 5: MSE vs. number of observations so far in the trajectory, in the complex dynamics setting, for: (a) standard test set, and (b) test set with out-of-distribution control policy.  $95\%$  confidence intervals are calculated over 5 seeds.
As baselines for comparison, we choose recent SDE-based methods that provide Bayesian uncertainty estimation: GRU-ODE-Bayes [De Brouwer et al., 2019] and CRU [Schirmer et al., 2022] (despite terminology, both implement an SDE rather than ODE). In these methods, concatenating the control signal to the observation results in poor learning, as the control becomes part of the model output and dominates the loss function. To enable effective learning for the baselines, we mask-out the control from the loss. Additionally, we design a dedicated LSTM model that supports irregular predictions, as described in Appendix F.2.
Out-of-distribution control (OOD): We simulate two benchmarks – one with complex eigenvalues and another with real eigenvalues (no periodicity). We train all models on a dataset of 1000 random trajectories, and test on a separate dataset – with different trajectories that follow the same distribution. In addition, we use an OOD test dataset, where the control is positively correlated with the observations:  $u_{t} = b_{t} + 0.5 \cdot Y_{t}$ . This can simulate, for example, forecasting of the same biochemical process after changing the medicine dosage policy.
Table 2 and Fig. 5a summarize the prediction errors. Before changing the control policy, NESDE achieves the best accuracy in the complex dynamics, and is on par with GRU-ODE-Bayes in the real dynamics. Notice that CRU, which relies on a real-valued linear model in latent space, is indeed sub-optimal under the complex dynamics, compared to NESDE and GRU-ODE-Bayes. The LSTM presents high errors in both benchmarks.
Once the control changes, all models naturally deteriorate. Yet, NESDE presents the smallest deterioration and best accuracy in the OOD test datasets – for both complex and real dynamics. In particular, NESDE provides a high prediction accuracy after mere 2 observations (Fig. 5b), making it a useful zero-shot model. The robustness to the modified control policy can be attributed to the model of NESDE in Eq. (6), which decouples the control from the observations.
In a similar setting in Appendix E.6, the control  $u$  used in the training data has continuous knowledge of  $Y$ . Since the model only observes  $Y$  in a limited frequency,  $u$  carries additional information about  $Y$ . This results in extreme overfitting and poor generalization to different control policies – for all methods except for NESDE, which maintains robust OOD predictions in this challenging setting.
Sample efficiency: We train each method over datasets with different number of trajectories. Each model is trained on each dataset separately until convergence. As shown in Fig. 6, NESDE achieves the best test accuracy for every training dataset, and learns reliably even from as few as 100 trajectories. The other methods deteriorate significantly in the smaller datasets. Note that in the real dynamics, LSTM fails regardless of the amount of data, as also reflected in Table 2.
GRU-ODE-Bayes achieves the best sample efficiency among the baselines. In Appendix E.2, we use a benchmark from the study of GRU-ODE-Bayes itself [De Brouwer et al., 2019], and demonstrate the superior sample efficiency of NESDE in that benchmark as well. Appendix E.3 extends the notion of sample efficiency to sparse trajectories: for a constant number of training trajectories, it reduces the number of observations per trajectory. NESDE demonstrates high robustness to the amount of data in that setting as well.
(a) Complex dynamics
(b) Real dynamics
Figure 6: Test MSE vs. train data size.  $95\%$  confidence intervals are calculated over 1000 test trajectories.
Regular LSTM: Appendix E.4 extends the experiments for regular data with constant time-steps. In the regular setting, LSTM provides competitive accuracy when observations are dense. However, if the signal is not observed every time-step, LSTM deteriorates significantly in unobserved time-steps (Fig. 12c), possibly because gradients have to be propagated over many steps. Hence, even in regular settings, LSTM struggles to provide predictions more frequent than the measurements.
6 Medication Dosing Regimes
As discussed in Section 1, many medical applications could potentially benefit from ODE-based methods. Specifically, we address medication dosing problems, where observations are often sparse, the dosing is a control signal, and uncertainty estimation is crucial. We test NESDE on two such domains, against the same baselines as in Section 5 (GRU-ODE-Bayes, CRU and an irregular LSTM). We also add a naive model with "no-dynamics" (predicts the last observed value).
The benchmarks in this section were derived from the MIMIC-IV dataset [Johnson et al., 2020]. Typically to electronic health records, the dataset contains a vast amount of side-information (e.g., weight and heart rate). We use some of this information as an additional input – for each model according to its structure (context-features for the hyper-network of NESDE, covariates for GRUODE-Bayes, state variables for CRU, and embedding units for the LSTM). Some context features correspond to online measurements and are occasionally updated. In both domains, we constraint the process eigenvalues  $\lambda$  to be negative, to reflect the stability of the biophysical processes. Indeed, the spectral representation of NESDE provides us with a natural way to incorporate such domain knowledge, which often cannot be used otherwise. For all models, in both domains, we use a 60-10-30 train-validation-test data partition. See more implementation details in Appendix F.
6.1 Unfractionated Heparin Dosing
Unfractionated Heparin (UH) is a widely used anticoagulant drug. It may be given in a continuous infusion to patients with life-threatening clots and works by interfering with the normal coagulation cascade. As the effect is not easily predicted, the drug's actual activity on coagulation is traditionally
Table 3: Test mean square errors (MSE) and negative log-likelihood (NLL, for models that provide probabilistic prediction) in the medication-dosing benchmarks.
Model	UH Dosing	Vancomycin Dosing
MSE	NLL	MSE	NLL
Naive	613.3 ± 13.48	-	112.2 ± 16.4	-
LSTM	482.1 ± 6.52	-	92.89 ± 11.3	-
GRU-ODE-Bayes	491 ± 6.88	4.52 ± 0.008	80.54 ± 11.8	6.38 ± 0.12
CRU	450.4 ± 8.27	4.49 ± 0.012	76.4 ± 12.8	3.87 ± 0.2
NESDE (ours)	411.2 ± 7.39	4.43 ± 0.01	70.71 ± 12.3	3.69 ± 0.13
monitored using a lab test performed on a blood sample from the patient: activated Partial Thromboplastin Time (aPTT) test. The clinical objective is to keep the aPTT value in a certain range. The problem poses several challenges: different patients are affected differently; the aPTT test results are delayed; monitoring and control are required in higher frequency than measurements; and deviations of the aPTT from the objective range may be fatal. In particular, underdosed UH may cause clot formation and overdosed UH may cause an internal bleeding [Landefeld et al., 1987]. Dosage rates are manually decided by a physician, using simple protocols and trial-and-error over significant amounts of time. Here we focus on continuous prediction as a key component for aPTT control.
Following the preprocessing described in Appendix F.1, the MIMIC-IV dataset derives 5866 trajectories of a continuous UH control signal, an irregularly-observed aPTT signal (whose prediction is the goal), and 42 context features. It is known that UH does not affect the coagulation time (aPTT) directly (but only through other unobserved processes, Delavenne et al. [2017]); thus, we mask the control mapping  $B$  to have no direct effect on the aPTT metric, but only on the latent variable (which can be interpreted as the body UH level). The control (UH) and observations (aPTT) are one-dimensional  $(m = 1)$ , and we set the whole state dimension to  $n = 4$ .
(a) Heparin
Figure 7: A sample of patients from (a) the UH dosing dataset, and (b) the VM dosing dataset. The lower plots correspond to medication dosage (UH in (a) and VM in (b)). The upper plots correspond to the continuous prediction of NESDE (aPTT levels in (a) and VM concentration in (b)), with  $95\%$  confidence intervals. In both settings, the prediction at every point relies on all the observations up to that point.
(b) Vancomycin
6.2 Vancomycin Dosing
Vancomycin (VM) is an antibiotic that has been in use for several decades. However, the methodology of dosing VM remains a subject of debate in the medical community [Rybak et al., 2009], and there is a significant degree of variability in VM dynamics among patients [Marsot et al., 2012]. The dosage of VM is critical; it could become toxic if overdosed [Filippone et al., 2017], and ineffective as an antibiotic if underdosed. The concentration of VM in the blood can be measured through lab test, but these tests are often infrequent and irregular, which fits into our problem setting.
Here, the goal is to predict the VM concentration in the blood at any given time, where the dosage and other patient measurements are known. Following the preprocessing described in Appendix F.1, the dataset derives 3564 trajectories of VM dosages at discrete times, blood concentration of VM  $(m = 1)$  at irregular times, and similarly to Section 6.1, 42 context features. This problem is less noisy than the UH dosing problem, as the task is to learn the direct dynamics of the VM concentration, and not the effects of the antibiotics. The whole state dimension is set to  $n = 2$ , and we also mask the control mapping  $B$  to have no direct effect on the VM concentration, where the latent variable that directly affected could be viewed as the amount of drug within the whole body, which in turn affects the actual VM concentration in the blood.
Fig. 7 displays sample trajectories predicted by NESDE in both domains. As summarized in Table 3, NESDE outperforms the other baselines in both UH and VM dosing tasks, in terms of both square errors (MSE) and likelihood (NLL). For the UH dosing problem, Fig. 8 also presents the errors
vs. prediction horizon (the time passed since the last observation). Evidently, NESDE provides the best accuracy in all the horizons. Its smoothness allows NESDE to obtain high accuracy in short horizons, and its robustness allows generalization to out-of-distribution horizons: while most of the data corresponds to horizons of 5-7 hours (see Fig. 14 in the appendix), NESDE provides reliable prediction at other horizons as well. By contrast, LSTM and GRU-ODE-Bayes have difficulty with short horizons; they only become competitive with the naive model after 6 hours. CRU provides more robust predictions, but is still outperformed by NESDE.
Despite the large range of aPTT levels in the data (e.g., the top  $5\%$  are above 100s and the bottom  $5\%$  are below 25s),  $50\%$  of all the predictions have errors lower than  $12.4s -$  an accuracy level that is considered clinically safe. Fig. 8 shows that indeed, up to 3 hours after the last lab test, the average error is smaller than 10s.
Running time: All Heparin and Vancomycin experiments were run on a single Ubuntu machine with eight i9-10900X CPU cores and Nvidia's RTX A5000 GPU. A single training process of NESDE required 8-24 hours (with variance mostly attributed to early stopping criterion). LSTM required similar running times (6-24 hours), whereas GRU-ODE-Bayes and CRU required 36-72 hours to train.
Figure 8: aPTT prediction errors in the UH problem, vs. the time passed since the last aPTT test.
Motivated by medical forecasting and control problems, we characterized a set of challenges in modeling continuous dosing dynamics: sample efficiency, uncertainty estimation, personalized modeling, continuous inference and generalization to different control policies. To address these challenges, we introduced the novel NESDE algorithm, based on a stochastic differential equation with spectral representation. We demonstrated the reliability of NESDE in a variety of synthetic and real data experiments, including high noise, little training data, contextual side-information and out-of-distribution control signals. In addition, NESDE demonstrated high prediction accuracy after as few as 2 observations, making it a useful zero-shot model.
We applied NESDE to two real-life high-noise medical problems with sparse and irregular measurements: (1) blood coagulation forecasting in Heparin-treated patients, and (2) Vancomycin levels prediction in patients treated by antibiotics. In both problems, NESDE significantly improved the prediction accuracy compared to alternative methods. This demonstrates NESDE's advantage in continuous forecasting of controlled SDE systems, which include medication dosing and controlled chemical and biological processes.
As demonstrated in the experiments, NESDE provides robust, reliable and uncertainty-aware continuous forecasting. This paves the way to development of decision making in continuous high-noise decision processes, including medical treatment, finance and operations management. Future research may address medical optimization via both control policies (e.g., to control medication dosing) and sampling policies (to control measurements timing, e.g., of blood tests).
Acknowledgements: This research was supported by VATAT Fund to the Technion Artificial Intelligence Hub (Tech.AI).
Michael Rybak, Ben Lomaestro, John C. Rotschafer, Jr. Moellering, Robert, William Craig, Marianne Billeter, Joseph R. Dalovisio, and Donald P. Levine. Therapeutic monitoring of vancomycin in adult patients: A consensus review of the American Society of Health-System Pharmacists, the Infectious Diseases Society of America, and the Society of Infectious Diseases Pharmacists. American Journal of Health-System Pharmacy, 66(1):82-98, 01 2009. ISSN 1079-2082. doi: 10.2146/ajhp080434.
Mona Schirmer, Mazin Eltayeb, Stefan Lessmann, and Maja Rudolph. Modeling irregular time series with continuous recurrent units. In International Conference on Machine Learning, pages 19388-19405. PMLR, 2022.
B. Sinopoli, L. Schenato, M. Franceschetti, K. Poolla, M.I. Jordan, and S.S. Sastry. Kalman filtering with intermittent observations. IEEE Transactions on Automatic Control, 49(9):1453-1464, 2004. doi: 10.1109/TAC.2004.834121.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N. Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need, 2017.
Yi-Jen Wang and Chin-Teng Lin. Runge-kutta neural network for identification of dynamical systems in high accuracy. IEEE Transactions on Neural Networks, 9(2):294-307, 1998.
Chao Yu, Jiming Liu, Shamim Nemati, and Guosheng Yin. Reinforcement learning in healthcare: A survey. ACM Computing Surveys (CSUR), 55(1):1-36, 2021.
A Observation Filtering: The Conditional Distribution and the Relation to Kalman Filtering 15
B Integrator Implementation 15
C The Dynamics Spectrum and Complex Eigenfunction Implementation 16
D Solver Analysis 17
E Extended Experiments 17
E.1 Ablation Study for Patient Individualization 17
E.2 Comparison to ODE-based Methods 18
E.3 Sparse Observations 19
E.4 Synthetic Data Experiments with Regular Observations 20
E.5 Interpretability: Inspecting the Spectrum 21
E.6 Model Expressiveness and Overfitting 21
F Medication Dosing Prediction: Implementation Details 22
F.1 Data preprocessing 22
F.2 LSTM Baseline Implementation 22
F.3 Extended Results 23
A Observation Filtering: The Conditional Distribution and the Relation to Kalman Filtering
As described in Section 4, the NESDE algorithm keeps an estimated Normal distribution of the system state  $X(t)$  at any point of time. The distribution develops continuously through time according to the dynamics specified by Eq. (6), except for the discrete times where an observation  $\hat{Y}(t)$  is received: in every such point of time, the  $X(t)$  estimate is updated to be the conditional distribution  $X(t)|\hat{Y}(t)$ .
Calculating the conditional Normal distribution: The conditional distribution can be derived as follows. Recall that  $X \sim N(\mu, \Sigma)$  (we remove the time index  $t$  as we focus now on filtering at a single point of time). Denote  $X = (Y, Z)^{\top}$  where  $Y \in \mathbb{R}^m$  (similarly to Eq. (5)) and  $Z \in \mathbb{R}^{n - m}$ ; and similarly,  $\mu = (\mu_Y, \mu_Z)^{\top}$  and
$$
\Sigma = \left( \begin{array}{c c} \Sigma_ {Y Y} & \Sigma_ {Y Z} \\ \Sigma_ {Z Y} & \Sigma_ {Z Z} \end{array} \right)
$$
First consider a noiseless observation  $(R = 0)$ : then according to Eaton [1983], the conditional distribution  $X|Y = \hat{Y}$  is given by  $X = (Y,Z)^{\top}$ ,  $Y = \hat{Y}$  and  $Z \sim N(\mu_Z', \Sigma_{ZZ}')$ , where
$$
\mu_ {Z} ^ {\prime} := \mu_ {Z} + \Sigma_ {Z Y} \Sigma_ {Y Y} ^ {- 1} (\hat {Y} - \mu_ {Y})
$$
$$
\Sigma_ {Z Z} ^ {\prime} := \Sigma_ {Z Z} - \Sigma_ {Z Y} \Sigma_ {Y Y} ^ {- 1} \Sigma_ {Y Z}
$$
In the general case of  $R \neq 0$ , we can redefine the state to include the observation explicitly:  $\tilde{X} = (\hat{Y}, X)^{\top} = (\hat{Y}, Y, Z)^{\top}$ , where  $\tilde{\mu}, \tilde{\Sigma}$  of  $\tilde{X}$  are adjusted by  $\mu_{\hat{Y}} = \mu_y$ ,  $\Sigma_{\hat{Y}\hat{Y}} = \Sigma_{YY} + R$ ,  $\Sigma_{\hat{Y}Y} = R$  and  $\Sigma_{\hat{Y}Z} = \Sigma_{YZ}$ . Then, the conditional distribution can be derived as in the noiseless case above, by simply considering the new observation as a noiseless observation of  $\tilde{X}_{1:m} = \hat{Y}$ .
The relation to the Kalman filtering: The derivation of the conditional distribution is equivalent to the filtering step of the Kalman filter [Kalman, 1960], where the (discrete) model is
$$
X _ {t + 1} = A \cdot X _ {t} + \omega_ {t} \quad (\omega_ {t} \sim N (0, Q))
$$
$$
\hat {Y} _ {t} = H \cdot X _ {t} + \nu_ {t} \quad (\nu_ {t} \sim N (0, R)),
$$
Our setup can be recovered by substituting the following observation model  $H \in \mathbb{R}^{m \times n}$ , which observes the first  $m$  coordinates of  $X$  and ignores the rest:
$$
H = \left( \begin{array}{c c c c c c c} 1 & & & & & 0 & \ldots & 0 \\ & 1 & & & & & \\ & & \ldots & & & | & | & | \\ & & & 1 & & 1 & 0 & \ldots & 0 \end{array} \right)
$$
and the Kalman filtering step is then
$$
K := \Sigma H ^ {\top} (H \Sigma H ^ {\top} + R) ^ {- 1}
$$
$$
\mu^ {\prime} := \mu + K (\hat {Y} - H \mu)
$$
$$
\Sigma^ {\prime} := \Sigma - K H \Sigma
$$
Note that while the standard Kalman filter framework indeed supports the filtering of distributions upon arrival of a new observation, its progress through time is limited to discrete and constant time-steps (see the model above), whereas our SDE-based model can directly make predictions to any arbitrary future time  $t$ .
B Integrator Implementation
Below, we describe the implementation of the integrator of the Eigen-SDE solver mentioned in Section 4.
Numerical integration given  $u(t)$ : In the presence of an arbitrary (continuous) control signal  $u(t)$ , it is impossible to compute the integral that corresponds with  $u(t)$  (Eq. (2)) analytically. On the
other hand,  $u(t)$  is given in advance, and the eigenfunction,  $\Phi (t)$ , is a known function that can be calculated efficiently at any given time. By discretizing the time to any fixed  $\Delta t$ , one could simply replace the integral by a sum term
$$
\int_ {t _ {0}} ^ {t} \Phi (\tau) ^ {- 1} u (\tau) d \tau \approx \sum_ {i = 0} ^ {\frac {t - t _ {0}}{\Delta t}} \Phi (t _ {0} + i \cdot \Delta t) u (t _ {0} + i \cdot \Delta t) \Delta t
$$
while this sum represents  $\frac{t - t_0}{\Delta t}$  calculations, it can be computed efficiently, as it does not require any recursive computation, as both  $\Phi(t)$  and  $u(t)$  are pre-determined, known functions. Each element of the sum is independent of the other elements, and thus the computation could be parallelized.
Analytic integration: The control  $u$  is often constant over any single time-interval  $\mathcal{I}$  (e.g., when the control is piecewise constant). In such cases, for a given interval  $\mathcal{I} = [t_0, t]$  in which  $u(t) = u_{\mathcal{I}}$ , the integral could be solved analytically:
$$
\int_ {t _ {0}} ^ {t} \Phi (\tau) ^ {- 1} u (\tau) d \tau = \int_ {t _ {0}} ^ {t} e ^ {- \Lambda \tau} V ^ {- 1} u _ {\mathcal {I}} d \tau = \int_ {t _ {0}} ^ {t} e ^ {- \Lambda \tau} d \tau V ^ {- 1} u _ {\mathcal {I}} = \frac {1}{\Lambda} \big (e ^ {- \Lambda t _ {0}} - e ^ {- \Lambda t} \big) V ^ {- 1} u _ {\mathcal {I}}
$$
one might notice that for large time intervals this form is numerically unstable, to address this issue, note that this integral is multiplied (Eq. (2)) by  $\Phi(t) = Ve^{\Lambda t}$ , hence we stabilize the solution with the latter exponent:
$$
\Phi (t) \frac {1}{\Lambda} \big (e ^ {- \Lambda t _ {0}} - e ^ {- \Lambda t} \big) V ^ {- 1} u _ {\mathcal {I}} = V \frac {1}{\Lambda} \big (e ^ {\Lambda (t - t _ {0})} - e ^ {\Lambda (t - t)} \big) V ^ {- 1} u _ {\mathcal {I}} = V \frac {1}{\Lambda} \big (e ^ {\Lambda (t - t _ {0})} - 1 \big) V ^ {- 1} u _ {\mathcal {I}}
$$
to achieve a numerically stable computation.
In addition to the integral over  $u(t)$ , we also need to calculate the integral over  $Q$  (Eq. (4)). In this case,  $Q$  is constant, and the following holds;
$$
\int_ {t _ {0}} ^ {t} \Phi (\tau) ^ {- 1} Q \left(\Phi (\tau) ^ {- 1}\right) ^ {\top} d \tau = \int_ {t _ {0}} ^ {t} e ^ {- \Lambda \tau} V ^ {- 1} Q \left(V ^ {- 1}\right) ^ {\top} \left(e ^ {- \Lambda \tau}\right) ^ {\top} d \tau = V ^ {- 1} Q \left(V ^ {- 1}\right) ^ {\top} \circ \int_ {t _ {0}} ^ {t} e ^ {- \tilde {\Lambda} \tau} d \tau
$$
where  $\circ$  denotes the Hadamard product, and
$$
\tilde {\Lambda} = \left( \begin{array}{c c c} 2 \lambda_ {1} & \dots & \lambda_ {1} + \lambda_ {n} \\ \vdots & \ddots & \\ \lambda_ {n} + \lambda_ {1} & \dots & 2 \lambda_ {n} \end{array} \right)
$$
In this form, it is possible to solve the integral analytically, similarly to the integral of the control signal, and again, we use the exponent term from  $\Phi(t)$  to obtain a numerically stable computation.
C The Dynamics Spectrum and Complex Eigenfunction Implementation