SlowGuess/NeurIPS_2023_OCR huggingface dataset:
generative variables,  $\theta = V(\nu ,\phi)$  where  $\nu \sim p(\nu)\Rightarrow \theta \sim q_{\phi}(\theta)$ . Then, building on Lemma 1, we can arrive at an unbiased estimate for the gradient of the modified ELBO which inherits all the properties discussed with the original approximation:
$$
\begin{array}{l} \nabla_ {\phi} \mathcal {L} (\phi) \approx \frac {N}{R} \sum_ {i = 1} ^ {M} \sum_ {k = 1} ^ {R} \nabla_ {\phi} \log p _ {V (\nu^ {(k)}, \phi)} (x _ {i} ^ {(j)} \mid T (t _ {i} ^ {(j)}, \epsilon^ {(k)}, \phi)) \\ - \left(t _ {1} ^ {(j + 1)} - t _ {1} ^ {(j)}\right) \frac {N}{2 R} \sum_ {k = 1} ^ {R} \nabla_ {\phi} | | r _ {V (\nu^ {(k)}, \phi), \phi} (T (t ^ {(k)}, \epsilon^ {(k)}, \phi), t ^ {(k)}) | | _ {C _ {V (\nu^ {(k)}, \phi)} ^ {2}} ^ {2} (t) \tag {25} \\ - D _ {K L} \left(q _ {\phi} (\theta) \mid \mid p (\theta)\right). \\ \end{array}
$$
where each  $t^{(k)}\sim \mathcal{U}(t_1^{(j)},t_1^{(j + 1)})$  , each  $\epsilon^{(k)}\sim \mathcal{N}(0,I)$  , and each  $\nu^{(k)}\sim p(\nu)$
D Detailed description of recognition network
This section details the deep kernel based encoder we used in all numerical studies. We found this particular encoding architecture to be useful for our purposes because it is interpretable and stable in training. With this being said, any encoder which can transform a batch of observations down to a reduced dimension latent state can be used in combination with Lemma 1 to arrive at an unbiased estimate for the gradient of the ELBO which retains all the properties discussed in the main body of this work.
Given a dataset,  $\mathcal{D} = \{(t_i,x_i)\}_{i = 1}^N$ , recall that the first step in our amortization strategy is to split the dataset into  $N / M$  non-overlapping partitions,
$$
\text {o r i g i n a l} \quad [ t _ {1}, t _ {2}, \dots , t _ {M}, t _ {M + 1}, \dots , t _ {N} ]
$$
$$
\text {r e i n d e x e d d a t a s e t :} \quad [ t _ {1} ^ {(1)}, t _ {2} ^ {(1)}, \dots , t _ {M} ^ {(1)}, t _ {1} ^ {(2)}, \dots , t _ {M} ^ {(N / M)} ]
$$
Recall that we would like to approximate the latent state over each partition using only the  $M$  observations in each partition,  $q_{\phi}(z(t) \mid x_1^{(j)}, x_2^{(j)}, \ldots, x_M^{(j)}) \approx p(z(t) \mid \mathcal{D})$  for  $t \in [t_1^{(j)}, t_1^{(j+1)}]$ . Going forward we will drop writing the superscript as we will be only working on a single partition,  $(t_1, x_1), (t_2, x_2), \ldots, (t_M, x_M)$ . The user first selects how many snapshots into the future they would like to use to estimate the latent state at the present time,  $K$ ; in our own studies we found choosing  $K \in [1, 5]$  worked well for the problems we considered.
Given some encoding network,  $\mathrm{ENC}_{\phi}$ , we compute:
$$
h _ {i} = \operatorname {E N C} _ {\phi} \left(x _ {i}, x _ {i + 1}, \dots , x _ {i + K}\right). \tag {26}
$$
where  $h_i \in \mathbb{R}^{2d}$  with  $d < D$ . We will describe the particular architecture design for  $\mathrm{ENC}_{\phi}$  in the context of each numerical study in Appendix G. We note that for our approach, it is often important to use at least a small number of neighbours (i.e.  $K > 0$ ) when estimating the latent state because we are limited to approximating MGPs over the latent state.
To explain this point more clearly, let us consider the example of inferring a latent SDE using video of a pendulum as we did in Section 4.4. If we choose a single frame near the center, we have no way of knowing from that frame alone if the pendulum is currently swinging left or right. In other words, if we were to build an encoder which takes in one single frame, the encoder should predict that the posterior given that single frame is multimodal. As our approach only allows for one to approximate the latent state using MGPs, this is not an option. Allowing the encoder to take in a few frames at a time remedies this issue. We also note that previous works for inferring latent differential equations made this same choice [1, 11, 12, 17].
Recall that we need to approximate the latent state at any time over the window of observations in the partition,  $t \in [t_1^{(j)}, t_1^{(j+1)}]$ . To accomplish this we effectively interpolate between encodings using a deep kernel [36]. Letting,
$$
H = \left[ \begin{array}{c} h _ {1} ^ {T} \\ h _ {2} ^ {T} \\ \vdots \\ h _ {M} ^ {T} \end{array} \right], \tag {27}
$$
where  $H \in \mathbb{R}^{M \times 2d}$  and  $t_{\mathrm{node}} = [t_i, t_1, \dots, t_M]$ , we construct the encoder for the mean and diagonal of the covariance over the latent state as,
$$
\left[ m _ {\phi} (t) ^ {T} \quad \log S _ {\phi} (t) ^ {T} \right] = k _ {\phi} \left(t, t _ {\text {n o d e}}\right) ^ {T} \left(k _ {\phi} \left(t _ {\text {n o d e}}, t _ {\text {n o d e}}\right) ^ {- 1} + \sigma_ {n} ^ {2} I\right) ^ {- 1} H. \tag {28}
$$
Here we note that the right hand side of (28) is a row vector of length  $2d$ . We use the notation  $\left[m_{\phi}(t)^T \log S_{\phi}(t)^T\right]$  to indicate that  $m_{\phi}(t)$  is given by the first  $d$  elements of this vector and  $\log S_{\phi}(t)$  is given by the next  $d$  elements. Here  $k_{\phi}$  is a so-called deep kernel and  $k_{\phi}(t, t_{\mathrm{node}}) \in \mathbb{R}^{M}$  and  $k_{\phi}(t_{\mathrm{node}}, t_{\mathrm{node}}) \in \mathbb{R}^{M \times M}$ . In addition  $\sigma_n \in \mathbb{R}^+$  is tuned as a part of the optimization procedure. While many options for the base kernel are possible, we made use of a squared exponential kernel,
$$
k _ {\phi} (t, t _ {*}) = \sigma_ {f} \exp \left(- \frac {| | \mathrm {D K} _ {\phi} (t) - \mathrm {D K} _ {\phi} (t _ {*}) | | ^ {2}}{2 \ell^ {2}}\right), \tag {29}
$$
where  $\sigma_f, \ell \in \mathbb{R}^+$  are positive constants tuned as a part of the optimization procedure and  $\mathrm{DK}_{\phi}$  is a neural network whose architecture we will describe in the context of the numerical studies. While the base kernel is stationary, the neural networks allow for the encoder to infer non-stationary relationships [36]. It is worth noting that without this deep-kernel our approach struggled to achieve good validation accuracy on the datasets we considered.
Advantages of this encoder design are that it can easily take in varying amounts of data, it is interpretable because  $[m_{\phi}(t_i)^T\log S_{\phi}(t_i)^T]\approx h_i^T$  , and it is cheap to compute so long as  $M$  is relatively small because evaluations of  $\mathrm{ENC}_{\phi}$  can be performed in parallel. Particular choices for  $\mathrm{ENC}_{\phi}$  and  $\mathrm{DK}_{\phi}$  are described in context in Appendix G.
E Approximate posterior on diffusion term
This section summarizes the log-normal parameterization we used to approximate the posterior over the diagonal drift function terms in the main body of the paper. First, we note that in the loss function we only require access to the product,  $C_{\theta}(t)^{-1} = L_{\theta}(t)\Sigma L_{\theta}(t)^{T}$  so, rather than parametrizing  $L_{\theta}$  on its own, we parametrize  $C_{\theta}^{-1}$ .
Specifically we parametrize  $C_{\theta}(t)^{-1} = \mathrm{diag}(\theta)$ , where  $\theta \in \mathbb{R}^d$ . The prior is defined as,
$$
p (\theta) = \prod_ {i = 1} ^ {d} \mathcal {L N} \left(\theta_ {i} \mid \tilde {\mu} _ {i}, \tilde {\sigma} _ {i} ^ {2}\right). \tag {30}
$$
We parametrize the approximate posterior as,
$$
q _ {\phi} (\theta) = \prod_ {i = 1} ^ {d} \mathcal {L N} \left(\theta_ {i} \mid \mu_ {i}, \sigma_ {i} ^ {2}\right), \tag {31}
$$
where  $\mu_{i}$  and  $\sigma_{i}$  are the variational parameters. The KL divergence between the posterior and prior is given by,
$$
D _ {K L} \left(q _ {\phi} (\theta) \mid | p (\theta)\right) = \sum_ {i = 1} ^ {d} \left(\log \tilde {\sigma} _ {i} - \log \sigma_ {i}\right) - \frac {1}{2} \left(d - \sum_ {i = 1} ^ {d} \frac {\sigma_ {i} ^ {2} - \left(\mu_ {i} - \tilde {\mu} _ {i}\right) ^ {2}}{\tilde {\sigma} _ {i} ^ {2}}\right). \tag {32}
$$
F Computing resources
Experiments were performed on an Ubuntu server with a dual E5-2680 v3 with a total of 24 cores, 128GB of RAM, and an NVIDIA GeForce RTX 4090 GPU. The majority of our code is written in PyTorch [37]. For benchmarking we made use of torchdiffreq [1], torchsde [12, 35], and pytorch_lightning. All code is available at github.com/coursekevin/arlatentsde.
G Details on numerical studies
This section contains more details on the numerical studies including the specific architecture design and hyperparameter selection for each experiment. For each experiment we design an encoder
consisting of two neural networks,  $\mathrm{ENC}_{\phi}(x_i,\ldots ,x_{i + K})$  and  $\mathrm{DK}_{\phi}(t)$  (see Appendix D), a decoder  $\mathrm{DEC}_{\theta}(z(t))$ , and a model for the SDE consisting of a drift,  $f_{\theta}(t,z)$ , and dispersion matrix,  $L_{\theta}(t)$ . For all numerical studies we place a log-normal prior on the dispersion matrix and assume that the approximate posterior is constant in time, see Appendix E. For all experiments we gradually increased the value of the KL-divergence (both the KL-divergence due to the SDE prior and the KL-divergence on the dispersion matrix parameters) from 0 to 1 using a linear schedule.
G.1 Orders of magnitude fewer function evaluations
In this section we provide a more detailed description of the numerical study in Section 4.1. As a reminder, we tasked a neural ODE (NODE) and our approach with building a predictive model for the Lotka-Volterra system given a dataset of time-series observations. The Lotka-Volterra equations are a system of nonlinear ODEs usually written as,
$$
\begin{array}{l} \dot {x} = \alpha x - \beta x y, \\ \cdot \quad s \end{array} \tag {33}
$$
$$
\dot {y} = \delta x y - \gamma y.
$$
In our experiment we chose  $\alpha = 2/3$ ,  $\beta = 4/3$ , and  $\delta = \gamma = 1$ . We also assumed that there was some small amount of Brownian noise given by  $\Sigma = \mathrm{diag}(10^{-3}, 10^{-3})$ . Using the initial condition  $x = 0.9$  and  $y = 0.2$ , we draw a sample from the system using the default adaptive stepping scheme in torchside from time 0 to 65 with an initial step size of 0.1 and an absolute and relative tolerance of  $10^{-5}$ . We then evaluate the solution at a frequency of  $50\mathrm{Hz}$  and added Gaussian noise with a standard deviation of 0.01. We use the first 50 seconds for training and reserve the remaining 15 seconds for validation.
Both the NODE and our approach use the same model for the gradient field and drift function respectively, see Figure 6a. The encoder and deep kernel architecture are provided in Figures 6b and 6c respectively. As mentioned in the main body of the paper, we set the decoder to be the identity function so as to force our model to learn the dynamics in the original coordinates. We selected a Gaussian likelihood with a constant standard deviation of 0.01.
In terms of hyperparameters we set the schedule on the KL-divergence to increase from 0 to 1 over 1000 iterations. We choose a learning rate of  $10^{-3}$  with exponential learning rate decay where the learning rate was decayed  $lr = \gamma lr$  every iteration with  $\gamma = \exp (\log (0.9) / 1000)$  (i.e. the effective rate of learning rate decay is  $lr = 0.9lr$  every 1000 iterations.). We used the nested Monte-Carlo approximation described in Equation (10) with  $R = 1$ ,  $S = 10$ , and  $M = 256$ . In terms of kernel parameters, we initialized  $\ell = 10^{-2}$ ,  $\sigma_f = 1$ , and  $\sigma_n = 10^{-5}$ . In terms of the diffusion term, set  $\mu_{i} = \sigma_{i} = 10^{-5}$  and  $\tilde{\mu}_i = \tilde{\sigma}_i = 1$ .
G.2 Adjoint instabilities experiment
In this section, we provide some additional details of the numerical study described in Section 4.2. Recall the parametric model for the Lorenz system in equations (11-13). As a reminder, given time-series dataset of observations, our goal was to infer the value of the parameters,  $\sigma$ ,  $\beta$ ,  $\rho$ , which were likely to have generated the data starting from an initial guess:  $\theta_0 = [\sigma_0, \beta_0, \rho_0]$ . The true value of the parameters was chosen as  $\theta_* = [10, 8/3, 28]$ . For all experiments we used the initial condition [8, -2, 36.05] as was suggested in [10]. We generated data by solving the differential equation using scipy's RK4(5) initial value problem solver with a relative and absolute tolerance of  $10^{-6}$  and  $10^{-8}$  respectively. We generated data at a frequency of  $200\mathrm{Hz}$  over the time intervals [0, 1], [0, 10], [0, 50], and [0, 100]. For each time interval we generated 5 datasets by adding independent Gaussian noise with a variance of 1 to the data.
To arrive at an initial guess we sample from the distribution  $\theta_0\sim \mathcal{N}(\theta_{*},(0.2\theta_{*})^{2})$ . For each time series length we tasked our approach with inferring the true value of the parameters given 5 different guesses for the initial condition (i.e. one guess / dataset). The reported gradients for our approach are given by the average  $\ell_2$ -norm of the gradient of the ELBO with respect to the parameters  $\sigma, \beta$ , and  $\rho$  after optimizing for 2000 iterations. For the adjoint method, we report the gradient of the function:
$$
\mathcal {L} (\theta) = \frac {1}{3 N} \sum_ {i = 1} ^ {N} \left(x _ {i} - x _ {\theta} \left(t _ {i}\right)\right) ^ {2} + \left(y _ {i} - y _ {\theta} \left(t _ {i}\right)\right) ^ {2} + \left(z _ {i} - z _ {\theta} \left(t _ {i}\right)\right) ^ {2}, \tag {34}
$$
Figure 6: Architecture diagrams for the drift, deep kernel, and encoder used in the Lotka-Volterra problem are provided in Figures (a), (b), and (c) respectively. Note that we have used the shorthand  $m(t_{i})$ ,  $\log S(t_{i})$  to show how we have split the columns of  $h_i$  in two. The value of  $[m(t_i),\log S(t_i)]$  only  $\approx h_i$  unless  $\sigma_n = 0$ , see Appendix D. Note the arrow from  $x_{i}$  to  $m(t_{i})$  indicates a residual connection (which was useful in this case because we are learning a SDE in the original data coordinates).
Figure 7: Architecture diagrams for the deep kernel and encoder used in the Lorenz system parameter tuning problem are provided in Figures (a) and (b) respectively. Note that we have used the shorthand  $m(t_{i})$ ,  $\log S(t_{i})$  to show how we have split the columns of  $h_i$  in two. The value of  $[m(t_i),\log S(t_i)]$  only  $\approx h_i$  unless  $\sigma_n = 0$ , see Appendix D. Note the arrow from  $x_{i}$  to  $m(t_{i})$  indicates a residual connection (which was useful in this case because we are learning a SDE in the original data coordinates).
at the starting iteration. Note we cannot provide average gradients over the entire optimization procedure for the adjoint based method because the initial gradients are too large when the time interval was [0, 50] or [0, 100]. The error bars are given by one standard deviation from the mean.
A description of the encoder architecture is provided in Figure 7. As was the case in the previous experiment, we set the decoder to be the identity function and placed a log-normal prior on the diffusion term.
In terms of hyperparameters, we choose  $M = 128$ ,  $R = 10$ , and  $S = 100$ . We selected a learning rate of 0.1 and decayed the learning rate  $lr = \gamma lr$  every iteration with  $\gamma = \exp(\log(0.9)/1000)$ . We linearly increased the value of the KL-divergence from 0 to 1 over the course of 100 iterations. In terms of kernel parameters we initialized  $\sigma_f = 1$ ,  $\ell = 10^{-2}$ , and  $\sigma_n = 10^{-5}$ . For the diffusion term, we set  $\mu_i = \sigma_i = 10^{-5}$  and  $\tilde{\mu}_i = \tilde{\sigma}_i = 10^{-5}$ .
G.3 MOCAP experiment
For this experiment we use the preprocessed dataset provided by [34].
Like in previous examples, we place a log-normal prior on the diffusion term. Like previous works making use of the benchmark, we assume a Gaussian likelihood. We place a log-normal prior on the variance of the likelihood. A description of the architecture is provided in Figure 8. Note that the architecture we chose is very similar to the architecture used by [34, 12].
In terms of hyperparameters, we chose a batch size of 512. We used a linear schedule to update the weighting on the KL-divergence from 0 and 1 over the course of 200 iterations. We make use of the nested Monte-Carlo scheme in equation (10) with  $R = 10$  and  $S = 10$ . We chose a learning rate of 0.01 and decayed the learning rate each iteration according to the schedule  $lr = \gamma lr$  with  $\gamma = \exp(\log(0.9)/1000)$ . In terms of kernel parameters, we initialize  $\sigma_n = 10^{-5}$ ,  $\ell = 10^{-2}$ , and  $\sigma_f = 1$ . In terms of the prior on the diffusion term we initialize  $\mu_i = \sigma_i = 10^{-5}$  and set  $\tilde{\mu}_i = \tilde{\sigma}_i = 1$ . In terms of the prior on the variance of the observations, we initialize  $\mu_i = \sigma_i = 10^{-2}$  and set  $\tilde{\mu}_i = \tilde{\sigma}_i = 1$ . We considered both Softplus [12] and tanh [34] nonlinearities and found that tanh nonlinearities provided improved validation performance. We train for 100 epochs testing validation accuracy every 10 epochs. We report the average test accuracy after training 10 models from different random seeds.
Previous studies tended to report mean-squared-error as, MSE ± ERROR. We report RMSE ± NEW ERROR so that error units are consistent with the units of the original dataset. To convert MSE to RMSE we used a first-order Taylor-series approximation,
$$
\mathrm {R M S E} = \sqrt {\mathrm {M S E}}
$$
$$
\text {N E W E R R O R} = \frac {1}{2} \text {E R R O R / R M S E} \tag {35}
$$
G.4 Neural SDE from video
In this section, we provide a more detailed description on the numerical study described in Section 4.4. We generated data by simulating a nonlinear pendulum with the equations,
$$
\dot {x} = p
$$
$$
\dot {p} = - \sin (x), \tag {36}
$$
for 30 seconds while sampling the state at a frequency of  $15\mathrm{Hz}$ . The architecture we used for this experiment is provided in Figure 9.
In terms of hyperparameters, we chose a batch size of 128. We gradually increased the weighting of the KL-divergence term using a linear schedule over the course of 1000 iterations. We used the nested Monte-Carlo method suggested in Equation 10 and set  $R = 20$  and  $S = 10$ . We chose a learning rate of 0.001 and decayed the learning rate each iteration according to the schedule  $lr = \gamma lr$  with  $\gamma = \exp(\log(0.9)/1000)$ . In terms of kernel parameters, we initialized  $\sigma_f = 1$ ,  $\sigma_n = 10^{-5}$ , and  $\ell = 10^{-2}$ . We placed a log-normal prior on the diffusion term and approximated the posterior using log-normal variational distribution. With regards to the prior on the diffusion term we initialize  $\mu_i = \sigma_i = 0.1$  and set  $\tilde{\mu}_i = \tilde{\sigma}_i = 10^{-5}$ .
(c)  $\mathrm{DK}_{\phi}(t)$
(b)  $\mathbf{DEC}_{\phi}(t)$
(d)  $f_{\theta}(z,t)$
Figure 8: Architecture diagrams for the encoder, decoder, deep kernel and drift function used in the MOCAP benchmark. We used very similar architectures to [34, 12]. Here  $\mu(t)$  indicates the mean of the likelihood,  $p_{\theta}(x \mid z(t))$ .
(a)  $\mathsf{ENC}_{\phi}(x_i)$
(c)  $\mathrm{DK}_{\phi}(t)$
Figure 9: Architecture diagrams for the encoder, decoder, deep kernel and drift function used in the Neural SDE from video example.
(d)  $f_{\theta}(z,t)$
H Numerical study on the effect of Monte-Carlo parameters
In this study we investigate the effect of varying the number of nested Monte Carlo samples on the rate of validation RMSE convergence. To do so, we consider the problem of building a predictive model for a four-dimensional predator-prey system. The system consists of two predators and two prey where there is a competitive dynamic between the predators. The equations governing the dynamics of the system are,
$$
\dot {x} _ {1} = x _ {1} \left(\alpha_ {1} - \beta_ {1} y _ {1} - \gamma_ {1} y _ {2}\right) \left(1 - \frac {x _ {1}}{k _ {1}}\right), \tag {37}
$$
$$
\dot {x} _ {2} = x _ {2} \left(\alpha_ {2} - \beta_ {2} y _ {1} - \gamma_ {2} y _ {2}\right) \left(1 - \frac {x _ {2}}{k _ {2}}\right), \tag {38}
$$
$$
\dot {y} _ {1} = y _ {1} \left(- \delta_ {1} + \epsilon_ {1} x _ {1} + \xi_ {1} x _ {2} - \nu_ {1} y _ {2}\right), \tag {39}
$$
$$
\dot {y} _ {2} = y _ {2} \left(- \delta_ {2} + \epsilon_ {2} x _ {1} + \xi_ {2} x _ {2} + \nu_ {2} y _ {1}\right), \tag {40}
$$
where  $\alpha_{i}$  is the grow rate of the prey  $x_{i},\beta_{i}$  is the rate that predator  $y_{1}$  is consuming prey  $x_{i},\gamma_{i}$  is the rate that predator  $y_{2}$  is consuming prey  $x_{i},k_{i}$  is the carrying capacity for prey  $x_{i},\delta_{i}$  is the death rate of predator  $y_{i},\epsilon_{i}$  is the conversion rate for predator  $y_{i}$  from  $x_{1},\xi_{i}$  is the conversion rate for predator  $y_{i}$  from prey  $x_{2}$ , and  $\nu_{i}$  represents the competitive effects on  $y_{i}$  caused by the other predator.
We simulated the system for 300 units of time collecting data at a frequency of  $10\mathrm{Hz}$ . We assume a Gaussian noise with a standard deviation of  $10^{-2}$ . For all experiments we used the following hyperparameters: a batch size of 256, 1000 warmup iterations, and a learning rate of  $10^{-3}$ . In terms of kernel parameters, we initialized  $\sigma_f = 1$ ,  $\ell = 10^{-2}$ , and  $\sigma_n = 10^{-5}$ . For the diffusion term we set  $\mu_i = \sigma_i = 10^{-5}$  and  $\tilde{\mu}_i = \tilde{\sigma}_i = 1$ . The architecture description for the neural networks used in this example are provided in Figure 11.
Results are summarized in Figure 10 below. On this problem we find we are able to achieve a reasonable validation accuracy for  $S = 10$ ,  $S = 50$ , and  $S = 100$ ; however, it is challenging to know beforehand what gradient variance will be acceptable for a particular data set. We see that increasing  $S$  increases the total number of function evaluations. Note that the total number of parallel evaluations of the forward model in all cases remains constant.
Figure 10: Validation RMSE versus iteration (L) and number of function evaluations (R).
(a)  $\mathrm{ENC}_{\phi}(x_i)$
(b)  $\mathrm{DK}_{\phi}(t)$
Figure 11: Architecture diagrams for the encoder, deep kernel and drift function used in the Monte-Carlo parameters study.
(c)  $f_{\theta}(z,t)$