SlowGuess/NeurIPS_2023_OCR huggingface dataset:
where, again, each  $\epsilon^{(k)}\sim \mathcal{N}(0,I)$  and each  $t^{(k,1)},t^{(k,2)},\ldots ,t^{(k,S)}\sim \mathcal{U}(t_1^{(j)},t_1^{(j + 1)})$  . In addition, because the integral over time is one-dimensional we used stratified sampling to draw from  $\mathcal{U}(t_1^{(j)},t_1^{(j + 1)})$  in order to further reduce the variance in the integral over time. In this case we often found we could choose  $R\sim 10$  and  $S\sim 10$  . To be clear, (10) is simply a method for variance reduction that we found to be useful; it is not a necessary component for our approach.
3 Limitations & Related Work
Summary of assumptions. In the previous sections we introduced an ELBO which, when maximized, leaves us with a generative model in the form of a nonlinear, latent SDE with time-dependent diffusion and an approximation to the latent state over the time-window of observations in the form of a Gaussian process. To reiterate, we only assume that the approximating posterior, i.e. the distribution over the latent state given a batch of observations, is a Gaussian process; this is an assumption that is commonly made in the context of nonlinear state estimation, for example [23, 24]. When making predictions, we sample from the nonlinear SDE which characterizes the generative model (1).
Stochastic adjoint sensitivities. Li et al. [12] proposed the stochastic adjoint sensitivity method, enabling the inference of latent SDEs using a wide range of approximate posteriors over the latent state. In our work we choose to approximate the posterior over the latent state using a MGP which enables us to eliminate the requirement of solving any differential equations entirely; as we have discussed extensively this choice enables dramatic computational savings. A limitation of our approach as compared to the stochastic adjoint sensitivities method is that our method should only be used to approximate the posterior over the latent state when it is approximately a MGP. Intuitively, this limitation is akin to the limitations of mean-field stochastic variational inference as compared
to stochastic variational inference with an expressive approximate posterior such as normalizing flows [25]. From our practical experience working on a range of test cases, this has not been a limiting factor. It is worth reiterating that this limitation applies only to the approximate posterior over the time window of observations; the predictive posterior can be a complex distribution defined by a nonlinear SDE with a Gaussian initial condition.
In addition, the stochastic adjoint sensitivity method allows for state dependent diffusion processes whereas our approach only allows for a time dependent diffusion process. In cases where a state dependent diffusion process is deemed necessary, our approach could be used to provide a good initial guess for the parameters of the drift function. It remains a topic of future work to determine if this limitation is mitigated by the fact that we are learning latent SDEs rather than SDEs in the original data space. Across the range of test cases we considered, we have not encountered a problem for which the assumption of a time-dependent diffusion matrix was limiting.
Latent neural ODEs. Chen et al. [1], Rubanova et al. [11], and Toth et al. [17] presented latent ordinary differential equations (ODEs) as generative models for high-dimensional temporal data. These approaches have two main limitations: (i) they encode all uncertainty in the ODE's initial condition and (ii) they rely on adjoint sensitivities, necessitating the solution of a sequence of initial value problems during optimization. As was discussed previously, SDEs provide a more natural modeling paradigm for estimating uncertainty, naturally capturing our intuition that uncertainty should accumulate over time [18]. Moreover, to reiterate, our work avoids solving differential equations entirely by relying on unbiased approximations of a one-dimensional integral instead; as we will show, this can result in a dramatic decrease in the number of required function evaluations in training as compared to methods based on adjoints. Moreover, we will show that our approach avoids the numerical instabilities of adjoint methods when they are used to approximate gradients of time averaged quantities over long time intervals for chaotic systems. It is worth mentioning that gradients computed by backpropagation of a forward solver are not consistent with the adjoint ODE in general [26] so we do not consider comparisons to such approaches here.
Weak form methods. Methods for inferring continuous time models of dynamical systems using the weak form of the differential equations were introduced in the context of learning ODEs with linear dependence on the parameters [27, 28]. More recently these methods were adapted for training neural ODEs more quickly than adjoint methods for time-series prediction problems [5]. These methods share some similarities to the present approach in how they achieve a computational speed-up – both methods transform the problem of solving differential equations into a problem of integration. In contrast to the present approach, these methods only allow for one to learn an ODE in the data coordinates (i.e. they do not allow for one to infer an autoencoder and a set of differential equations simultaneously). Moreover, these methods rely on a biased estimate for the weak form residual which will fail when observations become too widely spaced. In contrast, in the present approach, we rely on unbiased approximations to the evidence lower bound. Finally, these methods require the specification of a carefully designed test-space [29] – a consideration not required by our approach.
4 Numerical Studies
In this section we provide a number of numerical studies to demonstrate the utility of our approach. In the first study, we show that our approach can be used to train neural SDEs using far fewer evaluations of the model than adjoint methods. In the second study, we consider the problem of parameter tuning for a chaotic system over long time intervals. We show that our approach does not suffer from the numerical instabilities which are known to cause issues with adjoint methods on problems such as these. Finally we close this section with two practical test cases: the first demonstrating competitive performance on a motion capture benchmark and the second showing how our approach can be applied to learn neural SDEs from video. An additional numerical study exploring the effect of the nested Monte Carlo parameter,  $S$ , is provided in Appendix H. Details on computing resources are provided in Appendix F. All code required to reproduce results and figures is provided at github.com/coursekevin/arlatentsde.
Figure 2: Lotka-Volterra benchmarking result. In the left figure we see our method (ARCTA) requires more than one order of magnitude fewer evaluations of the model (NFE) than the standard neural ODE (NODE) to achieve a similar validation accuracy. In the right figure we have plotted a probabilistic prediction on the test set along with three samples from the predictive distribution.
4.1 Orders of magnitude fewer function evaluations in training
In this numerical study we consider the task of building a predictive model from noisy observations of a predator-prey system. We simulated the Lotka-Volterra equations for 50 seconds collecting data at a frequency of  $10\mathrm{Hz}$ . Observations were corrupted by Gaussian noise with a standard deviation of 0.01. Validation data was collected over the time interval [50, 65] seconds. We then attempt to build a predictive model from the data using a neural ODE (NODE) and our method, amortized reparametrization for continuous time auto-encoding (ARCTA), with the same model for the ODE and drift function respectively. To make comparisons with the NODE fair, we set the decoder to be the identity function. We assume the diffusion matrix is constant and place a log-normal prior on its diagonal elements. We approximate the posterior over these elements using a log-normal variational posterior. Details on the architecture and hyperparameters are provided in Appendix G.1. For this experiment, as well as subsequent experiments, we made use of the Adam optimizer [30].
We considered three different tolerances on the NODE adaptive stepping scheme. We trained our model as well as the NODEs using 10 different random seeds while recording the validation RMSE and the number of evaluations of the model. Looking to Figure 2, we see that our approach required more than an order of magnitude fewer evaluations of the model to achieve a similar RMSE on the validation set. This remains true even when the tolerance of the ODE solver is reduced such that the validation RMSE is substantially higher than our approach.
4.2 Numerical instabilities of adjoints
It is well-known that adjoint based methods produce prohibitively large gradients for long time averaged quantities of chaotic systems [10] and accordingly methods, such as least squares shadowing [31], have been introduced to address such concerns. In this section we reproduce this phenomena on a simple parameter tuning problem and show that our approach does not suffer these same issues.
Given the parametric form of the chaotic Lorenz equations,
$$
\dot {x} = \sigma (y - x) \tag {11}
$$
$$
\dot {y} = x (\rho - z) - y \tag {12}
$$
$$
\dot {z} = x y - \beta z \tag {13}
$$
along with an initial guess for the parameters,  $\sigma_0$ ,  $\rho_0$ , and  $\beta_0$ , our goal is to tune the value of parameters such that they align with the observed data.
For this experiment we collect data at a frequency of  $200\mathrm{Hz}$  and corrupt observations by Gaussian noise with a covariance of 1. We generate five independent datasets over the time intervals [0, 1], [0, 10], [0, 50], and [0, 100]. For each dataset we generated an initial guess for the parameters by sampling from a Gaussian whose mean is the true value of the parameters and standard deviation is  $20\%$  of the mean. For the adjoint methods we report the  $\ell_2$ -norm of the gradient with respect to the parameters at the initial guess. For our method (ARCTA) we optimize for 2000 iterations (which tended to be enough iterations to successfully converge to a reasonable solution) and report the average gradient across all iterations. Details on hyperparameters and our architecture design are provided in Appendix G.2. Results are summarized in Figure 3. While adjoints expectedly provide prohibitively large gradients as the length of the time series is increased, our approach remains numerically stable.
Figure 3: Stability of gradients in chaotic systems. The log-scale on vertical axis shows our approach remains stable for longer time series, while adjoint-based gradients become unusable at 50 and 100 seconds.
4.3 Motion capture benchmark
Method	Test RMSE
DTSBN-S [32]	5.90 ± 0.002†
npODE [33]	4.79†
NeuralODE [1]	4.74 ± 0.093†
ODE²VAE [34]	3.17 ± 0.221†
ODE²VAE-KL [34]	2.84 ± 0.343†
Latent ODE [11]	2.45 ± 0.057*
Latent SDE [12]	2.01 ± 0.050*
ARCTA (ours)	2.76 ± 0.168
Figure 4: MOCAP benchmarking results,  $\dagger$  from [34] and  $*$  from [12]. Our score is computed by training 10 models with different seeds and averaging on the test set. Looking to the table, we see that our method performs similarly to other state-of-the-art methods. The plot shows the predictive posterior on the test set for some select outputs. Other benchmark results were compiled in [34, 12]. RMSE was computed from MSE by taking the square root of the mean and transforming the error via a first-order Taylor-series approximation.
In this experiment we consider the motion capture dataset from [32]. The dataset consists of 16 training, 3 validation, and 4 independent test sequences of a subject walking. Each sequence consists of 300 time-series observations with a dimension of 50. We made use of the preprocessed data from [34]. Like previous approaches tackling this dataset, we chose a latent dimension of 6. We assume a Gaussian observation likelihood. We place a log-normal prior on the diagonal elements of the diffusion matrix and the noise on the observations. We approximate the posterior of the diffusion matrix and observation noise covariance using a log-normal approximate posterior. Details on our architecture design and hyperparameter selection are provided in Appendix G.3.
For our approach, we train 10 models and report their average performance on the test set due to the extremely limited number (4) of independent test sequences. Looking to Figure 4, we see that our approach provided competitive performance on this challenging dataset. This result, in combination with those presented previously demonstrating we require fewer function evaluations for similar
forecasting accuracy and improved gradient stability for chaotic systems, make clear the utility of the present work. It is possible to achieve state-of-the-art performance at a significantly reduced computational cost as compared to adjoint based methods.
4.4 Neural SDEs from video
Figure 5: Neural SDEs from video. Here we used five frames to estimate the initial state and then forecast in the latent space for 30 seconds. The bottom plot shows the latent SDE. The top row shows 10 samples from the predictive posterior overlaid on the data.
In this experiment we attempt to learn a latent SDE from video. We generated  $32 \times 32$  black and white frames of a nonlinear pendulum as it evolves for 30 seconds collecting data at a frequency of  $15\mathrm{Hz}$ . We transform the 1024 dimensional state down to two dimensions using a convolutional architecture. Details on the hyperparameters and architecture are provided in Appendix G.4. This problem is similar to the problem considered in [3] except the dynamical system we consider is nonlinear. In this prior work, the authors were forced to regularize the latent space so that one set of coordinates resembles a generalized velocity. In the present work, no such regularization is required.
We assume a Bernoulli likelihood on the pixels. Like in previous numerical studies we place a log-normal prior on the diagonals of the diffusion term and approximate the posterior using a log-normal variational distribution. After training we generate a forecast that is visualized in Figure 5. We see that we were successfully able to build a generative model for this simple video. This result demonstrates the broad applicability of the present approach to standard generative modeling tasks.
Here we have presented a method for constructing unbiased approximations to gradients of the evidence lower bound used to train latent stochastic differential equations with a time and memory cost that scales independently with the amount of data, the length of the time-series, and the stiffness of the model for the latent differential equations. We achieve this result by trading off the numerical precision of adaptive differential equation solvers with Monte-Carlo approximations to expectations using a novel amortization strategy and a recently derived change of variables for expectations under Markov Gaussian processes [13].
We have demonstrated the efficacy of our approach in learning latent SDEs across a range of test problems. In particular we showed that our approach can reduce the number of function evaluations as compared to adjoint methods by more than one order of magnitude in training while avoiding the numerical instabilities of adjoint methods for long time series generated from chaotic systems. In addition, we showed that our approach can be used for generative modeling of a simple video.
In the immediate future, there is significant room for future work in applying variance reduction schemes to the expectation over time to further reduce the total number of required function evaluations. There are also opportunities to explore the utility of the proposed approach for generative modeling on more realistic problems. Finally, there are opportunities to apply our work in the context of implicit densities [35].
This research is funded by a grant from NSERC.
A Expectations under linear SDEs
In this section we re-derive a result from [13] regarding how to rewrite expectations under linear stochastic differential equations. As we will explain in greater detail in Appendix B, it is this result that allowed [13] to rewrite the ELBO entirely in terms of quantities that do not require differential equation solvers.
Theorem 1. Consider the density,  $q(z(t))$ , of the solution to a linear SDE  $dz = (-A(t)z + b(t))dt + L(t)d\beta$  with initial condition  $z_0 \sim \mathcal{N}(m_0, S_0)$ , where  $A: \mathbb{R} \to \mathbb{R}^{d \times d}$  is symmetric,  $b: \mathbb{R} \to \mathbb{R}^d$ ,  $L: \mathbb{R} \to \mathbb{R}^{d \times d}$ , and  $\beta$  indicates Brownian motion with diffusion matrix  $\Sigma$ . Then the expected value of a bounded functional,  $f$ , satisfies
$$
\mathbb {E} _ {z (t) \sim q (z (t))} [ f (A (t), b (t), z (t)) ] = \mathbb {E} _ {z (t) \sim \mathcal {N} (m (t), S (t))} [ f (B (t), \dot {m} (t) + B (t) m (t), z (t)) ], \tag {14}
$$
where  $B(t) = \text{vec}^{-1}((S(t) \oplus S(t))^{-1}\text{vec}(L(t)\Sigma L(t)^T - \dot{S}(t)))$  and  $m(t)$  and  $S(t)$  indicate the mean and covariance, respectively, of the SDE solution at time  $t$ .
Proof. The solution of a linear SDE defines a Markov Gaussian process with marginal statistics,  $q(z(t)) \sim \mathcal{N}(m(t), S(t))$ , given by the solution to the ODEs,
$$
\dot {m} (t) = (- A (t) m (t) + b (t)), \tag {15}
$$
$$
\dot {S} (t) = - A (t) S (t) - S (t) A (t) ^ {T} + L (t) \Sigma L (t) ^ {T}, \tag {16}
$$
with initial condition  $m(0) = m_0$ ,  $S(0) = S_0$  [19]. Noticing that equation (16) defines a set of matrix Lyapunov equations in terms of  $A(t)$  allows us to express  $A(t)$  as a function of  $S(t)$  as follows,
$$
A (t) = \operatorname {v e c} ^ {- 1} \left(\left(S (t) \oplus S (t)\right) ^ {- 1} \operatorname {v e c} \left(L (t) \Sigma L (t) ^ {T} - \dot {S} (t)\right)\right), \tag {17}
$$
where  $\oplus$  is called the Kronecker sum and is defined as  $S\oplus S = I\otimes S + S\otimes I$  and  $\otimes$  indicates the standard Kronecker product. Letting  $B(t) = A(t)$  be the expression for  $A(t)$  written in terms of  $S(t)$ , we can rearrange Equation (15) to solve for  $b(t)$  as,
$$
b (t) = \dot {m} (t) + B (t) m (t). \tag {18}
$$
Substituting the expressions for  $A(t)$  and  $b(t)$  into Equation (14) yields the desired result.
Remark In the case that  $S(t)$  and  $L(t)\Sigma L(t)^T$  are diagonal as we assume throughout this work, Equation 17 simplifies to,
$$
A (t) = \frac {1}{2} S (t) ^ {- 1} \left(L (t) \Sigma L (t) ^ {T} - \dot {S} (t)\right). \tag {19}
$$
B Evidence lower bound derivation
Given a dataset of observations,  $\mathcal{D} = \{(x_i,t_i)\}_{i = 1}^N$ , a generative likelihood,  $p_{\theta}(x\mid z(t_i))$  a prior SDE,  $dz = f_{\theta}(z(t),t)dt + L_{\theta}(t)d\beta$ , and an approximation to the posterior of the latent state,  $dz = (-A_{\phi}(t)z(t) + b_{\phi}(t))dt + L_{\theta}(t)d\beta$ , it is possible to derive the ELBO [20, 12],
$$
\begin{array}{l} \mathcal {L} (\theta , \phi) = \sum_ {i = 1} ^ {N} \mathbb {E} _ {z \left(t _ {i}\right) \sim q _ {\phi} (z (t))} \left[ \log p _ {\theta} \left(x _ {i} \mid z \left(t _ {i}\right)\right) \right] \tag {20} \\ - \frac {1}{2} \int_ {0} ^ {T} \mathbb {E} _ {z (t) \sim q _ {\phi} (z (t))} \left[ | | - A _ {\phi} (t) z (t) + b (t) - f _ {\theta} (z (t), t) | | _ {C _ {\theta} (t)} ^ {2} \right] d t, \\ \end{array}
$$
where  $C_{\theta}(t) = (L_{\theta}(t)\Sigma L_{\theta}^{T}(t))^{-1}$ . The second term in (20) contains an integral over time of terms of the form which matches that of the expectation in (14). From this observation, Course and Nair [13] applied the identity in (14) to derive the reparametrized ELBO,
$$
\begin{array}{l} \mathcal {L} (\theta , \phi) = \sum_ {i = 1} ^ {N} \mathbb {E} _ {z (t _ {i}) \sim q _ {\phi} (z (t))} [ \log p _ {\theta} (x _ {i} \mid z (t _ {i})) ] \tag {21} \\ - \frac {1}{2} \int_ {0} ^ {T} \mathbb {E} _ {z (t) \sim q _ {\phi} (z (t))} \left[ | | r _ {\theta , \phi} (z (t), t) | | _ {C _ {\theta} (t)} ^ {2} \right] d t, \\ \end{array}
$$
$$
C _ {\theta} (t) = \left(L _ {\theta} (t) \Sigma L _ {\theta} (t) ^ {T}\right) ^ {- 1}
$$
$$
r _ {\theta , \phi} (z (t), t) = B (t) \left(m _ {\phi} (t) - z (t)\right) + \dot {m} _ {\phi} (t) - f _ {\theta} (z (t), t) \tag {22}
$$
$$
B (t) = \operatorname {v e c} ^ {- 1} \left(\left(S _ {\phi} (t) \oplus S _ {\phi} (t)\right) ^ {- 1} \operatorname {v e c} \left(L _ {\theta} (t) \Sigma L _ {\theta} (t) ^ {T} - \dot {S} _ {\phi} (t)\right)\right).
$$
To reiterate what was mentioned in the main body of the present work, the advantage of this reparametrized ELBO is that all expectations are now taken with respect to normal distributions – this has effectively eliminated the requirement of a differential equation solver. Unfortunately this approach requires storing the entire estimate for the latent state – making it scale poorly for long time series with complex dynamics. After amortizing as is suggested in the main body of the present work we arrive at the final ELBO,
$$
\begin{array}{l} \mathcal {L} (\theta , \phi) = \sum_ {j = 1} ^ {N / M} \sum_ {i = 1} ^ {M} \mathbb {E} _ {q _ {\phi} \left(z \left(t _ {i} ^ {(j)}\right) \mid x _ {1} ^ {(j)}, \dots , x _ {M} ^ {(j)}\right)} \left[ \log p _ {\theta} \left(x _ {i} ^ {(j)} \mid z \left(t _ {i} ^ {(j)}\right)\right) \right] \tag {23} \\ - \frac {1}{2} \int_ {t _ {1} ^ {(j)}} ^ {t _ {1} ^ {(j + 1)}} \mathbb {E} _ {q _ {\phi} (z (t) | x _ {1} ^ {(j)}, \dots , x _ {M} ^ {(j)})} \left[ | | r _ {\theta , \phi} (z, t) | | _ {C _ {\theta} (t)} ^ {2} \right] d t. \\ \end{array}
$$
Again, as is discussed at length in the main body of the present work, the advantage of such an amortization strategy is that it is possible to construct an unbiased approximation to the gradients of this ELBO that scales independently with the amount of the data, the length of the time series, and the stiffness of the underlying differential equations.
C Evidence lower bound with priors on generative parameters
In many circumstances, it will be advantageous to place priors on a subset of the parameters,  $\theta$ . Using the tools of stochastic variational inference, we can infer the posterior on these variables with only a marginal increase to the overall computational cost and with no impact to the asymptotic computational complexity.
For example, in all numerical studies we placed a log-normal prior on the diagonal of the diffusion process term  $C_{\theta}^{-1}(t)$  and approximated the posterior using a log-normal approximate posterior, see Appendix E for details. While not necessary, we found that doing so helped to stabilize training in the examples we considered. The particular choice of prior-posterior pair was made to ensure that the Kullback-Leibler (KL) divergence between the approximate posterior and the prior could be written in closed form; see Appendix E. Given a particular choice of prior,  $p(\theta)$ , along with an approximate posterior,  $q_{\phi}(\theta)$ , we can amend the previously derived ELBO as,
$$
\begin{array}{l} \mathcal {L} (\phi) = \sum_ {j = 1} ^ {N / M} \sum_ {i = 1} ^ {M} \mathbb {E} _ {q _ {\phi} (\theta) q _ {\phi} (z \left(t _ {i} ^ {(j)}\right) | x _ {1} ^ {(j)}, \dots , x _ {M} ^ {(j)})} \left[ \log p _ {\theta} \left(x _ {i} ^ {(j)} \mid z \left(t _ {i} ^ {(j)}\right)\right) \right] \tag {24} \\ - \frac {1}{2} \int_ {t _ {1} ^ {(j)}} ^ {t _ {1} ^ {(j + 1)}} \mathbb {E} _ {q _ {\phi} (\theta) q _ {\phi} (z (t) | x _ {1} ^ {(j)}, \dots , x _ {M} ^ {(j)})} \left[ | | r _ {\theta , \phi} (z, t) | | _ {C _ {\theta} (t)} ^ {2} \right] d t - D _ {K L} (q _ {\phi} (\theta) | | p (\theta)), \\ \end{array}
$$
where  $D_{KL}(q\parallel p)$  indicates the KL divergence between  $q$  and  $p$
Consider a reparametrization to the latent state as  $z(t) = T(t,\epsilon ,\phi)$  where  $\epsilon \sim p(\epsilon)\Rightarrow z(t)\sim q_{\phi}(z(t)\mid x_1^{(j)},x_2^{(j)},\ldots ,x_M^{(j)})$ . Also consider a reparametrization to the approximate posterior for the