Title: Individualized Dosing Dynamics via Neural Eigen Decomposition
Venue: NeurIPS 2023
Accepted: yes
Keywords: personalized medicine

dosing dynamics

sequential prediction

stochastic differential equations

Kalman filter

recurrent neural networks

medical drug control
Abstract: 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.

Individualized Dosing Dynamics via Neural Eigen Decomposition

Abstract

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.

1 Introduction

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,

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.

1.1 Related Work

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 outof-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):

where X : R → R n is a time-dependent state; A ∈ R n × n is a fixed dynamics operator; ˜ u : R → R n is the control signal; and dW : R → R n is a Brownian motion vector with covariance Q ∈ R n × n .

General SDEs can be solved numerically using the first-order approximation ∆ X ( t ) ≈ ∆ t · 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]:

where X ( t 0 ) is an initial condition, and Φ( t ) is the eigenfunction of the system. More specifically, if V is the matrix whose columns { v i } n i =1 are the eigenvectors of A , and Λ is the diagonal matrix

whose diagonal contains the corresponding eigenvalues λ = { λ i } n i =1 , then

If the initial condition is given as X ( t 0 ) ∼ N ( µ 0 , Σ 0 ) , Eq. (2) becomes

Note that if ∀ i : λ i < 0 and ˜ u ≡ 0 , we have µ ( t ) t →∞ - - - → 0 (stable system). In addition, if λ 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 ) ∈ R m . To predict Y ( t 0 ) at a certain t 0 , we can use noisy observations ˆ Y ( t ) (at given times t < t 0 ), as well as a control signal u ( t ) ∈ R k ( ∀ t < t 0 ) ; offline data of Y and u from other sequences; and one sample of contextual information C ∈ 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

(a)

(b)

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 ) ∈ R n (where n > m is a hyperparameter). More specifically:

where F C is a stochastic dynamics operator (which may depend on the context C ); Y is simply the first m coordinates of X ; ˆ Y is the corresponding observation; and ν C ( t ) is its i.i.d Gaussian noise with zero-mean and (unknown) covariance R C ∈ R m × 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 ˆ Y . The control u ( t ) is modeled separately from ˆ 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):

where I i = ( t i , t i +1 ) is a time interval, dW is a Brownian noise with covariance matrix Q i , and A i ∈ R n × n , B ∈ R n × k , Q i ∈ R n × n , α ∈ R n form the linear dynamics model corresponding to the

Output

Inputs

NESDE

N(μ(t)+a,∑(t）)

Observations

u(t):

Control

Observation

Initialcondition

filter

μti-1,Eti-1

Eigen-SDE

μ(t),∑(t)

solver

V,A

Context

Hypernet

Update

Q,B

(Individualization)

dynamics

R

t∈[ti-1,ti]

Time

:Context

HypernetInputs

C2

91(C1;0)

Hypernet

V,入,Q,B,a,R

Output

:predictionsi

:Filtered

Parametersforg2

Wc

μ（T),∑(T);

g2(μ(t),∑(t);Wc)

update times

Dynamics

t1

t2

interval I i . In terms of Eq. (1), we substitute A := A i and ˜ u := Bu -A i α . Note that if A i is a stable system and u ≡ 0 , the asymptotic state is µ ( t ) t →∞ - - - → α . To solve Eq. (6) within every I i , NESDE has to learn the parameters { A i , Q i } i , α, B .

The end of 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 ) | ˆ Y ( t i ) (this is a particular case of Kalman filtering, as shown in Appendix A). Then, the prediction continues for I i according to Eq. (6). Note that once X ( t 0 ) is initialized to have a Normal distribution, it remains Normallydistributed 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: Akey feature of NESDE is that A i is only represented implicitly through the parameters V, λ defining its eigenfunction Φ( 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 ∈ 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 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

λ

(b) Real

λ

(c) Imaginary

λ

prediction

4

NESDE

prediction

2

NESDE

prediction

5

2

0

0

0

2

5

-2

10

NESDE

-4

15

0

2

4

time

6

8

10

0

2

4

time

6

8

10

0

2

4

time

6

8

10

Proof Sketch. If the process X ( t ) follows Eq. (6), the Eigen-SDE solver output corresponds to the true distribution X ( t ) ∼ N ( µ ( t ) , Σ( 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, λ, Q, B, α 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 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 I i . Specifically, we use Hypernet [Ha et al., 2016]: one neural network g 1 with network parameters Θ , returns the weights

Algorithm 1 NESDE

of a second network g 2 , i.e., W = g 1 ( C ; Θ) . Then, g 2 maps the current state estimation to a set of dynamics parameters: ( V, λ, Q, B, α, R ) = g 2 ( X ; W ) = g 2 ( X ; g 1 ( C ; Θ)) (Fig. 3). The dynamics g 2 ( X ; W ) may be updated more often than the individualized weights W = g 1 ( C ; Θ) , 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, λ, Q are renewed every time interval. α (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 Θ (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 (sparser and possibly irregular) states and observations { ( Y seq ( t j ) , ˆ Y seq ( t j )) } seq,j (if Y is not available, we use ˆ Y instead as the training target).

The latent space dimension n and the model-update frequency ∆ 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 ) | µ ( t j ) , Σ( t j )) (where µ, Σ are predicted by NESDE sequentially from u, ˆ 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.

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 · Y t . b t ∼ U [0 , 0 . 5] is a piecewise constant additive noise (changing 10 times per trajectory).

(a) Same control distribution (b) Out of distribution control

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 · 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

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 λ 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

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 .

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 .

6.3 Results

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

(a) Heparin

(b) Vancomycin

70

predictions

[w/6n] po0lg w

20

aPTT [sec]

60

observations

15

50

10

40

5

predictions

UH rate [U/min]

30

0

observations

5

VMDose

1.0

10

0.5

0

0.0

0

500

1000

1500

2000

2500

0

1000

2000

3000

Time [min]

Time [min]

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 100 s and the bottom 5% are below 25 s ), 50% of all the predictions have errors lower than 12 . 4 s - 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 10 s .

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.

7 Conclusion

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).

25

20

Error

15

NESDE

GRU-ODE

10

LSTM

CRU

5

Naive

0

200

400

600

800

1000

1200

1400

Time sincelastobservation[m]

References

Christof Angermueller, David Dohan, David Belanger, Ramya Deshpande, Kevin Murphy, and Lucy Colwell. Model-based reinforcement learning for biological sequence design. In International conference on learning representations , 2019.

Gianluca Bontempi, Souhaib Ben Taieb, and Yann-Aël Le Borgne. Machine learning strategies for time series forecasting. Lecture Notes in Business Information Processing , 138, 01 2013. doi: 10.1007/978-3-642-36318-4_3.

Ricky TQ Chen, Yulia Rubanova, Jesse Bettencourt, and David Duvenaud. Neural ordinary differential equations. arXiv preprint arXiv:1806.07366 , 2018.

Huseyin Coskun, Felix Achilles, Robert DiPietro, Nassir Navab, and Federico Tombari. Long short-term memory kalman filters:recurrent neural estimators for pose regularization. ICCV , 2017. URL https://github.com/Seleucia/lstmkf_ICCV2017 .

Edward De Brouwer, Jaak Simm, Adam Arany, and Yves Moreau. Gru-ode-bayes: Continuous modeling of sporadically-observed 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 , volume 32. Curran Associates, Inc., 2019. URL https://proceedings.neurips.cc/ paper/2019/file/455cb2657aaa59e32fad80cb0b65b9dc-Paper.pdf .

Delavenne, E. Ollier, S. Chollet, F. Sandri, J. Lanoiselée, S. Hodin, A. Montmartin, J.-F. Fuzellier, P. Mismetti, and L. Gergelé. Pharmacokinetic/pharmacodynamic model for unfractionated heparin dosing during cardiopulmonary bypass. BJA: British Journal of Anaesthesia , 118(5):705-712, May 2017. ISSN 0007-0912. doi: 10.1093/bja/aex044. URL https://doi.org/10.1093/bja/ aex044 .

Jimmy Ba Diederik P. Kingma. Adam: A method for stochastic optimization. ICLR , 2015. URL https://arxiv.org/abs/1412.6980 .

Morris L. Eaton. Multivariate Statistics: a Vector Space Approach . John Wiley and Sons, 1983.

Edward J Filippone, Walter K Kraft, and John L Farber. The nephrotoxicity of vancomycin. Clinical Pharmacology & Therapeutics , 102(3):459-469, 2017.

Chang Gao, Junkun Yan, Shenghua Zhou, Bo Chen, and Hongwei Liu. Long short-term memorybased recurrent neural networks for nonlinear target tracking. Signal Processing , 164, 05 2019. doi: 10.1016/j.sigpro.2019.05.027.

Ido Greenberg and Shie Mannor. Detecting rewards deterioration in episodic reinforcement learning. In Proceedings of the 38th International Conference on Machine Learning , volume 139, pages 3842-3853. PMLR, 18-24 Jul 2021. URL https://proceedings.mlr.press/v139/ greenberg21a.html .

Ido Greenberg, Netanel Yannay, and Shie Mannor. Noise estimation is not optimal: How to use kalman filter the right way. arXiv preprint arXiv:2104.02372 , 2021.

Ido Greenberg, Yinlam Chow, Mohammad Ghavamzadeh, and Shie Mannor. Efficient risk-averse reinforcement learning. Advances in Neural Information Processing Systems , 2022.

David Ha, Andrew Dai, and Quoc V Le. Hypernetworks. arXiv preprint arXiv:1609.09106 , 2016.

Luis Herrera, Hector Pomares, I. Rojas, Alberto Guillén, Alberto Prieto, and Olga Valenzuela. Recursive prediction for long term time series forecasting using advanced models. Neurocomputing , 70:2870-2880, 10 2007. doi: 10.1016/j.neucom.2006.04.015.

Florian Herzog. Stochastic differential equations, 2013. URL https://ethz.ch/content/dam/ ethz/special-interest/mavt/dynamic-systems-n-control/idsc-dam/Lectures/ Stochastic-Systems/SDE.pdf .

Sepp Hochreiter and Jurgen Schmidhuber. Long short-term memory. Neural Computation , 1997. URL https://direct.mit.edu/neco/article/9/8/1735/6109/Long-Short-Term-Memory .

A Johnson, L Bulgarelli, T Pollard, S Horng, LA Celi, and R Mark. Mimic-iv, 2020.

E. Kalman. A New Approach to Linear Filtering and Prediction Problems. Journal of Basic Engineering , 82(1):35-45, 03 1960. ISSN 0021-9223. doi: 10.1115/1.3662552. URL https: //doi.org/10.1115/1.3662552 .

John F. Kolen and Stefan C. Kremer. Gradient Flow in Recurrent Nets: The Difficulty of Learning LongTerm Dependencies , pages 237-243. Wiley-IEEE Press, 2001. doi: 10.1109/9780470544037. ch14.

Rahul G. Krishnan, Uri Shalit, and David Sontag. Deep kalman filters, 2015.

C Seth Landefeld, E Francis Cook, Margaret Flatley, Monica Weisberg, and Lee Goldman. Identification and preliminary validation of predictors of major bleeding in hospitalized patients starting anticoagulant therapy. The American journal of medicine , 82(4):703-713, 1987.

Weihua Li, Sirish L. Shah, and Deyun Xiao. Kalman filters in non-uniformly sampled multirate systems: For fdi and beyond. Automatica , 44(1):199-208, jan 2008. ISSN 0005-1098. doi: 10.1016/j.automatica.2007.05.009. URL https://doi.org/10.1016/j.automatica.2007. 05.009 .

Xuanqing Liu, Si Si, Qin Cao, Sanjiv Kumar, and Cho-Jui Hsieh. Neural sde: Stabilizing neural ode networks with stochastic noise. arXiv preprint arXiv:1906.02355 , 2019.

James Lu, Kaiwen Deng, Xinyuan Zhang, Gengbo Liu, and Yuanfang Guan. Neural-ode for pharmacokinetics modeling and its advantage to alternative machine learning models in predicting new dosing regimens. Iscience , 24(7):102804, 2021.

Amélie Marsot, Audrey Boulamery, Bernard Bruguerolle, and Nicolas Simon. Vancomycin. Clinical pharmacokinetics , 51(1):1-13, 2012.

Thomas M. Moerland, Joost Broekens, and Catholijn M. Jonker. Model-based reinforcement learning: A survey. CoRR , abs/2006.16712, 2020. URL https://arxiv.org/abs/2006.16712 .

P. A. P. Moran and Peter Whittle. Hypothesis testing in time series analysis, 1951.

Shamim Nemati, Mohammad M Ghassemi, and Gari D Clifford. Optimal medication dosing from suboptimal clinical examples: A deep reinforcement learning approach. In 2016 38th Annual International Conference of the IEEE Engineering in Medicine and Biology Society (EMBC) , pages 2978-2981. IEEE, 2016.

Dominic A. Neu, Johannes Lahann, and Peter Fettke. A systematic literature review on stateof-the-art deep learning methods for process prediction. CoRR , abs/2101.09320, 2021. URL https://arxiv.org/abs/2101.09320 .

Se Yong Park and Anant Sahai. Intermittent kalman filtering: Eigenvalue cycles and nonuniform sampling. In Proceedings of the 2011 American Control Conference , pages 3692-3697, 2011. doi: 10.1109/ACC.2011.5991285.

Guy Revach, Nir Shlezinger, Xiaoyong Ni, Adria Lopez Escoriza, Ruud J. G. van Sloun, and Yonina C. Eldar. Kalmannet: Neural network aided kalman filtering for partially known dynamics, 2021.

Yulia Rubanova, Ricky T. Q. Chen, and David K Duvenaud. Latent ordinary differential equations for irregularly-sampled 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 , volume 32. Curran Associates, Inc., 2019. URL https://proceedings.neurips.cc/paper/2019/file/ 42a6845a557bef704ad8ac9cb4461d43-Paper.pdf .

David E. Rumelhart et al. Learning representations by back-propagating errors. Nature , 1986. URL https://www.nature.com/articles/323533a0 .

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.

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.

Appendices

Table of Contents

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 ˆ Y ( t ) is received: in every such point of time, the X ( t ) estimate is updated to be the conditional distribution X ( t ) | ˆ Y ( t ) .

Calculating the conditional Normal distribution: The conditional distribution can be derived as follows. Recall that X ∼ N ( µ, Σ) (we remove the time index t as we focus now on filtering at a single point of time). Denote X = ( Y, Z ) ⊤ where Y ∈ R m (similarly to Eq. (5)) and Z ∈ R n -m ; and similarly, µ = ( µ Y , µ Z ) ⊤ and

First consider a noiseless observation ( R = 0 ): then according to Eaton [1983], the conditional distribution X | Y = ˆ Y is given by X = ( Y, Z ) ⊤ , Y = ˆ Y and Z ∼ N ( µ ′ Z , Σ ′ ZZ ) , where

In the general case of R = 0 , we can redefine the state to include the observation explicitly: ˜ X = ( ˆ Y , X ) ⊤ = ( ˆ Y , Y, Z ) ⊤ , where ˜ µ, ˜ Σ of ˜ X are adjusted by µ ˆ Y = µ y , Σ ˆ Y ˆ Y = Σ Y Y + R , Σ ˆ Y Y = R and Σ ˆ Y Z = Σ Y Z . Then, the conditional distribution can be derived as in the noiseless case above, by simply considering the new observation as a noiseless observation of ˜ X 1: m = ˆ 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

Our setup can be recovered by substituting the following observation model H ∈ R m × n , which observes the first m coordinates of X and ignores the rest:

and the Kalman filtering step is then

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, Φ( t ) , is a known function that can be calculated efficiently at any given time. By discretizing the time to any fixed ∆ t , one could simply replace the integral by a sum term

while this sum represent t -t 0 ∆ t calculations, it can be computed efficiently, as it does not require any recursive computation, as both Φ( 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 I (e.g., when the control is piecewise constant). In such cases, for a given interval I = [ t 0 , t ] in which u ( t ) = u I , the integral could be solved analytically:

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 Φ( t ) = V e Λ t , hence we stabilize the solution with the latter exponent:

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;

where ◦ denotes the Hadamard product, and

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 Φ( t ) to obtain a numerically stable computation.

C The Dynamics Spectrum and Complex Eigenfunction Implementation

The form of the eigenfunction matrix as presented in Section 2 is valid for real eigenvalues. Complex eigenvalues induce a slightly different form; firstly, they come in pairs, i.e., if z = a + bi is an eigenvalue of A (Eq. (1)), then ¯ z = a -bi (the complex conjugate of z ) is an eigenvalue of A . The corresponding eigenvector of z is complex as well, denote it by v = v real + v im i , then ¯ v (the complex conjugate of v ) is the eigenvector that correspond to ¯ z . Secondly, the eigenfunction matrix takes the form:

For brevity, we consider only the elements that correspond with z , ¯ z . To parametrize this form, we use the same number of parameters (each complex number need two parameters to represent, but since they come in pairs with their conjugates we get the same overall number) which are organized differently. Mixed eigenvalues (e.g., both real and complex) induce a mixed eigenfunction that is a concatenation of the two forms. Since the complex case requires a different computation, we leave the number of complex eigenvalues to be a hyperparameter. Same as for the real eigenvalues setting, it is possible to derive an analytical computation for the integrals. Here, it takes a different form, as

the complex eigenvalues introduce trigonometric functions to the eigenfunction matrix. To describe the analytical computation, first notice that:

and thus:

Note that here we consider a two-dimensional SDE, for the general case the trigonometric matrix is a block-diagonal matrix, and the exponent becomes a diagonal matrix in which each element repeats twice. It is clear that similarly to the real eigenvalues case, the integral term that includes u (as shown above) can be decomposed, and it is possible to derive an analytical solution for an exponent multiplied by sine or cosine. One major difference is that here we use matrix product instead of Hadamard product. The integral over Q becomes more tricky, but it can be separated and computed as well, with the assistance of basic linear algebra (both are implemented in our code).

D Solver Analysis

Below is a more complete version of Proposition 1 and its proof.

Proposition 2 (Eigen-SDE solver optimality: complete formulation) . Let X ( t ) be a signal that follows Eq. (6) for any time interval I i = [ t i , t i +1 ] , and u ( t ) a control signal that is constant over I i for any i . For any i , consider the Eigen-SDE solver with the parameters corresponding to Eq. (6) (for the same I i ). Assume that the first solver ( i = 0 ) is initialized with the true initial distribution X (0) ∼ N ( µ 0 , Σ 0 ) , and for i ≥ 1 , the i 'th solver is initialized with the i -1 'th output, along with an observation filter if an observation was received. For any interval i and any time t ∈ I i , consider the prediction ˜ X ( t ) ∼ N ( µ ( t ) , Σ( t )) of the solver. Then, µ ( t ) minimizes the expected square error of the signal X ( t ) , and ˜ X ( t ) maximizes the expected log-likelihood of X ( t ) .

Proof. We prove by induction over i that for any i and any t ∈ I i , ˜ X ( t ) corresponds to the true distribution of the signal X ( t ) .

For i = 0 , X ( t i ) = X (0) corresponds to the true initial distribution, and since there are no 'interrupting" observations within I 0 , then the solution Eqs. (2) and (4) of Eq. (6) corresponds to the true distribution of X ( t ) for any t ∈ [ t i , t i +1 ) . Since u is constant over I 0 , then the prediction ˜ X ( t ) of the Eigen-SDE solver follows Eq. (4) accurately using the analytic integration (see Appendix B; note that if u were not constant, the solver would still follow the solution up to a numeric integration error). Regarding t 1 , according to Appendix A, ˜ X ( t 1 ) corresponds to the true distribution of X ( t 1 ) after conditioning on the observation ˆ Y ( t 1 ) (if there was an observation at t 1 ; otherwise, no filtering is needed). This completes the induction basis. Using the same arguments, if we assume for an arbitrary i ≥ 0 that ˜ X ( t i ) corresponds to the true distribution, then ˜ X ( t ) corresponds to the true distribution for any t ∈ I i = [ t i , t i +1 ] , completing the induction.

Now, for any t , since ˜ X ( t ) ∼ N ( µ ( t ) , Σ( t )) is in fact the true distribution of X ( t ) , the expected square error E [ SE ( t )] = E [ ( µ -X ( t )) 2 ] is minimized by choosing µ := µ ( t ) ; and the expected log-likelihood E [ ℓ ( t )] = E [log P ( X ( t ) | µ, Σ)] is maximized by µ := µ ( t ) , Σ := Σ( t ) .

E Extended Experiments

E.1 Ablation Study for Patient Individualization

To provide an insight over the importance of the dynamics-individualization, we perform an ablation study for the hypernetwork module. We use the same medical benchmarks as in Section 6, and fit a version of NESDE with neutralized hypernetwork module. In particular, we fix the context inputs of the module to be a vector of 1 s, and thus prevent any propagation from the context features to the model's output. The results are presented in Table 4, and show a great degradation in model

performance in the UH-dosing benchmark, approving that the hypernetwork indeed utilize the information within the context features. In the Vancomycin dosing benchmark, while we still observe a degradation comparing to NESDE, the version of NESDE without hypernetwork still outperforms LSTM in terms of MSE and the rest of the baselines (except NESDE) in terms of NLL.

E.2 Comparison to ODE-based Methods

Section 5 compares NESDE to GRU-ODE-Bayes [De Brouwer et al., 2019] - a recent ODE-based method that can provide an uncertainty estimation (which is a typical requirement in medical applications). Similarly to other recent ODE-based methods [Chen et al., 2018], GRU-ODE-Bayes relies on a non-linear neural network model for the differential equation. GRU-ODE-Bayes presents relatively poor prediction accuracy in Section 5, which may be partially attributed to the benchmark settings. First, the benchmark required GRU-ODE-Bayes to handle a control signal. As proposed in De Brouwer et al. [2019], we incorporated the control as part of the observation space. However, such a control-observation mix raises time synchrony issues (e.g., most training input samples include only control signal without observation) and even affect the training supervision (since the new control dimension in the state space affects the loss). Second, as discussed above, the piecewise linear dynamics of NESDE provide higher sample efficiency in face of the 1000 training trajectories in Section 5.

In this section, we explicitly study the sample efficiency of NESDE vs. GRU-ODE-Bayes in a problem with no control signal. Specifically, we generate data from the GitHub repository of De Brouwer et al. [2019]. The data consists of irregular samples of the two-dimensional Ornstein-Uhlenbeck process, which follows the SDE

where the noise follows a Wiener process, which is set in this experiment to have the covariance matrix

The process is sparsely-observed: we use a sample rate of 0 . 6 (approximately 6 observations for 10 time units). Each sampled trajectory has a time support of 10 time units. The process has two dimensions, and each observation can include either of the dimensions or both of them. The dynamics

of the process are linear and remain constant for all the trajectories; however, the stable 'center" of the dynamics of each trajectory (similarly to α in Eq. (6)) is sampled from a uniform distribution, increasing the difficulty of the task and requiring to infer α in an online manner.

Fig. 9 presents a sample of trajectory observations along with the corresponding predictions of the NESDE model (trained over 400 trajectories). Similarly to De Brouwer et al. [2019], the models are tested over each trajectory by observing all the measurements from times t ≤ 4 , and then predicting the process at the times of the remaining observations until the end of the trajectory.

-0.1

GRU_ODE

GRU_ODE

0.045

0.2

NESDE

NESDE

NLL

0.3

0.040

-0.4

0.035

0.5

0.030

0.6

100

200

400

800

1600

100

200

400

800

1600

train_size

train_size

train_size=400

train_size=400

0.2

0.045

GRU_ODE

0.3

0.040

NESDE

-0.4

NLL

0.5

M

E

0.035

0.030

-0.6

GRU_ODE

0.7

NESDE

0.025

0.020

4

5

6

7

8

9

10

4

5

6

7

8

6

10

time

time

To test for data efficiency, we train both models over training datasets with different numbers of trajectories. As shown in Fig. 10, the sparsely-observable setting with limited training data causes GRU-ODE-Bayes to falter, whereas NESDE learns robustly in this scenario. The advantage of NESDE over GRU-ODE-Bayes increases when learning from smaller datasets (Fig. 10, top), or when predicting for longer horizons (Fig. 10, bottom). This demonstrates the stability and data efficiency of the piecewise linear dynamics model of NESDE in comparison to non-linear ODE models.

E.3 Sparse Observations

This experiment addresses the sparsity of each trajectory. We use the same benchmark as in Section 5 and generate 4 train datasets, each one contains 400 trajectories, and a test set of 1000 trajectories. In each train-set, the trajectories have the same number of data samples, which varies between datasets (4,6,8,10). The test-set contains trajectories of varying number of observations, over the same support. For each train-set, we train all the models until convergence, and test them. Fig. 11 presents the MSE over the test set, for both the complex and the real eigenvalues settings. It is noticeable that even with very sparse observations, NESDE achieves good performance. Here, GRU-ODE-Bayes appears to be more sample-efficient than CRU and LSTM, but it is less sample efficient than NESDE.

0.6

NESDE

NESDE

GRUODE

GRUODE

0.4

LSTM

LSTM

E

E

S

M

0.4

CRU

S

M

CRU

0.3

0.2

4

6

8

10

4

6

8

10

Traindatadensity

Traindatadensity

[observations per trajectory]

[observations per trajectory]

(a) Complex dynamics

(b) Real dynamics

E.4 Synthetic Data Experiments with Regular Observations

While NESDE (and ODE-based models) can provide predictions at any point of time, a vanilla LSTM is limited to the predefined prediction horizon. Shorter horizons provide higher temporal resolution, but this comes with a cost: more recursive computations are needed per time interval, increasing both learning complexity and running time. For example, if medical measurements are available once per hour while predictions are required every 10 seconds, the model would have to run recursively 360 times between consecutive measurements, and would have to be trained accordingly in advance. We use the synthetic data environment from Section 5, in the complex dynamics setting, and test both regularly and out-of-distribution control (see Section 5). Here, we use LSTM models trained with resolutions of 1, 8 and 50 predictions per observation. All the LSTM models receive the control u and the current observation Y as an input, along with a boolean b o specifying observability: in absence of observation, we set Y = 0 and b o = 0 . The models consist of a linear layer on top of an LSTM layer, with 32 neurons between the two. To compare LSTMs with various resolutions, we work with regular samples, 10 samples, one at each second. The control changes in a 10 -2 seconds' resolution, and contains information about the true state.

In Fig. 12c we present a sample trajectory (without the control signal) with the predictions of the various LSTMs and NESDE. It can be observed that while NESDE provides continuous, smooth predictions, the resolution of the LSTMs must be adapted for a good performance. As shown in Fig. 12a, all the methods perform well from time t = 3 and on, still, NESDE and the low-resolution variants of LSTM attain the best results. The poor accuracy of the high-resolution LSTM demonstrates the accuracy-vs-resolution tradeoff in recursive models, moreover, GRUODE shows similar behavior in this analysis, which may hint on the recursive components within GRUODE.

(a) One-step prediction

(b) Out of distribution control

(c) Sample trajectory

1.0

NESDE:MSE=0.19

2.0

NESDE:MSE=0.19

GRU-ODE:MSE=0.35

GRU-ODE:MSE=0.43

3

signal

0.6

LSTM (1:1): MSE=0.18

CRU:MSE=0.20

1.0

LSTM (1:1): MSE=0.21

CRU:MSE=0.49

2

LSTM 1:50

LSTM 1:8

MSE

0.4

LSTM (1:8): MSE=0.19

LSTM (1:50): MSE=0.35

E

0.6

LSTM (1:8): MSE=0.22

LSTM (1:50): MSE=0.42

prediction

1

0

LSTM 1:1

NESDE

0.4

observation

-1

0.2

0.2

-2

-3

2

4

6

8

10

2

4

6

8

10

0

2

4

6

8

10

time

time

time

The out-of-distribution test results (Fig. 12b) show that a change in the control policy could result in major errors; while NESDE achieves errors which are close to Fig. 12a, the other methods deteriorate in their performance. Notice the scale difference between the figures. The high-resolution LSTM and the ODE-based methods suffer the most, and the low-resolution variants of the LSTM, demonstrate

robustness to the control change. This result is similar to the results we present in Section 5, although here we see similarities between the variants of the LSTM and the ODE-based methods.

E.5 Interpretability: Inspecting the Spectrum

In addition to explicit predictions at flexible times, NESDE provides direct estimation of the process dynamics, carrying significant information about the essence of the process.

For example, consider the following 3 processes, each with one observable variable and one latent variable: A 1 = ( -0 . 5 -2 2 -1 ) with the corresponding eigenvalues λ 1 ≈ -0 . 75 ± 1 . 98 i ; A 2 = ( -0 . 5 -0 . 5 -0 . 5 -1 ) with λ 2 ≈ ( -1 . 3 , -0 . 19) ⊤ ; and A 3 = ( 1 -2 2 -1 ) with λ 3 ≈ ± 1 . 71 i . As demonstrated in Fig. 4, the three processes have substantially different dynamics: roughly speaking, real negative eigenvalues correspond to decay, whereas imaginary eigenvalues correspond to periodicity.

For each process, we train NESDE over a dataset of 200 trajectories with 5-20 observations each. We set NESDE to assume an underlying dimension of n = 2 (i.e., one latent dimension in addition to the m = 1 observable variable); train it once in real mode (real eigenvalues) and once in complex mode (conjugate pairs of complex eigenvalues); and choose the model with the better NLL over the validation data. Note that instead of training twice, the required expressiveness could be obtained using n = 4 in complex mode (see Appendix C); however, in this section we keep n = 2 for the sake of spectrum interpretability.

As the processes have linear dynamics, for each of them NESDE learned to predict a consistent dynamics model: all estimated eigenvalues are similar over different trajectories, with standard deviations smaller than 0.1. The learned eigenvalues for the three processes are ˜ λ 1 = -0 . 77 ± 1 . 98 i ; ˜ λ 2 = ( -0 . 7 , -0 . 19) ⊤ ; and ˜ λ 3 = -0 . 03 ± 0 . 83 i . That is, NESDE recovers the eigenvalues class (complex, real, or imaginary), which captures the essence of the dynamics - even though it only observes one of the two dimensions of the process. The eigenvalues are not always recovered with high accuracy, possibly due to the latent dimensions making the dynamics formulation ambiguous.

E.6 Model Expressiveness and Overfitting

It is well known that more complex models are capable to find complex connections within the data, but are also more likely to overfit the data. It is quite common that a data that involves control is biased or affected by confounding factors: a pilot may change his course of flight because he saw a storm that was off-the-radar; a physician could adapt his treatment according to some measure that is off-charts. Usually, using enough validation data could solve the overfitting issue, although sometimes the same confounding effects show in the validation data, which results in a model that is overfitted to the dataset. When targeting a model for control adjustment, it is important that it would be robust to changes in the control; a model that performs poorly when facing different control is unusable for control tuning. To exemplify an extreme case of confounding factors in the context of control, we add a correlation between the control (observed at all times) to the predictable measure (observed sparsely), in particular at times that the predictable is unobserved. We harness the same synthetic data benchmark as in Section 5, and use regular time samples, and the same LSTM baselines as in Appendix E.4 but here we generate different two types of control signals:

Same Distribution (SD): at each time t , the control u ( t ) = b t -0 . 8 · Y t .

Out of Distribution (OOD): at each time t , the control u ( t ) = b t +0 . 8 · Y t .

b t is a random piecewise constant and Y t is the exact value of the measure we wish to predict. The first type is used to generate the train and the test sets, additionally we generate an out-of-distribution test-set using the second type. We observe in Fig. 13 that GRU-ODE-Bayes and the high-resolution LSTM achieve very low MSE over the SD as seen during training. CRU also achieves very low MSE, although not as much. The results over the OOD data show that the high performance over SD came with a cost - the better a model is over SD the worse it is over OOD. The results of LSTM 1:1 are not surprising, it sees the control signal only at observation-times, so it cannot exploit the hidden information within the control signal. NESDE does not ignore such information, while maintaining the robustness w.r.t. control.

10

0.6

NESDE: MSE=0.17

NESDE: MSE=0.28

0.5

GRU-ODE:MSE=0.002

CRU:MSE=0.13

8

GRU-ODE:MSE=5.56

CRU: MSE=0.71

MSI

E

0.3

0.4

LSTM (1:1): MSE=0.25

LSTM (1:8): MSE=0.04

LSTM (1:50): MSE=0.01

MSE

6

LSTM (1:1):MSE=0.29

LSTM (1:8):MSE=2.67

LSTM (1:50): MSE=2.75

4

0.2

0.1

2

0.0

0

2

4

6

8

10

2

4

6

8

10

time

time

(a) Same control distribution

(b) Out of distribution control

F Medication Dosing Prediction: Implementation Details

Below, we elaborate on the implementation details of Section 6.

F.1 Data preprocessing

Heparin: We derive our data from the MIMIC-IV dataset [Johnson et al., 2020], available under the PhysioNet Credentialed Health Data License. For the UH dosing dataset, we extract the patients that were given UH during their intensive care unit (ICU) stay. We exclude patients that were treated with discrete (not continuous) doses of UH, or with other anticoagulants; or that were tested for aPTT less than two times. The control signal (UH dosing rate) is normalized by the patient weight. Each trajectory of measurements is set to begin one hour before the first UH dose, and is split in the case of 48 hours without UH admission. This process resulted with 5866 trajectories, containing a continuous UH signal, an irregularly-observed aPTT signal, and discretized context features. Note that we do not normalize the aPTT values.

Vancomycin: The VM dosing dataset derived similarly, from patients who received VM during their ICU stay, where we consider only patients with at least 2 VMconcentration measurements. Each trajectory begins at the patient's admission time, and we also split in the case of 48 hours without VM dosage. Additionally, we add an artificial observation of 0 at time t = 0 , as the VM concentration is 0 before any dose was given (we do not use these observations when computing the error).

General implementation details: For each train trajectory, we only sample some of the observations, to enforce longer and different prediction horizons, which was found to aid the training robustness. Hyperparameters (e.g., learning rate) were chosen by trial-and-error with respect to the validation-set (separately for each model).

Context variables C are used in both domains. We extract 42 features, some measured continuously (e.g., heart rate, blood pressure), some discrete (e.g., lab tests, weight) and some static (e.g., age, background diagnoses). Each feature is averaged (after removing outliers) over a fixed time-interval of four hours, and then normalized.

F.2 LSTM Baseline Implementation

The LSTM module we use as a baseline has been tailored specifically to the setting:

It includes an embedding unit for the context, which is updated whenever a context is observed, and an embedded context is stored for future use.

The inputs for the module include the embedded context, the previous observations, the control signal and the time difference between the current time and the next prediction time.

Where the control signal is piecewise constant: any time it changes we produce predictions (even though no sample is observed) that are then used as an input for the model, to model the effect of the UH more accurately.

We train it with the same methodology we use for NESDE where the training hyperparameters chosen by the best performance over the validation data.

Architecture for the medication dosing benchmarks : The model contains two fully connected elements: one for the context, with two hidden layers of size 32 and 16 -dimensional output which is fed into a Tanh activation; the second one uses the LSTM output to produce a one-dimensional output, which is fed into a ReLU activation to produce positive outputs, its size determined by the LSTM dimensions. The LSTM itself has an input of 19 dimensions; 16 + 1 + 1 + 1 for the context, control, previous observations and the time interval to predict. It has a hidden size of 64 and two recurrent layers, with dropout of 0 . 2 . All the interconnections between the linear layers include ReLU activations.

Architecture for the synthetic data benchmarks : Here, there is no context, then the model contains one fully connected element that receives the LSTM output and has two linear layers of sizes 32 and 1 with a Tanh activation between them. The LSTM has an input of 3 dimensions; for the state, control signal, and the time interval to predict. It has a hidden size of 32 and two recurrent layers, with dropout of 0 . 2 .

F.3 Extended Results

The figures below present more detailed information for the experiments discussed in Section 6.

Peer reviews:
comments: SUMMARY:
This develops a robust SDE algorithm for medical forecasting. This is demonstrated on both synthetic and real data.

STRENGTHS:
The problem solved is clearly compelling and the performance is evaluated on a wide range of datasets, where it exhibits superior performance to reasonable alternatives. It would appear to be an overall robust method.

WEAKNESSES:
It would be nice if some of the components of the model such as hypernet were defined here rather than briefly described, even though I understand the desire to save space. It makes it so that this paper is more difficult to understand/review.

QUESTIONS:
SDEs are not my area of expertise. I expect the other reviewers will have more enlightening questions and suggestions.

LIMITATIONS:
No societal ramifications, but no limitations mentioned as far as shortcomings of the model are concerned.
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2023-07-05

comments: SUMMARY:
The paper presents a piece-wise linear neural SDE with control for biomedical forecasting. The key contribution is combining rigorous linear SDE solutions with hypernet parameterisation.

STRENGTHS:
The paper is fantastically written and very clear. 

The idea of linearising SDEs is relatively novel, and the overall model is sensible and principled.

The results are good, and show that a relatively simple model construction is perhaps ideal for medical data.

WEAKNESSES:
The method is somewhat incremental: it combines linear SDEs with hypernets.

QUESTIONS:
The imaginary/cyclicity argument could have been elaborated more in experiments. Is medical data cyclical? I'm not sure. 

How come in fig3 we get same performance with 2 or 8 datapoints?

LIMITATIONS:
No issues
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2023-07-12

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper proposes a novel network approach, named Neural Eigen Stochastic Differential Equations (NESDE), to solve sequential prediction problems mainly in medical dosing control field of applications.\n\nSTRENGTHS:\n$\\mathbf{1}.$ The proposal of NESDE is novel, along with the hypernet that determines the subsequent model parameters.\n\n$\\mathbf{2}.$ This paper is very well written. The motivation, i.e., to address challenges in sequential prediction of medical dosing control, and the limitation of current works, is very well explained. The related works and background fundamental theories are also very well stated with intuitive illustrations. \n\n$\\mathbf{3}.$ The theoretical analysis is comprehensive and sound (also seen in $\\mathbf{Appendix}$.) \n\n$\\mathbf{4}.$ The empirical analysis is detailed and thorough. Author designed the experiments to validity the algorithmic merits and also real life application in the medical dosing field.\n\nWEAKNESSES:\n$\\mathbf{Note: }.$ It is a unfamiliar application field to me, and thus it is hard to determine whether or not the proposed method is indeed a significantly novel approach in the field. \n\n$\\mathbf{1}.$ The major concern is whether or not the proposed method has a broader application. It seems that the direct counterpart is neural ordinary differential equations, and is it possible to discuss and compare the general applicabilities of NODEs versus NESDEs?\n\nQUESTIONS:\n$\\mathbf{1}.$ See $\\mathbf{1}.$ in $\\mathbf{Weakness}$.\n\n$\\mathbf{2}.$ The experiment with synthetic data seems to support, e.g., the sample efficiency of the proposed method. However, it lacks training efficiency/computational complexity in the experiments. Is it possible to quantify the computational cost when comparing with other methods?   \n\n$\\mathbf{Trivial}.$ What is the point of comparison in Figure 1?\n\nLIMITATIONS:\nThis paper does not provide broader impact or limitation statements, but i do not have any concerns on potential negative societal impact.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 4, "DATE": "2023-07-12"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper presents a method for predicting time series at irregular time points. Its approach systematically integrates control signals, a convenient attribute for modelling dosing regimes in medicine. The latent process follows a piecewise linear stochastic differential equation (SDE) with control. The SDE is parameterized using its spectral representation, allowing for effective and analytical propagation. The parameters of the SDE are learned by a Hypernet. Experimental validation of the method is conducted on artificial data and two medication dosing problems.\n\nSTRENGTHS:\n- The paper introduces a method that effectively models control, with experimental results indicating an improvement over baseline models lacking explicit control.\n- The well-structured experimental study underscores the applicability and value of the method for two specific medical dosing scenarios.\n- The authors propose an efficient technique to compute dynamics by parameterizing the system's eigenfunction, thereby enabling fast and analytical computation.\n- This parameterization strategy not only facilitates the integration of prior knowledge on the dynamics, as evidenced in the dosing experiments, but also enhances interpretability—which is a useful attribute in practise\n\nWEAKNESSES:\n- The related work sections seems comprehensive, though neural SDE approaches could also be discussed as the model is based on a SDE latent process.\n- Explicitly stating which equations represent ESDE() and Filter() in Algorithm 1 would enhance clarity.\n- Although the experimental study appears to be well-executed and thoughtfully designed, the paper could be enhanced by including more baseline comparisons. Specifically, a comparison with NeuralSDE models would be intriguing, as they also yield uncertainty estimates and aren't confined to the piece-wise linearity constraint.\n- Minor: My understanding is that the context does not change over time, though Figure 2 suggest this on the first sight.\n\nQUESTIONS:\n--\n\nLIMITATIONS:\nThe piece-wise linear assumption on the observation sequence $Y(t)$ seems a restrictive assumption and limits the type of dynamics that can be modeled with the method. (From my understanding, the model learns a piece-wise linear SDE for $X(t)$ where the observations $Y(t)$ is simply a subset of the features of $X(t)$ , i.e. $Y(t)=X(t)_{1:m}$  . Therefore, the sole non-linearity allowed in the observation sequence come from the piece-wise linear formulation.) A discussion on when the piece-wise linear assumption fails compared to other (more expressive) models is encouraged.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-26"}