SlowGuess/NeurIPS_2023_OCR huggingface dataset:
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  $\bar{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  $\bar{v}$  (the complex conjugate of  $v$ ) is the eigenvector that corresponds to  $\bar{z}$ . Secondly, the eigenfunction matrix takes the form:
$$
\Phi (t) = e ^ {a t} \left( \begin{array}{c c} | & | \\ v _ {r e a l} \cdot c o s (b t) - v _ {i m} \cdot s i n (b t) & v _ {i m} \cdot c o s (b t) + v _ {r e a l} \cdot s i n (b t) \\ | & | \end{array} \right)
$$
For brevity, we consider only the elements that correspond with  $z$ ,  $\bar{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:
$$
\Phi (t) = e ^ {a t} \left( \begin{array}{c c} | & | \\ v _ {r e a l} & v _ {i m} \\ | & | \end{array} \right) \left( \begin{array}{c c} c o s (b t) & s i n (b t) \\ - s i n (b t) & c o s (b t) \end{array} \right)
$$
$$
\Phi (t) ^ {- 1} = e ^ {- a t} \left( \begin{array}{c c} c o s (b t) & - s i n (b t) \\ s i n (b t) & c o s (b t) \end{array} \right) \left( \begin{array}{c c} | & | \\ v _ {r e a l} & v _ {i m} \\ | & | \end{array} \right) ^ {- 1}
$$
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  $\mathcal{I}_i = [t_i, t_{i+1}]$ , and  $u(t)$  a control signal that is constant over  $\mathcal{I}_i$  for any  $i$ . For any  $i$ , consider the Eigen-SDE solver with the parameters corresponding to Eq. (6) (for the same  $\mathcal{I}_i$ ). Assume that the first solver ( $i = 0$ ) is initialized with the true initial distribution  $X(0) \sim N(\mu_0, \Sigma_0)$ , and for  $i \geq 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 \in \mathcal{I}_i$ , consider the prediction  $\tilde{X}(t) \sim N(\mu(t), \Sigma(t))$  of the solver. Then,  $\mu(t)$  minimizes the expected square error of the signal  $X(t)$ , and  $\tilde{X}(t)$  maximizes the expected log-likelihood of  $X(t)$ .
Proof. We prove by induction over  $i$  that for any  $i$  and any  $t \in \mathcal{I}_i$ ,  $\tilde{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  $\mathcal{I}_0$ , then the solution Eqs. (2) and (4) of Eq. (6) corresponds to the true distribution of  $X(t)$  for any  $t \in [t_i, t_{i+1})$ . Since  $u$  is constant over  $\mathcal{I}_0$ , then the prediction  $\tilde{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,  $\tilde{X}(t_1)$  corresponds to the true distribution of  $X(t_1)$  after conditioning on the observation  $\hat{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 \geq 0$  that  $\tilde{X}(t_i)$  corresponds to the true distribution, then  $\tilde{X}(t)$  corresponds to the true distribution for any  $t \in \mathcal{I}_i = [t_i, t_{i+1}]$ , completing the induction.
Now, for any  $t$ , since  $\tilde{X}(t) \sim N(\mu(t), \Sigma(t))$  is in fact the true distribution of  $X(t)$ , the expected square error  $E[SE(t)] = E[(\mu - X(t))^2]$  is minimized by choosing  $\mu \coloneqq \mu(t)$ ; and the expected log-likelihood  $E[\ell(t)] = E[\log P(X(t)|\mu, \Sigma)]$  is maximized by  $\mu \coloneqq \mu(t), \Sigma \coloneqq \Sigma(t)$ .
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 1s, 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
Table 4: Test mean square errors (MSE) and negative log-likelihood (NLL) in the medication-dosing benchmarks. This is an extension of Table 3 with the additional results of NESDE without the hypernetwork.
Model	UH Dosing	Vancomycin Dosing
MSE	NLL	MSE	NLL
Naive	613.3 ± 13.48	-	112.2 ± 16.4	-
LSTM	482.1 ± 6.52	-	92.89 ± 11.3	-
GRU-ODE-Bayes	491 ± 6.88	4.52 ± 0.008	80.54 ± 11.8	6.38 ± 0.12
CRU	450.4 ± 8.27	4.49 ± 0.012	76.4 ± 12.8	3.87 ± 0.2
NESDE - no hypernet	529.7 ± 13.34	5.42 ± 0.067	87.32 ± 11.57	3.73 ± 0.13
NESDE (ours)	411.2 ± 7.39	4.43 ± 0.01	70.71 ± 12.3	3.69 ± 0.13
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.
Figure 9: A sample test trajectory of the sparsely-observable OU process. The observations and the NESDE predictions (based on training over 400 trajectories) are presented separately for each of the two dimensions of the process.
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
$$
d x _ {t} = \theta (\mu - x _ {t}) d t + \sigma d W t,
$$
where the noise follows a Wiener process, which is set in this experiment to have the covariance matrix
$$
C o v = \left( \begin{array}{c c} 1 & 0. 5 \\ 0. 5 & 1 \end{array} \right).
$$
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  $\alpha$  in Eq. (6)) is sampled from a uniform distribution, increasing the difficulty of the task and requiring to infer  $\alpha$  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 \leq 4$ , and then predicting the process at the times of the remaining observations until the end of the trajectory.
Figure 10: Top: losses of NESDE and GRU-ODE-Bayes over the OU benchmark, along with confidence intervals of  $95\%$  over the test trajectories. NESDE demonstrates higher data efficiency, as its deterioration in small training datasets is moderate in comparison to GRU-ODE-Bayes. Bottom: errors vs. time, given 400 training trajectories, where all the test predictions rely on observations from times  $t \leq 4$ . The advantage of NESDE becomes larger as the prediction horizon is longer.
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.
(b) Real dynamics
Figure 11: Test MSE vs. train observations-per-trajectory.  $95\%$  confidence intervals are calculated over 1000 test trajectories.
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
Figure 12: MSE for predictions, relying on the whole history of the trajectory for (a) the test set, and (b) out-of-distribution test set. The uncertainty corresponds to 0.95-confidence-intervals over 1000 trajectories. (c) Sample trajectory and predictions. The LSTM predictions are limited to predefined times (e.g., LSTM 1:1 only predicts at observation times), but their predictions are connected by lines for visibility. The shading corresponds to NESDE uncertainty (note that the LSTM does not provide uncertainty estimation).
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} = \left( \begin{array}{cc} -0.5 & -2\\ 2 & -1 \end{array} \right)$  with the corresponding eigenvalues  $\lambda_{1}\approx -0.75\pm 1.98i$ ;  $A_{2} = \left( \begin{array}{cc} -0.5 & -0.5\\ -0.5 & -1 \end{array} \right)$  with  $\lambda_{2}\approx (-1.3, - 0.19)^{\top}$ ; and  $A_{3} = \left( \begin{array}{cc}1 & -2\\ 2 & -1 \end{array} \right)$  with  $\lambda_{3}\approx \pm 1.71i$ . 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  $\tilde{\lambda}_1 = -0.77 \pm 1.98i$ ;  $\tilde{\lambda}_2 = (-0.7, -0.19)^\top$ ; and  $\tilde{\lambda}_3 = -0.03 \pm 0.83i$ . 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:
1. Same Distribution (SD): at each time  $t$ , the control  $u(t) = b_t - 0.8 \cdot Y_t$ .
2. Out of Distribution (OD): at each time  $t$ , the control  $u(t) = b_t + 0.8 \cdot 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.
(a) Same control distribution
(b) Out of distribution control
Figure 13: MSE for predictions under regular time samples, where the control signal is correlated to the measure we wish to predict, even in times when it is unobserved. (a) Shows the results for a test set that has the same correlation between the control and the predictable measure as in the train set. (b) present the MSE for a different test set, with different correlation. Notice the different scales of the graphs.
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 VM concentration 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:
1. 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.
2. 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.
3. 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