SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Stav Belogolovsky*
stav.belo@gmail.com
Danny Eytan
biliary.colic@gmail.com
Ido Greenberg*
gido@campus.technion.ac.il
Shie Mannor*
shie@ee.technion.ac.il
Dosing models often use differential equations to model biological dynamics. Neural differential equations in particular can learn to predict the derivative of a process, which permits predictions at irregular points of time. However, this temporal flexibility often comes with a high sensitivity to noise, whereas medical problems often present high noise and limited data. Moreover, medical dosing models must generalize reliably over individual patients and changing treatment policies. To address these challenges, we introduce the Neural Eigen Stochastic Differential Equation algorithm (NESDE). NESDE provides individualized modeling (using a hypernetwork over patient-level parameters); generalization to new treatment policies (using decoupled control); tunable expressiveness according to the noise level (using piecewise linearity); and fast, continuous, closed-form prediction (using spectral representation). We demonstrate the robustness of NESDE in both synthetic and real medical problems, and use the learned dynamics to publish simulated medical gym environments.
Sequential forecasting in irregular points of time is required in many real-world problems, such as modeling dosing dynamics of various medicines (pharmacodynamics). Consider a patient whose physiological or biochemical state requires continuous monitoring, while blood tests are only available with a limited frequency. Pharmacodynamics models often rely on an ordinary differential equation models (ODE) for forecasting. Additional expressiveness can be obtained via customized learned models, such as neural-ODE, which learns to predict the derivative of the process [Chen et al., 2018, Liu et al., 2019]. By predicting the derivative, neural-ODE can make irregular predictions at flexible time-steps, unlike regular models that operate in constant time-steps (e.g., Kalman filter, Kalman [1960] and recurrent neural networks, Rumelhart et al. [1986]).
However, real-world forecasting remains a challenge for several reasons. First, the variation between patients often requires personalized modeling. Second, neural-ODE methods are often data-hungry: they aggregate numerous derivatives provided by a non-linear neural network, which is often sensitive to noise. Training over a large dataset may stabilize the predictions, but data is often limited. Third, most neural-ODE methods only provide a point-estimate, while uncertainty estimation is often critical in medical settings. Fourth, for every single prediction, the neural-ODE runs a numeric ODE solver,
$^{\dagger}$ Department of Electrical and Computer Engineering, Technion-Israel Institute of Technology, Haifa, Israel  
 $^{\dagger}$ Department of Physiology and Biophysics, Faculty of Medicine, Technion-Israel Institute of Technology, Haifa, Israel
‡Nvidia Research
Table 1: A summary of the features of NESDE. In Sections 5,6, NESDE is tested on synthetic and real medical data. In addition, each component of NESDE is studied experimentally as specified.
Challenge	Solution	Experimental evaluation
Individualized modeling	Hyper-network with high-level features input	Hyper-network ablation (Appendix E.1)
Sample efficiency	Regularized dynamics: piecewise-linear with complex eigenvalues	Varying train size (Section 5, Appendix E.2); varying sparsity (Appendix E.3)
Uncertainty estimation	Probabilistic Kalman filtering	NLL evaluation (Sections 5,6)
Fast continuous inference	Spectral representation with closed-form solution	Time flexibility (Appendix E.4); interpretability (Appendix E.5)
Control generalization	Decoupling control from other inputs	Out of distribution control (Section 5, Appendix E.6)
along with multiple neural network calculations of the derivative. This computational overhead in inference may limit latency-sensitive applications.
A fifth challenge comes from control. In the framework of retrospective forecasting, a control signal (drug dosage) is often considered part of the observation [De Brouwer et al., 2019]. However, this approach raises difficulties if the control is observed at different times or more frequently than other observations. If the control is part of the model output, it may also bias the train loss away from the true objective. Finally, by treating control and observations together, the patterns learned by the model may overfit the control policy used in the data – and generalize poorly to new policies.
Generalization to out-of-distribution control policies is essential when the predictive model supports decision-making, as the control policy may be affected by the model. Such decision-making is an important use-case of sequential prediction: model-based reinforcement learning and control problems require a reliable model [Moerland et al., 2020, Angermueller et al., 2019], in particular in risk-sensitive control [Yu et al., 2021, Greenberg et al., 2022, Greenberg and Mannor, 2021].
Section 4 introduces the Neural Eigen-SDE algorithm (NESDE) for continuous forecasting, which is designed to address the challenges listed above. NESDE uses a hypernetwork to provide an individualized stochastic differential equation (SDE), regularized to be piecewise-linear and represented in spectral form. The SDE derives a probabilistic model similar to Kalman filtering, which provides uncertainty estimation. Finally, the SDE decouples the control signal from other observations, to discourage the model from learning control patterns that may be violated under out-of-distribution control policies. Table 1 summarizes all these features.
Section 5 tests NESDE against both neural-ODE methods and recurrent neural networks. NESDE demonstrates robustness to both noise (by learning from little data) and out-of-distribution control policies. In Appendix E.5, the SDE model of NESDE is shown to enable potential domain knowledge and provide interpretability – via the predicted SDE eigenvalues. Appendix E.4 demonstrates the disadvantage of discrete methods in continuous forecasting.
In Section 6, NESDE demonstrates high prediction accuracy in two medical forecasting problems with noisy and irregular real-world data: (1) blood coagulation prediction given Heparin dosage, and (2) prediction of the Vancomycin (antibiotics) levels for patients who received it.
Contribution:
- We characterize the main challenges in continuous forecasting for medication dosing control.
- We design the novel Neural Eigen-SDE algorithm (NESDE), which addresses the challenges as summarized in Table 1 and demonstrated empirically in Sections 5 and 6.
- We use NESDE to improve modeling accuracy in two medication dosing processes. Based on the learned models, we simulate gym environments for future research of healthcare control.
Classic filtering: Classic models for sequential prediction in time-series include ARIMA models [Moran and Whittle, 1951] and the Kalman filter (KF) [Kalman, 1960]. The KF provides probabilistic distributions and in particular uncertainty estimation. While the classic KF is limited to linear dynamics, many non-linear extensions have been suggested [Krishnan et al., 2015, Coskun et al., 2017, Revach et al., 2021, Greenberg et al., 2021]. However, such models are typically limited to a constant prediction horizon (time-step). Longer-horizon predictions are often made by applying the model recursively [Herrera et al., 2007, Bontempi et al., 2013], This poses a significant challenge to many optimization methods [Kolen and Kremer, 2001], as also demonstrated in Appendix E.4.
Limited types of irregularity can also be handled by KF with intermittent observations [Park and Sahai, 2011, Sinopoli et al., 2004] or periodical time-steps [Li et al., 2008].
Recurrent neural networks: Sequential prediction is often addressed via neural network models, relying on architectures such as RNN [Rumelhart et al., 1986], LSTM [Hochreiter and Schmidhuber, 1997] and transformers [Vaswani et al., 2017]. LSTM, for example, is a key component in many SOTA algorithms for non-linear sequential prediction [Neu et al., 2021]. LSTM can be extended to a filtering framework to alternately making predictions and processing observations, and even to provide uncertainty estimation [Gao et al., 2019]. However, these models are typically limited to constant time-steps, and thus suffer from the limitations discussed above.
Neural-ODE models: Parameterized ODE models can be optimized by propagating the gradients of a loss function through an ODE solver [Chen et al., 2018, Liu et al., 2019, Rubanova et al., 2019]. By predicting the process derivative and using an ODE solver in real-time, these methods can choose the effective time-steps flexibly. Uncertainty estimation can be added via process variance prediction [De Brouwer et al., 2019]. However, since neural-ODE methods learn a non-linear dynamics model, the ODE solver operates numerically and recursively on top of multiple neural network calculations. This affects running time, training difficulty and data efficiency as discussed above. While neural-ODE models have been studied for medical applications with irregular data [Lu et al., 2021], simpler models are commonly preferred in practice. For example, the effects of Heparin on blood coagulation is usually modeled by either using discrete models [Nemati et al., 2016] or manually based on domain knowledge [Delavenne et al., 2017].
Our method uses SDE with piecewise linear dynamics (note this is different from a piecewise linear process). The linear dynamics per time interval permit efficient and continuous closed-form forecasting of both mean and covariance. Schirmer et al. [2022] also rely on a linear ODE model, but only support operators with real-valued eigenvalues (which limits the modeling of periodic processes), and do not separate control signal from observations (which limits generalization to out-of-distribution control). Our piecewise linear architecture, tested below against alternative methods including De Brouwer et al. [2019] and Schirmer et al. [2022], is demonstrated to be more robust to noisy, sparse or small datasets, even under out-of-distribution control policies.
2 Preliminaries: Linear SDE
We consider a particular case of the general linear Stochastic Differential Equation (SDE):
$$
d X (t) = [ A \cdot X (t) + \tilde {u} (t) ] + d W (t) \tag {1}
$$
where  $X: \mathbb{R} \to \mathbb{R}^n$  is a time-dependent state;  $A \in \mathbb{R}^{n \times n}$  is a fixed dynamics operator;  $\tilde{u}: \mathbb{R} \to \mathbb{R}^n$  is the control signal; and  $dW: \mathbb{R} \to \mathbb{R}^n$  is a Brownian motion vector with covariance  $Q \in \mathbb{R}^{n \times n}$ .
General SDEs can be solved numerically using the first-order approximation  $\Delta X(t) \approx \Delta t \cdot dX(t)$ , or using more delicate approximations [Wang and Lin, 1998]. The linear SDE, however, and in particular Eq. (1), can be solved analytically [Herzog, 2013]:
$$
X (t) = \Phi (t) \left(\Phi (t _ {0}) ^ {- 1} X (t _ {0}) + \int_ {t _ {0}} ^ {t} \Phi (\tau) ^ {- 1} \tilde {u} (\tau) d \tau + \int_ {t _ {0}} ^ {t} \Phi (\tau) ^ {- 1} d W (\tau)\right) \tag {2}
$$
where  $X(t_0)$  is an initial condition, and  $\Phi (t)$  is the eigenfunction of the system. More specifically, if  $V$  is the matrix whose columns  $\{v_{i}\}_{i = 1}^{n}$  are the eigenvectors of  $A$ , and  $\Lambda$  is the diagonal matrix
whose diagonal contains the corresponding eigenvalues  $\lambda = \{\lambda_{i}\}_{i = 1}^{n}$ , then
$$
\Phi (t) = V e ^ {\Lambda t} = \left( \begin{array}{c c c c c} | & | & | & | & | \\ v _ {1} \cdot e ^ {\lambda_ {1} t} & \dots & v _ {i} \cdot e ^ {\lambda_ {i} t} & \dots & v _ {n} \cdot e ^ {\lambda_ {n} t} \\ | & | & | & | & | \end{array} \right) \tag {3}
$$
If the initial condition is given as  $X(t_0)\sim N(\mu_0,\Sigma_0)$  , Eq. (2) becomes
$$
X (t) \sim N (\mu (t), \Sigma (t))
$$
$$
\mu (t) = \Phi (t) \left(\Phi (t _ {0}) ^ {- 1} \mu_ {0} + \int_ {t _ {0}} ^ {t} \Phi (\tau) ^ {- 1} \tilde {u} (\tau) d \tau\right), \Sigma (t) = \Phi (t) \Sigma^ {\prime} (t) \Phi (t) ^ {\top} \tag {4}
$$
where  $\Sigma^{\prime}(t) = \Phi (t_0)^{-1}\Sigma_0\bigl (\Phi (t_0)^{-1}\bigr)^\top +\int_{t_0}^{t}\Phi (\tau)^{-1}Q\bigl (\Phi (\tau)^{-1}\bigr)^\top d\tau$
Note that if  $\forall i:\lambda_i < 0$  and  $\tilde{u}\equiv 0$ , we have  $\mu (t)\xrightarrow{t\to\infty}0$  (stable system). In addition, if  $\lambda$  is complex, Eq. (4) may produce a complex solution; Appendix C explains how to use a careful parameterization to only calculate the real solutions.
3 Problem Setup: Sparsely-Observable SDE
We focus on online sequential prediction of a process  $Y(t) \in \mathbb{R}^m$ . To predict  $Y(t_0)$  at a certain  $t_0$ , we can use noisy observations  $\hat{Y}(t)$  (at given times  $t < t_0$ ), as well as a control signal  $u(t) \in \mathbb{R}^k$  ( $\forall t < t_0$ ); offline data of  $Y$  and  $u$  from other sequences; and one sample of contextual information  $C \in \mathbb{R}^{d_c}$  per sequence (capturing properties of the whole sequence). The dynamics of  $Y$  are unknown and may vary between sequences. For example, sequences may represent different patients, each with its own dynamics;  $C$  may represent patient information; and the objective is "zero-shot" learning
(b)
Figure 1: Samples of sparsely observed SDEs: the Brownian noise and the sparse observations pose a major challenge for learning the underlying SDE dynamics. Efficient learning from external trajectories data is required, as the current trajectory often does not contain sufficient observations.
upon arrival of a sequence of any new patient. In addition, the observations within a sequence are both irregular and sparse: they are received at arbitrary points of time, and are sparse in comparison to the required prediction frequency (i.e., continuous forecasting, as illustrated in Fig. 1).
To model the problem, we assume the observations  $Y(t)$  to originate from an unobservable latent process  $X(t) \in \mathbb{R}^n$  (where  $n > m$  is a hyperparameter). More specifically:
$$
d X (t) = F _ {C} \left(X (t), u (t)\right), Y (t) = X (t) _ {1: m}, \hat {Y} (t) = Y (t) + \nu_ {C} (t) \tag {5}
$$
where  $F_{C}$  is a stochastic dynamics operator (which may depend on the context  $C$ );  $Y$  is simply the first  $m$  coordinates of  $X$ ;  $\hat{Y}$  is the corresponding observation; and  $\nu_{C}(t)$  is its i.i.d Gaussian noise with zero-mean and (unknown) covariance  $R_{C} \in \mathbb{R}^{m \times m}$  (which may also depend on  $C$ ). Our goal is to predict  $Y$ , where the dynamics  $F_{C}$  are unknown and the latent subspace of  $X$  is unobservable. In cases where data of  $Y$  is not available, we measure our prediction accuracy against  $\hat{Y}$ . The control  $u(t)$  is modeled separately from  $\hat{Y}$ , is not part of the prediction objective, and does not depend on  $X$ .
4 Neural Eigen-SDE Algorithm
Model: In this section, we introduce the Neural Eigen-SDE algorithm (NESDE, shown in Algorithm 1 and Fig. 2). NESDE predicts the signal  $Y(t)$  of Eq. (5) continuously at any required point of time  $t$ . It relies on a piecewise linear approximation which reduces Eq. (5) into Eq. (1):
$$
\forall t \in \mathcal {I} _ {i}: d X (t) = \left[ A _ {i} \cdot \left(X (t) - \alpha\right) + B \cdot u (t) \right] + d W (t) \tag {6}
$$
where  $\mathcal{I}_i = (t_i, t_{i+1})$  is a time interval,  $dW$  is a Brownian noise with covariance matrix  $Q_i$ , and  $A_i \in \mathbb{R}^{n \times n}$ ,  $B \in \mathbb{R}^{n \times k}$ ,  $Q_i \in \mathbb{R}^{n \times n}$ ,  $\alpha \in \mathbb{R}^n$  form the linear dynamics model corresponding to the
Figure 2: NESDE algorithm: Hypernet determines the SDE parameters; Eigen-SDE solver uses them to predict the next time-interval; the filter updates the state upon arrival of a new observation, which initiates a new interval. For more frequent updates of the dynamics, the initial condition becomes the last prediction (dotted arrow).
Figure 3: Zoom into the individualization module of NESDE: Hypernet uses the context and the estimated state to determine the SDE parameters for the solver.
interval  $\mathcal{I}_i$ . In terms of Eq. (1), we substitute  $A \coloneqq A_i$  and  $\tilde{u} \coloneqq Bu - A_i\alpha$ . Note that if  $A_i$  is a stable system and  $u \equiv 0$ , the asymptotic state is  $\mu(t) \xrightarrow{t \to \infty} \alpha$ . To solve Eq. (6) within every  $\mathcal{I}_i$ , NESDE has to learn the parameters  $\{A_i, Q_i\}_i, \alpha, B$ .
The end of  $\mathcal{I}_{i-1}$  typically represents one of two events: either an update of the dynamics  $A$  (allowing the piecewise linear dynamics), or the arrival of a new observation. A new observation at time  $t_i$  triggers an update of  $X(t_i)$  according to the conditional distribution  $X(t_i)|\hat{Y}(t_i)$  (this is a particular case of Kalman filtering, as shown in Appendix A). Then, the prediction continues for  $\mathcal{I}_i$  according to Eq. (6). Note that once  $X(t_0)$  is initialized to have a Normal distribution, it remains Normally-distributed throughout both the process dynamics (Eq. (4)) and the observations filtering (Appendix A). This allows NESDE to efficiently capture the distribution of  $X(t)$ , where the estimated covariance represents the uncertainty.
Eigen-SDE solver (ESDE) - spectral dynamics representation: A key feature of NESDE is that  $A_{i}$  is only represented implicitly through the parameters  $V, \lambda$  defining its eigenfunction  $\Phi(t)$  of Eq. (3) (we drop the interval index  $i$  with a slight abuse of notation). The spectral representation allows Eq. (4) to solve  $X(t)$  analytically for any  $t \in \mathcal{I}_i$  at once: the predictions are not limited to predefined times, and do not require recursive iterations with constant time-steps. This is particularly useful in the sparsely-observable setup of Section 3, as it lets numerous predictions be made at once without being "interrupted" by a new measurement.
The calculation of Eq. (4) requires efficient integration. Many SDE solvers apply recursive numeric integration [Chen et al., 2018, De Brouwer et al., 2019]. In NESDE, however, thanks to the spectral decomposition, the integration only depends on known functions of  $t$  instead of  $X(t)$  (Eq. (4)), hence recursion is not needed, and the computation can be paralleled. Furthermore, if the control  $u$  is constant over an interval  $\mathcal{I}_i$  (or has any other analytically-integrable form), Appendix B shows how to calculate the integration analytically. Piecewise constant  $u$  is common, for example, when the control is updated along with the observations.
In addition to simplifying the calculation, the spectrum of  $A_{i}$  carries significant meaning about the dynamics. For example, negative eigenvalues correspond to a stable solution, whereas imaginary ones indicate periodicity, as demonstrated in Fig. 4. Note that the (possibly-complex) eigenvalues must be constrained to represent a real matrix  $A_{i}$ . The constraints and the calculations in the complex space are detailed in Appendix C.
Note that if the process  $X(t)$  actually follows the piecewise linear model of Eq. (6), and the model parameters are known correctly, then the Eigen-SDE solver trivially returns the optimal predictions. The complete proposition and proof are provided in Appendix D.
Proposition 1 (Eigen-SDE solver optimality). If  $X(t)$  follows Eq. (6) with the same parameters used by the Eigen-SDE solver, then under certain conditions, the solver prediction at any point of time optimizes both the expected error and the expected log-likelihood.
(a) Complex  $\lambda$
(b) Real  $\lambda$
(c) Imaginary  $\lambda$
Figure 4: Sample trajectories with different types of dynamics (and random control signal that is not displayed). As demonstrated in Appendix E.5, NESDE directly estimates the dynamics  $\lambda$  and reveals their type.
Proof Sketch. If the process  $X(t)$  follows Eq. (6), the Eigen-SDE solver output corresponds to the true distribution  $X(t) \sim N(\mu(t), \Sigma(t))$  for any  $t$ . Thus, both the expected error and the expected log-likelihood are optimal.
Updating solver and filter parameters: NESDE provides the parameters  $V, \lambda, Q, B, \alpha$  to the Eigen-SDE solver, as well as the noise  $R$  to the observation filter. As NESDE assumes a piecewise linear model, it separates the time into intervals  $\mathcal{I}_i = (t_i, t_{i+1})$  (the interval length is a hyperparameter), and uses a dedicated model to predict new parameters at the beginning  $t_i$  of every interval.
The model receives the current state  $X(t_{i})$  and the contextual information  $C$ , and returns the parameters for  $\mathcal{I}_i$ . Specifically, we use Hypernet [Ha et al., 2016]: one neural network  $g_{1}$  with network parameters  $\Theta$ , returns the weights
Algorithm 1 NESDE
Input: context  $C$  ; control signal  $u(t)$  ; update times  $\mathcal{I}\in$ $\mathbb{R}^T$  ; prediction times  $\{P_{\mathcal{I}_i}\}_{i\in \mathcal{I}}$    
Initialize:  $\mu ,\Sigma ,\alpha ,R\gets \mathrm{Prior}(C)$    
for  $\mathcal{I}_i$  in  $\mathcal{I}$  : do  
 $V,\lambda ,Q,B,\alpha ,R\gets \mathrm{Hypernet}(C,\mu ,\Sigma)$    
for  $t$  in  $P_{\mathcal{I}_i}$  do  
 $\mu_t,\Sigma_t\gets \operatorname {ESDE}(\mu ,\Sigma ,u,t;V,\lambda ,Q,B)\quad //Eq.(4)$    
predict:  $\tilde{Y}_t\sim \mathcal{N}(\mu_t + \alpha ,\Sigma_t + R)$    
if given observation  $\hat{Y}_t$  then  
 $\mu ,\Sigma \gets \mathrm{Filter}(\mu_t,\Sigma_t,R,\hat{Y}_t)\quad //A p p e n d i x A$    
end if  
end for  
end for
of a second network  $g_{2}$ , i.e.,  $W = g_{1}(C; \Theta)$ . Then,  $g_{2}$  maps the current state estimation to a set of dynamics parameters:  $(V, \lambda, Q, B, \alpha, R) = g_{2}(X; W) = g_{2}(X; g_{1}(C; \Theta))$  (Fig. 3). The dynamics  $g_{2}(X; W)$  may be updated more often than the individualized weights  $W = g_{1}(C; \Theta)$ , as the context  $C$  may change with a lower frequency. For the initial state, where  $X$  is unavailable, we learn a state prior from  $C$  by a dedicated network; this prior helps NESDE to function as a "zero-shot" model.
The Hypernet module implementation gives us control over the non-linearity and non-stationarity of the model. The importance of the Hypernet module for individualized modeling is demonstrated in an ablation test in Appendix E.1. In our current implementation, only  $V, \lambda, Q$  are renewed every time interval.  $\alpha$  (asymptotic signal) and  $R$  (observation noise) are only predicted once per sequence, as we assume they are independent of the state. The control mapping  $B$  is assumed to be a global parameter.
Training: The learnable parameters of NESDE are the control mapping  $B$  and Hypernet's parameters  $\Theta$  (which in turn determine the rest of the parameters). To optimize them, the training relies on a dataset of sequences of control signals  $\{u_{seq}(t_j)\}_{seq,j}$  and (sparse and possibly irregular) states and observations  $\{(Y_{seq}(t_j),\hat{Y}_{seq}(t_j))\}_{seq,j}$  (if  $Y$  is not available, we use  $\hat{Y}$  instead as the training target).
The latent space dimension  $n$  and the model-update frequency  $\Delta t$  are determined as hyperparameters. Then, we use the standard Adam optimizer [Diederik P. Kingma, 2015] to optimize the parameters with respect to the loss  $NLL(j) = -\log P(Y(t_j)|\mu (t_j),\Sigma (t_j))$  (where  $\mu ,\Sigma$  are predicted by NESDE sequentially from  $u,\hat{Y}$ ). Each training iteration corresponds to a batch of sequences of data, where the  $NLL$  is aggregated over all the samples of the sequences. Note that our supervision for the training is limited to the times of the observations, even if we wish to make more frequent predictions in inference.