Amortized Reparametrization: Efficient and Scalable Variational Inference for Latent SDEs
Authors: Kevin Course
Prasanth B. Nair
Venue: NeurIPS 2023 poster (2023)
Keywords: variational inference
differential equations
dynamical systems
neural ordinary differential equations
latent stochastic differential equations
TL;DR: We present a new method for training latent stochastic differential equations far more quickly than methods based on adjoints using a new amortization strategy.

Abstract: We consider the problem of inferring latent stochastic differential equations (SDEs) with a time and memory cost that scales independently with the amount of data, the total length of the time series, and the stiffness of the approximate differential equations. This is in stark contrast to typical methods for inferring latent differential equations which, despite their constant memory cost, have a time complexity that is heavily dependent on the stiffness of the approximate differential equation. We achieve this computational advancement by removing the need to solve differential equations when approximating gradients using a novel amortization strategy coupled with a recently derived reparametrization of expectations under linear SDEs. We show that, in practice, this allows us to achieve similar performance to methods based on adjoint sensitivities with more than an order of magnitude fewer evaluations of the model in training.

Amortized Reparametrization: Efficient and Scalable Variational Inference for Latent SDEs
Kevin Course University of Toronto
Prasanth B. Nair University of Toronto prasanth wwwr@utorontoocoommmmm
kevin. cour course@mail. www.ootooooooom ca
Abstract
We consider the problem inferring latent stochastic differential equations (SDEs) with time and memory cost that scales independently with the amount of data, the total length the time series, and the stiffness of the approximate differential equations. This is in stark contrast typical methods for inferring latent differen- tial equations which, despite their constant memory cost, have a time complexity that is heavily dependent the stiffness of the approximate differential equation. We achieve this computational advancement by removing the need to solve differ- ential equations when approximating gradients using novel amortization strategy coupled with a recently derived reparametrization of expectations under linear SDEs. We show that, in practice, this allows to achieve similar performance methods based on adjoint sensitivities with more than an order of magnitude fewer evaluations of the model in training.
Introduction
Recent years have seen the rise of continuous time models for dynamical system modeling M. As compared to traditional autoregressive style models [2], continuous time models are useful because they can deal with non-evenly spaced observations, they enable multileee/hhieracchical and adaptive prediction schemes, and because physics is (mostly) done in continuous time. For example, recent developments inferring continuous time differential equations from data has been met by flurry of work in endowing models with physics informed priors 3 5..
Despite their advantages, continuous time models remain significantly lenging to train than their autoregressive counterparts due to their reliance on adjoint methods for estimating gradients. Adjoint methods introduce significant computational burden in training because they require solving pair of initial value problems to estimate gradients. Solving such initial value problems as part of an iterative optimization procedure is computationally demanding for the following reasons:
computationally chal-
Gradient based updates to models differential equations can cause them to become extremely stiff. This will have the effect of causing the cost per iteration to explode mid-optimization.
(ii) With the exception of parareal methods (6), differential equation solvers fundamentally iterative sequential methods. This makes them poorly suited to being parallelized on modern parallel computing hardware.
In accordance with such challenges number of methods have been introduced to speed up training of continuous time models including regularizing dynamics 7 g and replacing ordinary differential equation (ODE) solvers with integration methods where possible (9. Despite the computational advancements brought about by such innovations, continuous time models remain expensive to train in comparison to discrete time models.
37th Conference on Neural Information Processing Systems (NcurIPS 2023).

---

In addition to these computational challenges, it is well-known that adjoint methods suffer from stability issues when approximating gradients of time averaged quantities over long time intervals for chaotic systems 10
In the current work, we present a memory and time efficient method for inferring nonlinear, latent stochastic differential equations (SDEs) from nigh-diimnssioal time-series datasets. In contrast to standard approaches for inferring latent differential equations that rely on adjoint sensitivities II I 12), our approach removes the requirement of solving differential equations entirely. We accomplish this advancement by coupling a novel amortization strategy with recently derived reparametrization for expectations under Markov Gaussian processes [[3]. We show that our approach can be used to approximate gradients of the evidence lower bound used train latent SDEs with a time and memory cost that is independent of the amount of data, the length of the time series, and the stiffness of the approximate differential equations. The asymptotic complexity for our approach is compared to well-known methods from the literature in [ We note that our method has a constant cost that is chosen by the user. Moreover, we will show that our method embarrassingly parallel (i.e. all evaluations of the model can be performed in parallel over each iteration) whereas, we reiterate, differential equation solvers are iterative sequential methods.
Method
Time
Memory
Deterministic adjoints (Neural ODE) [II
O(J)
O(1)
Stochastic adjoints [12]
O log J) O(J)
0(1)
Backprop through solver [4]
O(J)
Amortized reparametrization (ours)
O
O(R)
Table 1: Asymptotic complexity comparison for approximating gradients. Units are given by the number of evaluations of the differential equation (gradient field for ODEs or drift and diffusion function for SDEs). Here J is the number of sequential evaluations and is the number of parallel evaluations. // is adaptively chosen by the differential equation solver and is function of the stiffness of the differential equation meaning it can change (and possibly explode) while optimizing. In contrast, It is fixed constant used to control the variance of gradient approximations. While we could choose R = and would still arrive with unbiased approximations for gradients, in practice we found choosing R x 10 worked well for the problems we considered.
The applications of our method span various standard generative modeling tasks, such as auto- encoding, denoising. inpainting, and super-resolution [15). particularly tailored for high- high-dimensional time-series. Crucially, the computational efficiency of our approach not only enables the allocation of more computational resources towards hyperparameter tuning but also democratizes access to state-of-the-ar art methods by making them feasible to train on lower performance hardware.
In the next section we provide a description of the theory underpinning our work with the main result of a stochastic, unbiased estimate for gradients appearing in Lemma [ In Section 4 we provide number of numerical studies including learning latent neural SDEs from video and performance benchmarking on a motion capture dataset. Notably we show that we are able to achieve comparable performance to methods based on adjoints with more than one order of magnitude fewer evaluations of the model in training (Section 4. We also demonstrate that our approach does not suffer from the numerical instabilities that plague adjoint methods for long time-series with chaotic systems (Section 4.2) Finally, we close with discussion of the limitations of our approach as well as some suggestions for future work. All code is available at gigthbb.commouuuseeeeiiiaaaendd
2 Method
2.1 Problem description
Consider time-series dataset D - (xx,, (i where i, E R is the time stamp associated with the observation, It E D. For example v; may be an indirect observation of some underlying dynamical system, a video-frame, or a snapshot of a spatio-temporal field. We assume that each observation was generated via the following process: first latent trajectory, 2t)), is sampled from general,

---

nonlinear SDE with dime-dependent diffusion,
dz = z))t + To (t)d3,
with initial condition (=() and then each x; is generated from the conditional distribution po(x =(4)) where =(4;) is the realization of the latent trajectory at time stamp 4. Here fo : IR x R2 R is called the drift function, La IR + Rdxd is called the dispersion matrix, and [ denotes Brownian motion with diffusion matrix E. We will model both the drift function of the SDE and the conditional distribution using neural networks in this work. This combination of SDE and conditional distribution define the generative model that we wish to infer.
Like in the standard generative modeling setting, the "true" parameters defining the latent SDE and the conditional likelihood, well as the specific realization of the latent state, remain unknown. In addition, the posterior over the latent state at a point in time, (=(((() DD), is intractable. Our objectives are two-fold, we wish to: (i) infer likely setting for the parameters, 0, that is well- aligned with the observed data and (ii) infer a parametric model for the posterior over the latent state a particular point in time as a function of small window of the observations, for example qo (=())) 2i, 1i......, Ti+ M) R (=()t D) for & E [t,, M and 1 S. M << N. This latent variable model depicted in Figure In this work will tackle this problem statement using the machinery of stochastic variational inference [16].
N
Figure 1:
The solid lines indicate the data generating process that depends on the parameters, 6. The dashed lines indicate the approximate variational posterior that depends on the parameters, ạ).
Before proceeding with an explanation of our approach, it is worthwhile taking a moment to consider why assuming the latent state is realization of a SDE, as opposed an ordinary differential equation (ODE) with random initial conditions аيτي [7], merits the additional mathematical machinery. First, there is long history in using SDEs in the physical sciences model systems (even deterministic systems) that lack predictability due to super-sensitiviy to parameters and or initial conditions (18). One reason for the preference towards this modeling paradigm is that deterministic ODEs with random initial conditions only permit uncertainty to enter in their initial condition. This means that we are assuming that all randomness which may affect the latent trajectory at all future points in time are precisely encoded in the uncertainty in the initial condition. In contrast, SDEs make the less restrictive assumption that uncertainty accumulates over time.
2.2 Evidence lower bound
As is the case with standard latent variable models, training proceeds by specifying approximate posteriors over latent variables (in this case the latent state, =(t)) and then minimizing the Kullback- Leibler (KL) divergence between the approximate and true posterior. An added complication in our specific case comes from the fact that the latent state is defined by stochastic process rather than random-vector as is more typical. While it is possible to choose more general approximate posteriors [12. in this work we make the choice to approximate the posterior over the latent state as a Markov Gaussian Process (MGP) defined by the linear SDE,
dz = (-=(AAph) (-Ax(1)= (1) + Lo)d3,
with Gaussian initial condition, =(tg) Nm Sa) where my Rd and Sy E dRdd x is symmetric positive definite. Here Ag : - Rddd and by IR + Rd are symmetric matrix and vector-valued functions of time, respectively. Before proceeding, it is worth emphasizing that this choice of approximate posterior is not as restrictive as may appear since it only used to approximate the posterior (i.c. the latent state given the observations p(z(t) D) for t E [t,, tx]) When forecasting
3

---

beyond the data window, we will use the nonlinear SDE in ①.. In this way, the approximate posterior primarily serves as a useful intermediary in training the generative model rather than the end aud-product.
For our purposes, MGPs are useful because the marginal distribution ga (=(()) = (mss (1), (s()) (i.e. the distribution of the latent state time () is Gaussian whose mean and covariance given
by the solution to the following set of ODEs 19).
(1) = - AAs()mm + bo((),
(1) = - -A)(() - (())) + To(())To(()"
where m(0) mg and 8(0) So.
Using the reparametrization trick from
13 the evidence lower bound (ELBO) can be written as,
N
[(0,0) 6)
Σ
E (e(t))))())))) Pog =(t;))]
1=1
(4)
1
(2(0))0(((10) [Inr,o(=(0))))))ecc) dl,
2
where
= (B))())))))B((((g())))))))))))) =(t)) + m) - fo()
(5)
Co(1) (L)()) B(I) vec (())(1) @
So())) indi- cates the Kronecker sum, vec Raxd + Rd maps matrix into a vector by stacking columns, and B AB e Rdxd
vec-i Rdd² + Rdxd converts vector into a matrix such that vvec See Appendices A and B for a detailed derivation. We note that computing the residual, re,ó, scales linearly in the dimension of the state long as (t) and (t) diagonal. To ensure our approach remains scalable, we will make this assumption throughout the remainder this work. In this case,
have
1 2
B(t)
St)()...
(6)
Often we may wish to place additional priors onto subset of the generative model parameters, 8, and infer their posterior using stochastic variational inference as well. In this case we add the KI-divergence between the approximate posterior and prior onto the ELBO, K L(ga(0) || p(0)) and rewrite all expectations with respect to go (z())) and (f); details are provided Appendix
Remark 1: Despite having defined the approximate posterior in terms of an SDE with parameters Ag and bg. the BLBO only depends on the mean and covariance of the process at particular point in time, my and Sg. For this reason we can parametrize mi and Sy directly while implicitly optimizing with respect to Ag and by. In addition, we can efficiently compute my and Šo using automatic differentiation.
Remark 2: Despite the fact that the prior and approximate posterior are SDEs, all expectations in the ELBO are taken with respect to normal distributions. Moreover, in contrast the approach in [20) 21 there are no differential equality constraints instead we have been left with an integral over the window of observations.
Taken together, these observations allow us to infer the parameters of the generative model (a nonlinear, latent SDE with additive diffusion (T). without the use of a forward solver.
2.3 Amortization strategy
The implicit assumption in the ELBO in is that the state mean and covariance will be approximated over the entire window of observations. This can pose serious computational bottleneck with long or complicated time-series. In this section we propose a novel amortization strategy that will allow us effectively eliminate this cost by requiring that we only approximate the posterior over short. partitions of total the data time-window at once.

---

Rather than attempting to compute and store the posterior over the entire latent trajectory. we will instead construct an approximation to the posterior over small window of observations function of those observations. Consider a reindexing of the dataset by splitting it into N/M non-overlapping
partitions where I M << N,
original indexing:
[1, 12
A.M (2)
AN
reindexed dataset:
( 1/2
(NA0) M
In the case that N is not evenly divisible by M allow the final split to contain less elements. We approximate the latent state over each partition using only the M observations in each partition,
(s(=(() =f 222 "."",..."""" ) x p(=(() ')) for L E ((",,(it
())
This can be interpreted as N/M
probabilistic encoder over the time interval of the partition of observations. Letting E N/aS+1 = LM the ELBO can be compactly rewritten as, L(0, Φ) En 2 (0, Φ), where
M
(6) (0,
E 20(2(0
1= ()) [log po( ( )((
E [o(=())2 ......))))))))))) (1, dt.
An additional advantage of this amortization strategy is that it allows approach to scale to multiple trajectories without an increase to the overall computational cost. If there are multiple trajectories, we can reindex each trajectory independently and subsequently sum all sub loss functions.
To reiterate, the probabilistic encoder a function which takes M observations from particular partition along with a time stamp, f, and outputs mean vector and covariance matrix estimate for the latent state at that particular time. In principle, any function which can transform a batch of snapshots and a time stamp into mean and covariance could be used as an encoder in our work. In our implementation, we use deep neural networks to encode each x!")) using i E where I I [a I22 .)) ..... M. contains some temporal neighbours of x; into single latent vector. This approach yields a set of latent vectors associated with each observation in the partition h, for = I.2.....M.. We then interpolate between each latent vector using deep kernel based architecture to construct the posterior approximation for any time stamp in the partition; see Appendix D for details. We emphasize this is choice of encoding architecture that we found convenient, is straightforward to incorporate an autoregressive or transformer based encoder in our methodology 22
(
An important consideration is the selection of the partition parameter, M. In practice, M should be large enough so that the assumption a linear SDE for the approximate posterior is appropriate (i.e. we have enough observations in a partition so that the assumption of Gaussian process over the latent state is reasonable). For example, as we will see in an upcoming numerical study, the context of inferring latent SDEs from videos, we will need to choose M to be large enough so that we can reasonably infer both the position and velocity of the object in the video.
2.4 Reparametrization trick
While the previous sections have demonstrated how to eliminate the need for differential equa- tion solver by replacing the initial value problem with an integral, in this section we show how the reparametrization trick can be combined with the previously described amortization strategy to construct unbiased gradient approximations for the ELBO with time and memory cost that scales independently with the amount of data, the length of the time series, and the stiffness of the approximation to the differential equations. Consider a reparametrization to the latent state of the
form z(!) = T(t, Φ) where ~ p(e) so that x1) ~ (z(1) x 202
()) We can rewrite x
the second term in the evidence lower bound as,
(+1)
(5+3)
E (gg()(0)
[[re,, (zr), dt =
(2p Ept)) (To o), (). dt
= (6(+1) Eptestn (Tt,,
(8)
5

---

where p(() is uniform distribution, U _1+t2) and p(e) ~ (O, is Gaussian. With this rearrangement, we can derive the main result this work.
Lemma 1. An unbiased approximation of the gradient of the evidence lower bound, denoted as Li, Φ), with an (R)) time and memory cost can be formulated as follows:
M
R
V V0,0 ②)
V es log po(
T ,0,
R
(9)
(e
N
( Vo (T(t(h) (t) 0) t (
2R
where each t()) ~ u(t?), ,01 1 and each c ~ I).
The proof follows by applying the standard reparametrization trick amortized objective in
15 estimating gradients of the
In practice found choosing R 100 worked well for the problems considered. Note that in terms of elapsed time, 100 evaluations of this objective, which can be computed in parallel, is far cheaper than 100 evaluations of the SDE forward model evaluated part of an iterative sequential SDE solver. Moreover we found that adaptive stepping schemes required far more evaluations of the SDE forward model than our stochastic approach (see Section
Remark 1:
In the case that evaluations of the SDE drift term were relatively cheap compared to decoder evaluations (for example in the case the dimension of the latent state is much smaller than the dimension of the data), we found it useful to increase the number of samples used to approximate the integral over time without increasing the number of samples from the variational posterior. To do so, we made use of nested Monte Carlo scheme to approximate the second term in the ELBO,
Remark 2:
(o()+1) -2(") 10++)+(()) Ep(chpp) (T ().
((+1)
((()
R
9
ro. (T(t(ht)) (,) ,),,
(u (ua.iny (10)
RS
where, again, each (k) (O, I) and each ((k.1) (k.22) (k,S) u(te() t (((tt)). addi- tion, because the integral over time is one-dimensional we used stratified sampling to draw from often (((()),(++) in order to further reduce the variance in the integral over time. In this
found we could choose R ~ 10 and S ~ 10. To be clear, (IO) is simply a method for variance reduction that we found be useful; is not necessary component for our approach.
3 Limitations & Related Work
Summary of assumptions. In the previous sections we introduced an ELBO which, when maxi- mized, leaves us with generative model in the form of 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 batch of observations, is a Gaussian process; this is an assumption that is commonly made 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. et al. proposed the stochastic adjoint sensitivity method, enabling the inference of latent SDEs using wide range of approximate posteriors over the latent state. 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 approach as compared to the stochastic adjoint sensitivities method is that our method should only be used approximate the posterior over the latent state when it approximately a MGP. Intuitively, this limitation is akin to the limitations of mean ficlddfield stochastic variational inference compared
our

---

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 Gausian 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 deemed necessary, our approach could be used to provide a good initial guess for the parameters of the drift function. It remains 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 problem for which the assumption of a time-dependent diffusion matrix was limiting.
Latent neural ODEs. Chen et al. M. Rubanova et al. [II], and Toth et al. [1 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 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 [8]. Moreover, to reiterate, our work avoids solving differential equations entirely by relying on unbiased approximations
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 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 difficrnniill equations were introduced in the context of learning ODEs with linear dependence on the parameters (27 22) 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 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 the present approach, these methods only allow for one 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 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 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 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 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 Appendix [F] All code required to reproduce results and figures provided at guthub..com/comsseeevmm/aaannnn
7

---

0.5
ARCTA
NODE, tol=10
logio Validation RMSE
NODE, tol=10 6
-1.0
-1.5
-2.0
Train
Test.
35
40
45
50
65
60
65
1og
Time
Figure Lotka- Voltera benchmarking result. 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 similar validation accuracy. In the right figure we have plotted 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 predictive model from noisy observations of predator-pre prey system. We simulated the Lotka- Volterra equations for 50 seconds collecting data frequency of 10Hz. Observations were corrupted by Gaussian noise with a standard deviation of 0.01. Validation data was collected over the time inverval [50, 65) seconds. We then attempt to build a predictive model from the data using 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 log-normal prior on its diagonal elements. We approximate the posterior over these elements using log-normal variational posterior. Details 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 well the NODEs using 10 different random seeds while recording the validation RMSE and the number of evaluations of the model. Looking to Figure we see that our approach required more than an order of magnitude fewer evaluations of the model to achieve similar RMSE on the validation set. This remains true even when the tolerance of the ODE solver is reduced such that the validation RMSE 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 [00 and accordingly methods, such as least squares shadow- ing [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,
i = a x.) z(p 2) y غ = xy Bz
(11)
(12)
(13)
along with an initial guess for the parameters, do. Po, and 3o our goal to tune the value of parameters such that they align with the observed data.

---

For this experiment we collect data at a fre- quency of 200Hz and corrupt observations by Gaussian noise with a covariance of 1. We gen- erate 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 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 fn-morm 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 reasonable solution) and report the average gradient across all iterations. Details on hyper- parameters and our architecture design pro- vided in Appendix G.2 Results are summarized Figure 3 While adjoints expectedly provide prohibitively large gradients as the length of the time series increased, our approach remains numerically stable.
1045
ARCTA
Adjoints
10a5
1030
log10 V
1626
1020
1045
168
16
1.0
10.0 50.0
100.0
T(*)
Figure 3: Stability of gradients in chaotic systems. The log-scale on vertical axis shows approach remains stable for longer time series, while adjoint- based gradients become unusable 50 and 100 seconds.
4.3 Motion capture benchmark
Method
Test RMSE
DTSBN-S 32
5.90 = 0.002
npODE
3
4.79
25
NeuralODE
4.74 0.093t
0
ODE VAF 34
3.17 = 0.221
25
ODE VAE-KLL
34
2.84 = 0.343t
R1
Latent ODE
2.45 + 0.057*
Latent SDE
12
2.01 + 0.050*
ARCTA (ours)
2.76 + 0.168
5
20
25
30
((s)
Figure 4: MOCAP benchmarking results, from 33 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--rt methods. The plot shows the predictive posterior on the test set for some select outputs. Other benchmark results were compiled in [34 [2] RMSE was computed from MSE by taking the square root of the mean and transforming the error via a first-order Taylor-series series approximation.
In this experiment we consider the motion capture dataset from 132 The dataset consists of 16 training, 3 validation, and 4 independent test sequences of subject walking. Each sequence consists of 300 time-series observations with a dimension of 50. We made use of the preprocessed data chose latent dimension of 6. We
from (34]. Like previous approaches tackling this dataset, assume Gaussian observation likelihood. We place a log-normal prior on the diagonal elements of the diffusion matrix and the noise the observations. We approximate the posterior of the diffusion matrix and observation noise covariance using log-normal approximate posterior. Details on our architecture design and hyperparameter selection provided in Appendix G.33
For our approach, we train 10 models and report their average performance on test set due 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 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-arr performance significantly reduced computational cost as compared to adjoint based methods.
4.4 Neural SDEs from video
2(t)
((s)
Figure 5: Neural SDEs from video. Here we used five frames to estimate the intial 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 predictive posterior overlaid on the data.
In this experiment we attempt to learn a latent SDE from video. We generated 32 32 black and white frames of a nonlinear pendulum as it evolves for 30 seconds collecting data at frequency of 15Hz. We transform the 1024 dimensional state down two dimensions using a convolutional architecture. Details the hyperparameters and architecture are provided in Appendix G.4 This problem similar to the problem considered in except the dynamical system we consider nonlinear. In this prior work, the authors were forced to regularize the latent space so that one set of coordinates resembles a generalized velocity. the present work, no such regularization required.
We assume 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 log-normal variational distribution. After training we generate a forecast that is visualized in Figure [5] We see that we were successfully able build a generative model for this simple video. This result demonstrates the broad applicability of the present approach to standard generative modeling tasks.
Conclusions
Here we have presented 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-Cate approximations to expectations using a novel amortization strategy and a recently derived change of variables for expectations under Markov Gaussian processes [3
We have demonstrated the efficacy our approach in learning latent SDEs across a range test problems. In particular showed that approach 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 the expectation over time further reduce the total number required function evalu- ations. There are also opportunities to explore the utility of the proposed approach for generative modeling on more realistic problems. Finally, there are opportunities apply our work the context of implicit densities [35].
10

---

Acknowledgments and Disclosure of Funding
This research funded by a grant from NSERC.
References
[1] Ricky Q. Chen, Yulia Rubanova, Jesse Bettencourt, David Duvenaud. Neural Ordinary Differential Equations. Advances neural information processing systems, pages 6571-6583, 2018.
[2] Sepp Hochreiter and Jürgen Schmidhuber. Long short-term memory. Neurul computation, 9(8):1335-1780
1997. Publisher: MIT press.
[3] Sam Greydanus, Misko Dzamba, and Jason Yosinski. Hamiltonian Neural Networks. In Advances in Neural Information Processing Systems, pages 15353-15363, 2019.
Miles Cranmer, Sam Greydanus. Stephan Hoyer, Peter Battaglia, David Spergel, and Shirley Ho. La- grangian Neural Networks. ICLR 2020 Workshop on Integration of Deep Neural Models and Differential Equations, 2020.
[5] Kevin Course, Trefor Evans, and Prasanth B. Nair. Weak Form Generalized Hamiltonian Learning. In Advances in Neural Information Processing Systems, volume 33, pages 18716-18722 2020.
[6] Emmanuel Lorin. Derivation and analysis of parallel-in-time neural ordinary differential equations. Annals of Mathematics and Artificial Intelligence, 88: 1035-1099, 2020. Publisher: Springer.
Jacob Kelly, Jesse Bettencourt, Matthew Johnson, and David Duvenaud. Learning differential equations that are easy to solve. Advances in Neural Information Processing Systems, 33:4370-4380, 2020.
[8] Chris Finlay, Jörn-Henrik Jacobsen, Levon Nurbekyan, and Adam Oberman. How to train your neural ODE: the world of Jacobian and kinetic regularization. In International conference on machine learning, pages 1464,1,4, 2020.
[9] Patrick Kidger, Ricky T. Q Chen, and Terry J Lyons. "Hey, that's not ODE": Faster ODE Adjoints via Seminorms. Proceedings of the 38th International Conference on Machine Learning, volume 139 of Proceedings of Machine Learning Research, pages 5443 -452,, July 2021.
[10] Daniel J Lea. Myles R Allen, and Thomas WN Haine. Sensitivity analysis of the climate of chaotic system. Tellus A: Dynamic Meteorology and Oceanography, 52(5):523-532 532, 2000. Publisher: Taylor & Francis.
[11] Yulia Rubanova, Tian Qi Chen, and David K Duvenaud. Latent Ordinary Differential Equations for Seggularl-Saamplee Time Series. Advances in Neural Information Processing Systems, pages 5321 5331, 2019.
[12] Xuechen Li, Ting-Kam Kam Leonard Wong. Ricky T. Q Chen, and David Duvenaud. Scalable Gradients for Stochastic Differential Equations. In Silvia Chiappa and Roberto Calandra, editors. Proceedings of the Twenry Third International Conference Artificial Intelligence and Staristics, volume 108 of Proceedings of Machine Learning Research, pages 3870-3882, August 2020.
[13] Kevin Course Prasanth B. Nair. State estimation of physical system with unknown governing equations. Nature, 622(7982):261-22 267-2277, October 2023.
[14] Mike Giles and Paul Glasserman. Smoking adjoints: Fast monte carlo greeks. Risk, 19(1):88-92 2006.
[15] Diederik Kingma and Max Welling. Auto-Encoding Variational Bayes. In International Conference on Learning Representations, 2014.
[16] Matthew D Hoffman, David M Blei, Chong Wang, and John Paisley. Stochastic variational inference. The Journal of Machine Learning Research, 14(1):1303-1347, 2013. Publisher: JMIRR. org.
[17] Peter Toth, Danilo Rezende, Andrew Jaegle, Sébastien Racanière, Aleksandar Botev, and Irina Higgins. Hamiltonian Generative Networks. International Conference on Learning Representations, 2020.
[18] James Glimm and David Sharp. Stochastic Partial Differential Equations: Selected Applications In
Continuum Physics. In Stochastic Partial Differential Equations: Six Perspectives, Mathematical Surveys and Monographs, pages 3 44, American Mathematical Society, 1997.

---

[19] Simo Särkkä and Arno Solin. Applied stochastic differentiul equations, volume 10. Cambridge University Press, 2019.
[20] Cedric Archambeau, Dan Cornford Manfred Opper, and John Shawe-Taylor Gaussian process approx- imations of stochastic differential equations. In Gaussian Processes in Practice, pages 1-16. PMLR. 2007.
[21] Cédric Archambeau, Manfred Opper, Yuan Shen, Dan Cornford, and John Shawe-taylor. Variational Inference Diffusion Processes. In Advances in Neural Information Processing Systems, volume 20.
Curran Associates, Inc., 2007.
[22] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. Advances in neural information processing systems, 30, 2017.
[23] Timothy D. Barfoot. State estimation for robotics a matrix Lie group approach. Cambridge University Press, Cambridge, United Kingdom, 2017. ISBN 178-1-107-15939-6
[24] Timothy D Barfoot, James R Forbes, and David J Yoon. Exactly sparse Gaussian variational inference with application to derivative-fre free batch nonlinear state estimation. The International Journal of Robotics Research, 39(13)14473 1502, 2020. Publisher: SAGE Publications Sage UK: London, England.
[25] Danilo Rezende and Shakir Mohamed. Variational inference with normalizing flows. In International conference machine learning, pages 1530-1138... PMLR, 2015.
[26] Mihai Alexe and Adrian Sandu. On the discrete adjoints of adaptive time stepping algorithms. Journal of computational and applied mathematics, 233(4):1005 1020, 2009. ISSN 0377-0427 doi: 10.1016/j.cam. 2009.08.109 Place: Kidlington Publisher: Elsevier B.
[27] Hayden Schaeffer and Scott G. McCalla. Sparse model selection via integral terms. Physical Review E, 96 (2):023302, August 2017. doi: 10. 1..110//hyssceE..6... 20..0s22e..99022332..
[28] Yannis Pantazis and loannis Tsamardinos. A unified approach for sparse dynamical system inference from temporal measurements. Bieinformatics, 35(18))3387-3308 September 2019. ISSN 1367-4803. doi: 10. 1..0990boonfoomattess/0000....5.
[29] Daniel A Messenger and David M Bortz Weak SINDy: dalerkin-based data driven model selection. Multiscale Modeling & Simulation, 19(3):1474 1497, 2021. Publisher: SIAM.
[30] Diederik P Kingma and Jimmy Ba. Adam: method for stochastic optimization. International Conference on Learning Representations, 2015.
[31] Qiqi Wang. Rui Hu, and Patrick Blonigan. Least squares shadowing sensitivity analysis of chaotic limit cycle oscillations. Journal of Computational Physics, 267:210-224, 2014. Publisher: Elsevier.
[32] Zhe Gan, Chunyuan Li, Ricardo Henao, David Carlson, and Lawrence Carin. Deep temporal sigmoid belief networks for sequence modeling. Advances in Neural Information Processing Systems, 28, 2015.
[33] Markus Heinonen, Cagatay Yildiz, Hearik Mannerström, Jukka Intosalmi, and Harri Lähdesmääi.. Learning unknown ODE models with Gaussian processes. In International Conference on Machine Learning, pages 1959-1968. PMLR, 2018.
[34] Cagatay Yildiz, Markus Heinonen, and Harri Lahdesmaki. ODE2VAE: Deep generative second order ODEs with Bayesian neural networks. Advances in Neural Information Processing Systems, 32, 2019.
[35] Patrick Kidger, James Foster, Xuechen Li, and Terry J Lyons. Neural sdes as infimiteddimensional gans. In International Conference on Machine Learning, pages 5453-5463 PMLR, 2021.
[36] Andrew Gordon Wilson, Zhiting Hu, Ruslan Salakhutdinov. and Eric P Xing. Deep kernel learning. In Artificial intelligence and statistics, pages 370-378. PMLR, 2016.
[37] Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, Alban Desmaison, Andreas Kopf, Edward Yang. Zachary Vetoto, Martin Raison, Alykhan Tejani, Sasank Chilamkurthy, Benoit Steiner. Lu Fang. Junjie Bai, and Soumith Chintala. PyTorch: Imperative Style, Riih-Peerfrmmncce Deep Learning Library. In H. Wallach, H. arochelle, A. Beygelzimer, F. d' Alché- Buc, E. Fox, and R. Gameett editors, Advances in Newral Information Processing Systems 32, pages 8024-8035. Curran Associates, Inc., 2019.
12

---

Appendices
Amortized Reparametrization: Efficient and Scalable Variational Inference for Latent SDEs
A Expectations under linear SDEs
In this section we re-derive a result from [3] 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 [3] rewrite the ELBO entirely terms of quantities that do not require differential
equation solvers.
Theorem 1. Consider the density, q(2(t)), the solution to a linear SDE de = (-A(t)= b(t)) - L()dd with initial condition za ~ (m,, So), where : R - Rdxd is symmetric, b IR + TRd, : + gdxd, and 3 indicates Brownian motion with diffusion matrix E. Then the expected value of bounded functional, f. satisfies
(14) where B(t) vec-1((S(() @ ((e)) - S(t)) and m(t) and St) indicate the mean and covariance, respectively, of the SDE solution as time t.
(lt(t))())())) ((A((),)())
(()(())) Bl (BB) + B())m(1);(1)),,
Proof. The solution of a linear SDE defines Markov Gaussian process with marginal statistics, ((((()) ~ ((m((),S(()) given by the solution to the ODEs,
ri(() = (-A(t)m + b(t)),
(15)
$(t) -(t()(((t) -(())((t) L(t)EL(1)*
(16)
with initial condition m(0) mo. S(0) So [9]. Noticing that equation (16) defines of matrix Lyapunov equations in terms of A(t) allows us to express A(t) as function of S(t) as follows,
A(t) vec (S(()) @ S(t)) vec - -5(t))
(17)
where is called the Kronecker sum and defined as S 8 S @ I and indicates the standard Kronecker product. Letting Btt) = A(t) be the expression for A(t) written in terms of ((),, we can rearrange Equation (15) solve for b(t) as,
b(l) = (l) B(1)m(I).
(18)
Substituting the expressions for A(t) and b(t) into Equation yields the desired result.
Remark In the case that S(t) and ((t) diagonal assume throughout this work, Equation 17 simplifies to,
A(t) = 2 -U))) 5()).
(19)
B Evidence lower bound derivation
Given a dataset of observations, D ((x,, - generative likelihood, po(x =(1,)) a prior SDE, dz = fz(2()), ft + (o)d), and an approximation to the posterior of the latent state, bo(l) (t) (t) is possible derive the ELBO 20 12).
dz =
((),0) Σ E Be(l))=aa(=(1)) Plog x(44))
(20)
2 J Ee(t)rns(t(e) [||
(1) + b(t) fo(=(t), dt,
13

---

where Co(t) (Lo(t)LL (t)) The second term in 20 contains an integral over time of terms of the form which matches that of the expectation (1).. From this observation, Course and Nair [33 applied the identity in 14 to derive the reparametrized ELBO,
[ x)
Poog =x(,)))
(21)
/ 2
1
(2(0mgo(1(0)) [Ira,(()(0))))))mc)) dt,
where,
(t( = (Lo()))(o(()")
ro, (=))() = ((() - ==(1) + (1)(1) (0((())) B(I) vvc ((S @ (e()))"voo(Lo(()LLol) (c(1))
(22)
To reiterate what mentioned in the main body of the present work, the advantage of this reparametrized ELBO 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 is suggested in the main body of the present work we arrive at the final ELBO,
N/M M
((0,0)
Σ
E 2o(=()
(.)) logpo(z))) =
j=1
(23)
2 1 (o(((1) . [1ro,o(2. dr.
1
E
Again, 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 cireumstances, it will be advantageous to place priors on subset of the parameters, Using the tools of stochastic variational inference, can infer the posterior 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 normal prior on the diagonal of the diffusion process term C7 (t) and approximated the posterior using log-normal approximate posterior, Appendix E for details. While not necessary, we found that doing so helped to stabilize training in the examples 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 [ Given particular choice of prior, p(0), along with an approximate posterior, go (0), we can amend the previously derived ELBO as,
N/M Σ
[(0)=
Σ
E 9o
(9 (....))) log (ogpo(a ((((()))]
j=1
(24)
1
E
[1ro, (1),)) dt DKI(gs(0) (l)(
2
2,0
where DxL. (q p) indicates the KL divergence between g and p.
Consider a reparametrization to the latent state as z(t) = ett,e, where p(e) z(1) (s()(t) 2 (1)
..... x (1)
Also consider reparametrization to the approximate posterior for the
14

---

generative variables, 0 = V(v, 6) where V p(v) =
θ ~ (o( Then, building on Lemma
|
we can arrive at an unbiased estimate for the gradient of the modified ELBO which inherits all the properties discussed with the original approximation:
M
R
V VoL(o)
Pv PV (oto) P (x)
R
1=1 k=1
R
(25)
-(4 ((++1)
(3)
V (o (()), ,())))))))))))) ,0, )
2R
vall w)
k=1
Dxx p(0)).
where each (()) ~u(()) ( each (8) ~ NO 7), and each (R)) p(v).
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 useful for our purposes because it interpretable and stable in training. With this being said, any encoder which can transform batch of observations down to reduced dimension latent state can be used combination with Lemma D 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 dataset, D= = [(t,, = recall that the first step in our amortization strategy is to split the dataset into N/M non-overlapping partitions.
original indexing: [ti, 12, t..,,tt,,,aa tss IN
(1) ..... ot
MM (N/10)
reindexed dataset:
Recall that we would like to approximate the latent state over each partition using only the M (1)
observations in each partition, ga(=(() =.(" 422 ................................ ()) s p(z(() D) for l E [160, ((+11)) Going forward we will drop writing the superscript as we will be only working on a single partition, (t1, 21), (f2, 22).... (tM, 2M). 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; our own studies we found choosing K [1,5] worked well for the problems considered.
Given some encoding network, ENC,, we compute:
h; ENC, Fi+1, (T), x(+K).
(26)
where h4 E RR2d with d < D. We will describe the particular architecture design for ENC, 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.c. 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 consider the example of inferring a latent SDE using video of pendulum as we did in Section 4.4 If we choose single frame near the center, we have no way of knowing from that frame alone the pendulum currently swinging left 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 multimodal. As our approach only allows for one to approximate the latent state using MGPs, this is not option. Allowing the encoder to take in a few frames at time remedies this issue. We also note that previous works for inferring latent differential equations made this same choice [2 [7].
Recall that we need to approximate the latent state at any time over the window of observations the partition, E (10) ((--1). To accomplish this we effectively interpolate between encodings using a deep kemel (36] Letting.
h
H
(27)
AT
15

---

where H E IR 2x2d and fnode = [ts, f....., tarl, we construct the encoder for the mean and diagonal of the covariance over the latent state as,
[m)(t) log Sp(1)²] = (t, (made) (ka (mode) -
(28)
Here we note that the right hand side of (28) is row vector of length 2d. We use the notation
(t)( log (t))) to indicate that (t) is given by the first d elements of this vector and log (t) is given by the next d clements. Here Ki is so-called deep kernel and kp(t, fnede) E RM and Kig (fnodo (node) IR MMMM xM In addition on IR+ is tuned as part of the optimization procedure. While many options for the base kernel are possible, we made use a squared exponential kernel,
[IDK& DK DK (1.)11
& (t, t.) f exp
(29)
2
where o,,6 E IR+ are positive constants tuned as a part of the optimization procedure and DK, is 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 -tatttoot 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 [ma((;)² log (1.)²) Rs 1i. and it is cheap to compute long M is relatively small because evaluations of ENC, can be performed in parallel. Particular choices for ENC, and g are described in context in Appendix G
E Approximate posterior on diffusion term
This section summarizes the normallal parameterization we used to approximate the posterior over the diagonal drift function terms in the main body the paper. First, we note that the loss function we only require access to the product, Ca(t) = (t) so, rather than parametrizing Lg on its own, we parametrize Cz!
Specifically we parametrize Co(t)- -1 = (i), where 0 E Rd. The prior is defined as,
p(0) = II CN (i, i=1
(30)
We parametrize the approximate posterior as,
(0) II CN((0, p4 ()) C-l
(31)
where br and a; are the variational parameters. The KL divergence between the posterior and prior is given by,
1 d
o
(14
(2,)2
DKI(gs (()) Σ (logā; loga,)
Σ i=1
(32)
07
F Computing resources
Experiments were performed on an Ubuntu server with a dual E5-2680 v3 with 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 torchdiffeq [I, torchsde 11 and
pytorch_ lightning. All code is available at gathubbcom//oouskkiiiaaaattnnttt
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 encoder
16

---

consisting of two neural networks, (as............... -(+KK and (t, (see Appendix D), decoder (=C)( and model for the SDE consisting of a drift, fo(t, and dispersion matrix, Lo For all numerical studies we place log-normal prior on the dispersion matrix and assume that the approximate posterior is constant in time, see Appendix For all experiments we gradually increased the value of the KI-diiergence (both the KL--divereence due to the SDE prior and the KI--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, tasked neural ODE (NODE) and our approach with building predictive model for the Lotka- Volterra system given a dataset of time-series observations. The Lotka- Volterra equations system of nonlinear ODEs usually written as,
t az Bay,
(33)
ý dxy 7V.
In our experiment we chose a 2/3, B 4/3, and 8 = = We also assumed that there was some small amount of Brownian noise given by Σ = diag(10 3000) 10- Using the initial condition 0.9 and y = 0.2, we draw a sample from the system using the default adaptive stepping scheme in torchsde 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 standard deviation of 0.01. We use the first 50 seconds for training and reserve the remaining 15 seconds for validation.
frequency of 50Hz and added Gaussian noise with a
Both the NODE and our approach use the same model for the gradient field and drift function respectively, see Figure 6al The encoder and deep kernel architecture are provided in Figures 6b and respectivvey 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 Gaussian likelihood with a constant standard deviation of 0.01.
In terms of hyperparameters we set the schedule 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 Ir = --r every iteration with 7 = exp(Iog(0.9) /1000) (i.e. the effective rate of learning rate decay is Ir = 0.9/r every 1000 iterations.). We used the nested Monte-Carde approximation described in Equation (TO) with R 1, S 10, and M 256. In terms of kernel parameters, we initialized & = 10 of 1, and on = 10 In terms of the diffusion term, set [4s = 0; = 10-5 and رم; = न: = 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 (31.1)). As a reminder, given time-series dataset of observations, our goal was to infer the value of the parameters, were likely to have generated the data starting from an initial guess: 6u = lou. Po. Pu]. The true value of the parameters was chosen as 8_ = [10, 8/3, 28). For all experiments we used the initial condition [8. 28.32,36.00 as was suggested in [00 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 frequency of 200Hz over the time intervals [0, 1], [0, 10). [o, 50), and [0, 100). For each time interval we generated 5 datasets by adding independent Gaussian noise with variance of 1 to the data.
8, which
To arrive at an initial guess we sample from the distribution 0g ~ (0-,,, (0.20.) ))) 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 fz-norm of the gradient of the ELBO with respect to the parameters o, B, and p after optimizing for 2000 iterations. For the adjoint method, we report the gradient of the function:
1 (())
Σ
(x: - xz (1:)) + (y) vo + (z+ - =0)))))
(34)
3N
17

---

FC-64
ReLU
F2
FC-64
FC-32
FC-32
ReLU
ReLU
ReLUU
FC-64
FC-32
FC-32
+
ReLU
ReLU
ReLU
FC-2
FC-4
FC-1
drift function
m(4;
log S(4)
nonstationary transform
(a) at
(b) (x.)
(c) (Kx
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(t4), log (t+:) to show how have split the columns of h; in two. The value of (m(()), log Stt,)
only = h; unless on = 0, see Appendix D Note the arrow from x, to m(t+) indicates a residual connection (which was useful in this case because we are learning a SDE in the original data coordinates).
FC-32
ReLU
t
FC-32
FC-32
ReLU
tanh
+
FC-6
FC-1
mt.))
log St.,)
nonstationary transform
(a) (2:)
(b) (1)
Figure 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.), log St,,) to show how we have split the columns of h; in two. The value of (mn(,), log S(t,) only = h; unless o, = 0, see Appendix D Note the arrow from to m(t,) indicates residual connection (which was useful in this case because we are learning a SDE in the original data coordinates).
18

---

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, = 10, and = 100. We selected a learning rate of 0.1 and decayed the learning rate Ir yIr every iteration with y exp(log(0.9)) (0.))/0000)0000)) 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 0f = 1, l 10-?, and nn 10- For the diffusion term, we set μ44 o, 10-5 and Hr o, 10
G.3 MOCAP experiment
For this experiment we use the preprocessed dataset provided by [34].
Like in previous examples, we place log-normal prior 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 g 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 and over the course of 200 iterations. We make use of the nested Monte-Carlo scheme in equation (1O) 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 Ir = ylr with
7 (oxp(oog00.9)/141 10000). In terms of kernel parameters,
initialize on 10 L = 10 10-2, and
o = In terms of the prior on the diffusion term we initialize μ4; = a. = 10 and set fı: = 0, = 1. In terms of the prior on the variance of the observations, we initialize μt; o; 10-2 and set ju; 5; 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 squa---uure--error as, MSE + ERROR. We report RMSE + NEW ERROR that error units are consistent with the units of the original dataset. To convert MSE to RMSE we used first-order Taylor-series approximation,
RMSE = VMSE
(35)
NEW ERROR /RMS
G.4 Neural SDE from video
In this section, provide more detailed description the numerical study described in Section 4.4 We generated data by simulating a nonlinear pendulum with the equations,
t =
(36)
p sin(2)) (x),
for 30 seconds while sampling the state frequency of 15Hz. The architecture used for this experiment is provided in Figure
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 learning rate 0.001 and decayed the learning rate each iteration according to the schedule = 7ir with 7 exp(log(0.9)) (0.9)/ / /1000). In terms of kernel parameters, and f f= 10-2 We placed 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 μ4 o, 0.1 and set jı, ō, 10-5
initialized o, 1, o2 = 10-5,
19

---

Ti+1
2i+2
z(t)
FC-30
FC-15
+
tanh
tanh
FC-30
FC-30
tanh
tanh
FC-15
FC-30
tanh
tanh
FC-6
FC-50
h4
p(t)
(b) (FC
(a) (2,)
x
FC-32
FC-30
tanh
tanh
FC-1
FC-6
nonstationary transform
drift
(c) (t)
(d) fo(
Figure 8: Architecture diagrams for the encoder, decoder, deep kernel and drift function used in the
MOCAP benchmark. We used very similar architectures to [3 [2]. Here u(t) indicates the mean of the likelihood, (z zt))..
20

---

Fi+1
31+5
Conv-16-3
BN-ReLU
x(t)
Pool-Max-2
FC-(64 x 32 x 32)
Conv-32-3
ReLU
BN-ReLU
FC-(64 x 32 32)
Pool-Max-2
ReLU
FC- 128
ConvT-32-4
BN-ReLU
BN-ReLU
FC-4
ConvT-1 1-1
hq
logit
(a) (x:)
(b) DEC, (t)
FC-32
FC-64
ReLU
tanh
FC-32
FC-64
ReLU
tanh
FC-1
FC-2
nonstationary transform
drift
(c) (t)
(d) fz,
Figure 9: Architecture diagrams for the encoder, decoder, deep kernel and drift function used in the Neural SDE from video example.
21

---

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 predictive model for dimensinnii 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,
31
1 31 (a1 B191 - 71122
(37)
h1
32
82 22 (a2 B2y11 - (2122)
(38)
k2
ú1 y1 01 + 61æ9 + E122
v192).
(39)
1/2 (-02 + 1 + C2 + M231).
(40)
where o; is the grow rate the prey vi, 3 is the rate that predator y1 consuming prey 2i 7/ the rate that predator 1/2 is consuming prey Fg, kę is the carrying capacity for prey Jy 6, is the death rate of predator yi, & is the conversion rate for predator yi from x1, & is the conversion rate for predator y/ from prey 22, and иn represents the competitive effects yı caused by the other predator.
We simulated the system for 300 units of time collecting data at frequency of 10Hz. We assume Gaussian noise with standard deviation of 10- For all experiments we used the following hyperparameters: batch size of 256, 1000 warmup iterations, and learning rate of 10-3 In terms of kernel parameters, we initialized o f = 10 2 and on 10 For the diffusion term we set = = 10 and iu; = ō. = 1. The architecture description for the neural networks used in this example provided in Figure
Results are summarized in Figure below. On this problem we find we are able to achieve a reasonable validation accuracy for S 10. S 50, and S 100; however, 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.
R=
8
50
R=1, 8
50
log10 Validation RMSE
-02
logio Validation RMSE
-0.2
-0.4
-0.4
-0.6
-0.6
-0.8
-0.8
-1.0
-1.0
-1.2
-1.2
500
1500
50000
100000
150000
200000
Iteration
NFE
Figure 10: Validation RMSE versus iteration (L.) and number of function evaluations (R).
22

---

Z;
FC-32
ReLU
FC-32
ReLU
FC-4
mtt
log St+,)
(a) (2.)
t
FC-32
FC-64
ReLU
ReLU
FC-32
FC-64
ReLU
ReLU
FC-1
FC-2
nonstationary transform
drift function
(h) (t)
(c) fo(=,t)
Figure 11: Architecture diagrams for the encoder, deep kernel and drift function used in the Monte-
Carlo parameters study.
23


Reviews:
[{"review_id": "9BMFcq9Qoj", "reviewer": "Reviewer_8z7R", "summary": "The present work proposes a time and memory efficient method to perform inference in a directed probabilistic model in the presence of a latent stochastic process with intractable posterior (path) distribution. Learning is performed using a variational Auto-Encoding approach, in which an approximate latent posterior process is selected from the class of Markov-Gaussian processes that are solutions to linear non-autonomous (Itô) SDEs. An evidence lower bound objective (ELBO) is derived that removes the need for numerical solvers when approximating gradients. Included experiments on toy and  real world time series data approve similar performances compared to methods based on adjoint sensitivities but with two-orders of magnitude fewer evaluations and with a time and memory cost that scales independently with the amount of data, the total length of the time series and the stiffness of the approximate differential equations.", "questions": "- (l. 32f.) \"... replacing ordinary differential equation (ODE) solvers with integration methods ...\" But don't those numerical ODE solver build on integration methods?\n- In terms of the total number of time series samples included, the experiments dealt with are probably rather small-scale. To really highlight the added value of the proposed method, I would like to see large-scale experiments like the PhysioNet 2012 experiment in [12].\n- Minor typos:\n  - (l. 28) parareal -> parallel\n  - (l. 54/67) time-series -> time series\n  - (l. 66) Description -> description", "limitations": "The authors address the limitations of the approach regarding the models expressivity compared to existing stochastic adjoint sensitivity methods in Chapter 3 of their work.\nI would appreciate some remarks on limitations concerning the recognition module, i.e., the encoder. As stated in Section D, data sparsity is only accepted to a certain extent and cf. (l 495f.) \"It is worth noting that without this deep-kernel our approach struggled to achieve good validation accuracy on the datasets we considered\". This suggests that the correct choice of the encoder is crucial and, on the other hand, raises doubts about the contribution of the latent stochastic dynamic. Is the proposed approach limited to interpolation and time series classification tasks or may we solve also tackle extrapolation problems?\nNo evidence was found that the authors addressed any potential negative societal impact. Please make the effort to include this mandatory information in your work.", "rating": 3, "confidence": 4, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "I agree, a huge drawback of those existing models is the dependency on stochastic optimization methods that require backpropagading gradients through the numerical (stochastic) differential equation solver and hence rely on adjoint methods.\nIt is the attempt itself tackling the road blocks by combining a variety of ideas that rely on model simplification, that I particularly appreciate. Especially under the lens of Bayesian inference it is recommended to question the need for unbounded expressivity in case of learning latent dynamics with continuous time models. This consideration is interesting on its own and adds progress to the community, e.g., regarding the applicability on low performance hardware.\n\nThe authors propose improvement in time and memory cost by combining the following three major strategies: (1) simplifying the function class available for learning latent dynamics while reducing the expressiveness of the latent (SDE) models used, (2) exploiting the nature of sequential data such as time series that allow divide-and-conquer strategies such as the analysis of small sliding observation windows, and (3) introducing an optimized objective (i.e. an ELBO) by exploiting the reparameterization trick within the model used.\n\nThis method builds on a novel combination of well-known techniques, and is original to the best of my knowledge. Further, the paper is well positioned to existing work, the content is well organized and, with a few exceptions, has no spelling or grammatical flaws. Empirical evidence is accessed through experiments on toy data as well as established benchmark data from real world applications; however, I would classify them small-scale.", "weaknesses": "It is not the value of the approach I am concerning here, but the shortcomings related to the clarity of the submission and the technical soundness which I will address with examples in the following.\n\n1. (2.1 Problem description)\n   - I am missing a clear definition of the observation model, e.g., cf [Vrettas et al. 2015; Variational mean-field algorithm for efficient inference in large systems of stochastic differential\nequations ; Section II.A].\n   - Considering time $t_i \\in \\mathbb{R}$ instead of $\\mathbb{R}_{+}$ is not wrong but unusal.\n   - Minor ambiguities make understanding difficult in total: $dz$ vs $dz(t)$ and $\\beta$ vs $\\beta(t)$; the definition of $\\theta$ is missing.\n   - Is there any motivations for the stochastic variational inference approach?\n   - I would $q_\\phi(z(t)| x_i,\\dots,x_{i+M}) \\approx p_{\\theta}(z(t)|\\mathcal{D})$ as an local approximate posterior with respect to the time variable; are there necessary and/or sufficient conditions for an observed stochastic process that guarantee the validity of this approximation?\n2. (2.2 Evidence lower bound)\n   - (l. 96f) The following statement: \"An added complication in our specific case comes from the fact that the latent state is defined by a stochastic process rather than a random-vector as is more typical.\" is poorly worded. You probably intend to say that, typically, stochastic variational inference considers vector-valued random variables to define latent states. But your approach requires function valued random variables, i.e., stochastic processes. What are those complications expected here?\n   - It would be advantageous in combination with Eq. (1) to mention that the time dependency of the drift coefficient is necessary to account for a possible non-stationary observation process $x(t)$ that is driven by the latent stochastic dynamics of $z(t)$; cf. [32; Section 3].\n   - (Eq. (3)): To derive this result, the authors refer to unpublished work without stating the authors, cf. [14], but they included the reference in the supplementary material submitted. Notwithstanding the fact that I cannot call this good practice, this result is based on known facts and the derivation could have been included in the supplementary material of the present work, for example. Besides that, the derivation in [14] only concerns the case of a constant dispersion matrix $\\Sigma$ with respect to time $t$ whereas $L(t)$, e.g., in Eq. (1) depends on time.\n   - (Eq. (4)): Section A of the supplementary material indicates that this result follows by applying properties concerning a Lyapunov equation and by using the Kronecker sum. Following the arguments myself raised difficulties; in my opinion some given identities only hold in case of $A(t)$ is symmetric, i.e., $A(t) = A\\top(t)$. Last but not least titling the result Theorem overstates its contribution.\n   - (Remark 2): I would like to see the form of the prior SDE. In the setting of stochastic variational inference it is important to well design the dispersion matrix between prior and posterior SDE in order to guarantee that the KL divergence is finite in value.\n   - Honestly I still don't get how the authors include the fact, that $m_\\varphi$ and $\\dot{m}_\\varphi$ as well as \n    $S_\\varphi$ and $\\dot S_\\varphi$ are related through an ODE? Because all four are included into $r_\\theta, _\\varphi$ and e.g., Section G.1 only shows architecture design for $m$ and $S$.\n3. (Experiments)\n    - (Figure 2): Can you quantify \"a similar validation accuracy\"? Does the right plot only contain the ARCTA probabilistic prediction? Section F.1 does not contain the promised information, did you mean Section G.1?\n    - (Section 4.2; l. 259ff.): The evaluation protocol is unclear to me; why are there multiple adjoint methods included and do you only optimize ARCTA for 2000 iterations? Can the quality be evaluated using error measures, e.g. MSE? A very similar experiment was presented in [13] proposing an adjoint method for SDEs. The results are promising. Can you please take these results into account?\n    - (Section 4.3): From my point of view, the motion capture data set under consideration is quite delicate, since I know the shortcomings very well from my own experiments with it in the past. In my opinion it is the overall size of time series contained (16 training, 3 validation and 4 testing) that frequently causes severe overfitting. So I am still bothered with severe instabilities when reproducing the results from [30 & 13]. Therefore, I also view the results presented here with caution. There are 4 times 50 = 200 predictive posterior outputs by the model, evaluated on the test dataset; given four test time series you only provide three single dimension, that is strange. Secondly, I would like to know why the evaluation of the results is carried out in relation to the RMSE, requiring all benchmark results to be translated.\n\nFinally, I would like to reiterate my appreciation for the authors' interest in expanding the field of learning latent SDEs in the setting of stochastic variational inference. I'm sure that by working on the deficiencies mentioned, the work will gain in value and will deliver the desired contribution to the community in a new round.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your comments.\n\n**“The authors refer to unpublished work without stating the authors, cf. [14], but they included the reference in the supplementary material submitted. Notwithstanding the fact that I cannot call this good practice, this result is based on known facts and the derivation could have been included in the supplementary material of the present work, for example.\"**\n \nFor papers which are in review but not available as a non-anonymous preprint, the recommendation from the 2023 call for papers is to, “include a copy of the cited anonymized submission in the supplementary material and write ‘Anonymous et al. [1] concurrently show.’” The paper included in the supplementary material is currently under review and not available as a non-anonymous preprint. Our apologies but we are not sure what known facts you are referring to since no supporting references are listed in your comment.\n \n**\"No evidence was found that the authors addressed any potential negative societal impact. Please make the effort to include this mandatory information in your work.\"**\n \nIn the paper checklist we answered n/a for the question on broader impacts. The paper checklist states, “if you develop a generic algorithm for optimizing neural networks, you do not need to mention that this could enable people to train models that generate Deepfakes faster.” In our perspective, our work can be seen as providing a method for training neural SDEs more quickly than the status quo. \n \n**“Parareal -> Parallel”**\n \nSorry, this is not a typo. Parareal methods are a class of time-marching schemes with a history spanning over two decades; see, for example, reference [6]. \n\n**“Following the arguments myself raised difficulties; in my opinion some given identities only hold in case of A(t) is symmetric. Last but not least titling the result Theorem overstates its contribution.”**\n\nThe theorem is applicable when A(t) is nonsymmetric. We believe that this is an important theoretical result with significant practical impact. This result allows us to recast expectations under linear SDEs entirely in terms of their marginal distributions without requiring differential equality constraints like in [D]. From a practical perspective, this allows us to evaluate expectation under linear SDEs using either ($A$, $b$, ($m(0)$, $S(0)$)) or $(m(t), S(t))$. \n\n[D] Archambeau, C., Opper, M., Shen, Y., Cornford, D., Shawe-taylor, J. “Variational Inference for Diffusion Processes.” In: Advances in Neural Information Processing Systems, vol. 20. (2007).\n\n**“Honestly I still don't get how the authors include the fact, that m and dm/dt as well as S and dS/dt and are related through an ODE?”**\n\nApologies if we have misunderstood your question.\nIn lines 102-104 we provide an explanation for this standard result on SDEs along with a link to a textbook for the benefit of readers who are new to this topic.\n\n**”Can you quantify \"a similar validation accuracy\"? Does the right plot only contain the ARCTA probabilistic prediction? Section F.1 does not contain the promised information, did you mean Section G.1?”**\n\nApologies for the confusion. This should have been section G.1. The right plot only contains one ARCTA prediction. The left plot contains the validation accuracy vs. number of function evaluations. \n\nWe trained 10 ARCTA models and 10 NODEs with different tolerances for the adaptive stepping schemes on the same dataset with different random seeds. The purpose of this study was to show that our approach requires a fewer number of evaluations of the model on average to achieve a similar validation accuracy. To make our claim more conservative, we now state our approach \"requires more than one order of magnitude fewer evaluations of the model...\" Since submitting the original study we have also improved our hyperparameter selection for our approach. The new Figure 2 is attached. \n \n **“A very similar experiment was presented in [13] proposing an adjoint method for SDEs. The results are promising. Can you please take these results into account?”** \n\nThe numerical study in [13] considers a fundamentally different problem. The authors in this work collect data for a short time-window of 1 second. The purpose of the numerical study in our work is to demonstrate that adjoint sensitivities for chaotic systems can blow up when the time-window of interest is increased (we show this for the case when the time-window is set to 10, 50, and 100 seconds). See references [10] & [11] for more details on this phenomena. \n\n**“In my opinion it is the overall size of time series contained (16 training, 3 validation and 4 testing) that frequently causes severe overfitting ... Therefore, I also view the results presented here with caution. There are 4 times 50 = 200 predictive posterior outputs by the model, evaluated on the test dataset; given four test time series you only provide three single dimension, that is strange.”** \n\nWe evaluated the test accuracy on the entire test dataset. We plot predictions on three trajectories for illustrative purposes (instead of plotting 200 trajectories).  Your comment that the dataset is sensitive due to the fact that there are only 3 validation and 4 testing trajectories is well-taken. We have updated our claim to: “Looking to the table, we see that our method performs similarly to other state-of-the-art methods.” We converted to RMSE due to concerns that MSE might exaggerate performances differences between approaches on a quick reading. \n\n**I would appreciate some remarks on limitations concerning the recognition module, i.e., the encoder. As stated in Section D, data sparsity is only accepted to a certain extent and cf. (l 495f.)**\n\nRegarding the encoder design when we say, \"It is worth noting that without this deep-kernel our approach struggled to achieve good validation accuracy on the datasets we considered,\" this was intended to contrast the encoder with and without a deep kernel specifically.  \n"}, {"review_id": "jGGC7d6lr1", "reviewer": "Reviewer_oJgT", "summary": "This paper proposes a new inference method for latent SDE model. Due to the continuous nature of SDE, inference methods for latent SDE models are usually expensive and the cost would grow as the dataset goes larger or longer. Existing methods also would require differential equation solvers, which makes them unfriendly to GPU parallelization. The approach proposed in this work uses an amortized Markov Gaussian process as a variational posterior for inference over the latent SDE. The proposed variational objective allows for factorization over time steps and therefore enables alleviates the need to iterate over all time steps. Under such a framework, the cost at each iteration is only dependent on the number of Monte Carlo samples utilized for gradient estimation. The authors then empirically evaluate the efficiency of their method in a wide range of tasks, in which the proposed method shows much better performance in terms of computational cost, numerical stability, and convergence compared with existing methods.", "questions": "I am not familiar with the literature of latent SDE but I wonder whether the proposed variational posterior is more biased than existing approaches. Assuming computational cost is not an issue, would the baseline approaches be more favorable to the proposed method?", "limitations": "Nothing I can think of other than the limitation discussed by the authors at line 182.", "rating": 7, "confidence": 2, "soundness": 3, "presentation": 3, "contribution": 4, "strengths": "- The paper is well-motivated and the proposed approach, to the best of my knowledge, is very sensible and well supported by empirical evaluations.\n\n- The authors provide a thorough discussion of existing methods and include them in the experiments as baseline methods.\n\n- The paper is well written and the proposed method is presented clearly.", "weaknesses": "Some technical details seem to be discussed too briefly in the main text, for example, the implementation of the encoder and the use of deep kernel for \"interpolate between latent vectors\". I can hardly understand the meaning without referring to the appendix. In Remark 2, the authors mention \"make use of a nested Monte Carlo scheme along with stratified sampling to reduce the number of decoder evaluations and the variance of gradients.\" but the meaning of stratified sampling and nested Monte Carlo in this context is not discussed.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your comments. We appreciate that you found our approach to be well-motivated, sensible\nand well-supported by empirical evaluations. Responses to your specific questions and\nconcerns are provided below.\n\n**Some technical details seem to be discussed too briefly in the main text, for example, the implementation of the encoder and the use of deep kernel for “interpolate between latent vectors”. I can hardly understand the meaning without referring to the appendix. In Remark 2, the authors mention “make use of a nested Monte Carlo scheme along with stratified sampling to reduce the number of decoder evaluations and the variance of gradients.” But the meaning of stratified sampling and nested Monte Carlo in this context is not discussed.**\n\nThe point that some technical details are discussed too briefly in the main text \nis well-taken. The revisions which address this concern are provided \nbelow along with the line number where this new section would appear in the original text.\n\n[L141-145] “To reiterate, the probabilistic encoder is a function which takes in $M$\nobservations from a particular partition along with a time stamp, $t$,\nand outputs a mean and covariance as an estimate for the latent state\nat that particular time.\nIn principle, any function which can transform a batch of snapshots and a time\nstamp into a mean and covariance could be used as an encoder in our work.\nIn our implementation, we use deep neural networks to encode $x_i^{(j)}$ for\n$i\\\\in\\\\mathcal{I}$ where $\\\\mathcal{I}$ contains some temporal neighbours\nof $x_i$ into a single latent vector.\nThis approach yields a set of latent vectors associated with each observation\nin the partition $h_i$ for $i=1,2,\\\\dots, M$.\nWe then interpolate between each latent vector using a deep kernel based\narchitecture to construct the posterior approximation for any time stamp\nin the partition; see Appendix~D for details.”\n\n[L167-173] “In the case that \nevaluations of the SDE drift term were relatively \ncheap compared to decoder evaluations (for example in the case the dimension\nof the latent state is much smaller than the dimension of the data), \nwe found it useful to increase the number of samples used to approximate the \nintegral over time without increasing the number of samples from the \nvariational posterior.\nTo do so we made use of a nested Monte Carlo scheme\nto approximate the second term in the ELBO,\n\n$$\n( t_1^{(j+1)} - t_1^{(j)} )E_{p(\\\\epsilon)p(t)}||r_{\\\\theta,\\\\phi}(T(t,\\\\epsilon, \\\\phi),t)||^2_{C_\\\\theta(t)}\n$$\n$$\\\\quad \\\\approx \\\\frac{( t_1^{(j+1)} - t_1^{(j)} )}{RS} \\\\sum\\_{k=1}^R \\\\sum\\_{l=1}^S  ||r_{\\\\theta,\\\\phi}(T(t^{(k,l)},\\\\epsilon^{(k)}, \\\\phi),t^{(k, l)})||^2_{C_\\\\theta(t^{(k, l)})}\n$$\n\nwhere, again, each $\\\\epsilon^{(k)} \\\\sim \\\\mathcal{N}(0, I)$ and each \n$t^{(k,1)}, t^{(k, 2)}, \\\\dots, t^{(k, S)}\\\\sim \\\\mathcal{U}(t_1^{(j)}, t_1^{(j+1)})$.\nIn addition, because the integral over time is one-dimensional we used \nstratified sampling to draw from $\\\\mathcal{U}(t_1^{(j)}, t_1^{(j+1)})$\nin order to further reduce the variance in the integral over time.”\n\n**I wonder whether the proposed variational posterior is more biased than existing approaches. Assuming computational cost is not an issue, would the baseline approaches be more favorable to the proposed method?**\n\nThank you for this question. A summary of the assumptions made by our approach as compared reference [13] is provided on lines 175-194.\nTo summarize, the assumptions we make with respect \nto this approach are that (1) the posterior over the latent state can be approximated by \na Gaussian process and (2) the diffusion matrix is strictly time-dependent and not \na function of the state.\n\nRegarding the first assumption, it is important to emphasize that this assumption \ndoes not require that the *prior* (or generative model) define a Gaussian process, it only \nrequires that the *posterior* (i.e. the state given a batch of observations) be well \napproximated by a Gaussian process.\nThis is an assumption that is commonly used in the state estimation literature for \na wide variety of nonlinear state estimation tasks [A,B]. Because we are seeking \nto infer a latent state using more observations than states, we believe that \nthis assumption does not limit our approach as compared to [13].\n\nWith regards to the second assumption, it is true that our approach is strictly \nmore biased than [13]. Our approach is not capable of learning latent SDEs \nwith state dependent diffusion processes [L189]. \n\nAll this is to say that our approach makes more assumptions than [13]. \nHowever, we would like to emphasize that our approach is significantly more efficient than the approach of [13]. This enables the application of our approach to complex, high-dimensional systems on a limited computational budget. Even in situations where there are no strict limits on the computational budget, the computational efficiency of our approach would allow for more detailed hyperparameter tuning studies to maximize performance. \n\nIn order to help ensure that readers don’t leave with the same questions, we have provided the following summary of the assumptions made by our approach in the updated draft.\n\n[L175] “**Summary of Assumptions.**\nIn the previous sections we introduced an ELBO which, when maximized,\nleaves us with a generative model in the form of a nonlinear latent SDE \nwith time-dependent diffusion and an approximation to the latent \nstate in the form of a Gaussian process.\nTo reiterate, we only assume that the approximating posterior, i.e. the \ndistribution over the latent state given a batch of observations, is a Gaussian \nprocess; this is an assumption that is commonly made in the context \nof nonlinear state estimation, for example [A, B].\nWhen making predictions, we sample from the nonlinear SDE which characterizes \nthe generative model.”"}, {"review_id": "O0kXjtXpWe", "reviewer": "Reviewer_6Z23", "summary": "\nThis paper introduces a method for identifying latent stochastic models from discrete time observations.\nThe authors perform\nunbiased approximations to gradients of the\nevidence lower bound to identify parameters and estimate the state for latent stochastic differential equations. \nThy propose to combine a recently introduced (by the same authors - provided anonymous paper) reparametrisation of the ELBO\nin terms of time dependent linear SDEs and amortisation to reduce the computational demands.\nThe resulting approach does not require numerical integration of the approximate SDE or the ODEs defining the evolution of the central moments (as required by previous approaches), and seems to perform on par or better than existing frameworks. \n\nAlthough the method is limited by their assumptions of linear covariances and additive noise, the merging of the amortisation with the reparametrisation of the ELBO seems a rather interesting contribution. The performance of the method on the numerical experiments seem to be on par with existing approaches. However in most numerical examples the employed noise in the system was negligible when considering the dynamic range of the employed systems.\n", "questions": "- How do you select the partitions of the temporal axis and how does this selection influence the obtained results?\n\n- In line 498 in the supplement you mention that the encoder design is interpretable. Can you explain what you mean with this?\n\n- Could you please explicitly mention in the appendix what observation model was employed in each of the numerical experiments?\n\n- The time complexity of the algorithm depends on the number of evaluations $R$. Can you provide a systematic evaluation of how the number of evaluations $R$ influences the performance of the framework?\n\n- What kind of issues do emerge in your framework if you don't split the observations into chunks?\n\n- How does the method perform for decreasing sampling frequency?\n\n- Does the method require the entire state of the system to be accessible to the observation model, or could also a lower dimensional projection of the latent state work?", "limitations": "\n- The proposed framework is limited only to systems with diagonal covariances.\n\n- As far as I understand, they employed only a Gaussian observation model in all experiments.\n\n---\n\n[1] Duncker, Lea, et al. \"Learning interpretable continuous-time models of latent stochastic dynamical systems.\" International Conference on Machine Learning. PMLR, 2019.\n\n", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- The authors eliminate the need for numerically integrating the approximate differential equations by replacing the initial value problem with an integral resulting in reduced computational demands.\n\n- The approach avoids numerical instabilities often encountered in adjoint methods.\n\n- the proposed solution has a time and memory cost that depends only on the number of parallel evaluations of the ELBO, i.e., scales independently with the amount of data, the length of the time series, and the stiffness of the approximation to the differential equations.", "weaknesses": "\n- The approximate process is a linear SDE, and thus cannot probably capture effectively nonlinear dynamics that may result in multimodal transition densities.\n\n\n- To my understanding the proposed framework requires a diagonal covariance matrix (but I am happy to be corrected if this is not the case) for the diffusion process, as well as a diagonal time dependent covariance for the approximate process. How would you tackle problems with interactions in the diffusion component with your framework?\n\n\n\n- I consider the diffusion term for the experiment of the Lorenz system substantially small ($\\sigma = 10^{-5}$) considering the state space volume the state of the system spans.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your thoughtful comments. Our response to specific questions and comments  are provided below. One important clarification is that we only make the assumption \nthat the approximate posterior (i.e. the latent state given the observations) is Gaussian.\nWe only use this approximate posterior to encode the observations to the latent space.\nThe generative model is not required to be a Gaussian process. In fact, the generative model that we use in this work is a nonlinear SDE with additive noise, which makes our approach fairly flexible and capable of dealing with complex, real-world dynamical systems [L67-L74]. \n\n**The approximate process is a linear SDE, and thus cannot probably capture effectively nonlinear dynamics that may result in multimodal transition densities.**\n\nBecause the generative model is a nonlinear SDE, predictions will result in \nmultimodal transition densities [L184-L186]. \nTo help ensure that readers don’t leave with the same confusion, we have provided the following paragraph in revised main text. \n\n[L175] “**Summary of Assumptions.**\nIn the previous sections we introduced an ELBO which, when maximized,\nleaves us with a generative model in the form of a nonlinear latent SDE \nwith time-dependent diffusion and an approximation to the latent \nstate in the form of a Gaussian process.\nTo reiterate, we only assume that the approximating posterior, i.e. the \ndistribution over the latent state given a batch of observations, is a Gaussian \nprocess; this is an assumption that is commonly made in the context \nof nonlinear state estimation, for example [A,B].\nWhen making predictions, we sample from the nonlinear SDE which characterizes \nthe generative model.”\n\n**To my understanding the proposed framework requires a diagonal covariance matrix ...?**\n\nThe framework, in general, does not require the covariance matrix be diagonal.\nThe diffusion process and the covariance are required to be diagonal in order to scale \nas $O(d)$ where $d$ is the dimension of the latent state.  \nIf we are willing to give up some computational efficiency, we could use a full covariance \nmatrix and full diffusion matrix ($O(d^3)$ time cost).\n\n**I consider the diffusion term for the experiment of the Lorenz system substantially small considering the state space volume the state of the system spans.**\n\nWe agree that the diffusion term is small. The purpose of this section was to \ndemonstrate the numerical instability of adjoint methods for long time series and \nto show that our approach does not suffer from these same instabilities.\n\n**How do you select the partitions of the temporal axis ...?**\n\nIt is challenging to systematically select the partitions of the temporal axis. The \nclosest analogy to this hyperparameter in the standard supervised learning setting is the batch size; however, this isn’t the full story.\n\nIntuitively, there is a trade-off between the complexity of the encoder and having enough observations \nin each partition such that the assumption of a Gaussian process for the latent state \nis reasonable. For example, given two observations in each partition, the encoder only needs \nto interpolate between the latent states over a short time window (meaning the dependence \non $t$ could be relatively simple). \nHowever, using only \ntwo observations might make it challenging to approximate the latent state using \na Gaussian process (think of the pendulum problem, it is difficult to estimate the \nvelocity using only a few frames).\nFurthermore, like the batch size in the standard supervised learning setting, the parameter \n$M$ has an effect on the variance of the gradients.\n\nFor our work, we maximized the number of elements in each partition \nsuch that fast approximations of the gradients of the ELBO can be calculated on the hardware available to us.\n\n**Could you please explicitly mention in the appendix what observation model was employed ...?**\n\nThank you for pointing this out – we have added an explicit description \nof the likelihood for each problem in the appendix. For your reference, \nwe assumed a Gaussian likelihood for 4.1, 4.2, & 4.3 and a Bernoulli likelihood for 4.4 (because\nthe frames where black and white).\n\n**Can you provide a systematic evaluation of how the number of evaluations \ninfluences the performance of the framework?**\n\nThank you for this question. First, it is worth mentioning that even for one evaluation of the model we are left with an unbiased estimate for the gradient of the ELBO. Increasing the number of evaluations only serves to decrease the variance of gradient estimates.\n\nWe have added an additional numerical study to the supplementary material to explore this effect. For the problem we considered, we found that convergence with $S=10$, $S=50$, and $S=100$ was similar. The numerical study is described in greater detail in the “General Response” section. \n\n**What kind of issues do emerge in your framework if you don’t split the observations into chunks?**\n\nWe provide a brief discussion in the manuscript~[L125-128]. The main issue with \nnot splitting the observations into chunks is that we need to compute and store \nan approximation to the latent state over the entire time series. \nIn addition, splitting the observations into chunks also allows us to efficiently approximate the ELBO using a batch of datapoints rather than the entire dataset.\n\n**How does the method perform for decreasing sampling frequency?**\n\nThis will be highly problem dependent. From a theoretical perspective our \napproach should perform similarly to other approaches for inferring latent SDEs \nas the sampling frequency is decreased.\n\n**Does the method require the entire state of the system to be accessible to the observation model...?**\n\nThe method should be able to approximate the data so long as the data can be approximated\nby a nonlinear SDE with additive diffusion. Regarding reconstructing hidden states, this \nis a challenging inverse problem beyond the scope of the current work.\n"}, {"review_id": "Ro4qPSojLj", "reviewer": "Reviewer_ayS4", "summary": "The authors develop a reparameterization scheme for affine SDEs which is advantageous in memory cost and time. Specifically, the authors consider splitting the computation into a series of different chunks such that each can be effectively parallelized and integrated individually by taking the expectation with respect to a uniform random variable over the time interval. This is possible since the SDE is linear and a mean and covariance matrix provide the full description of the data at all time marginals and individual trajectories are not needed. The points are then combined using an interpolation technique that is based upon deep kernels. The authors finally consider experiments to highlight the utility of the method. They consider synthetic and a real experiment on time series data that compares the number of function evaluations and the accuracy of the methods.", "questions": "How does the interpolation between the windows affect the quality of the sample path? This point was not discussed thoroughly and it seems like it might have impact on the sample quality.\n\nWere there any comparisons to non-adjoint based back-propagation through the solver? For example, using just an Euler solver, how does the method compare with backpropagation through that? \n\nLine 147: “While the previous sections have demonstrated how to replace a differential equation with an integral …” Can the authors please expand upon this? When solving a differential equation, the integral is approximated as a sum, as the authors have done. Is this comment due to the fact that you are not using an adaptive solver but you are using an expectation to compute the integral? Or is it more that you do not require the value of $z(t)$ at some later time, only the expectation?\n\nIs there an ablation on changing the window size/number of windows? How does this parameter affect the performance?\n\nWhy is the drift parameterized as a neural network of z on line 519? I thought the drift was linear in z.\n\nIs there a way to verify that the linear SDE assumption is a valid one? Are there some equivalence classes one can derive for the case studied here and a possibly more general class of SDEs?", "limitations": "The authors provide a discussion on the issues associated with the linear SDE assumption.", "rating": 5, "confidence": 2, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "The idea of transforming the sequential computations into parallelizable expectations is appealing and the numerical results suggest this approach provides a significant computational speed up. This is a useful extension beyond the usual sequential nature of ODE which cannot be parallelized and could be particularly useful and impactful in certain cases. Additionally, empirical results suggest that the modifications provide increased stability. The authors also show the utility of the method on a real dataset where the performance is superior compared to most of the baselines. Finally, the paper is fairly well written and the presentation is intuitive.", "weaknesses": "The method only applies to SDEs where the marginal descriptions are known, as the authors make an assumption that the evolution of the latent space is given by a Markov Gaussian process. However, this is not the case for most SDEs, which makes the algorithm limited in scope. From the paper, it is not clear where such an assumption would hold that the latent distribution should be Gaussian everywhere. This question was not studied, and I think that's the major weakness of the paper. While the computations are easier to perform, it is mainly due to the simpler model. SDEs are particularly useful when sampling from a complicated distribution, but in this case the latent distribution Gaussian for all time marginals. Maybe to strengthen this, it would be helpful to show the types of stochastic processes in the observable space that can be modeled (e.g. by applying Ito's lemma to the latent process). This would the provide a more concrete statement on the expressiveness of the model.\n\nThe experimental evaluation is somewhat limited since the evaluation is only on two synthetic datasets and one real dataset. I would be curious to see how the other methods compare on the pendulum dataset. Additionally, I think some more synthetic datasets would be useful to describe the utility/use of the method, particularly on non-affine SDEs and seeing the performance in approximation. It would be helpful to see a comparison to backpropagating through the solver rather than the adjoint method, which the authors mention is numerically unstable. Understanding how the behavior of backpropagation through the solver (since this is possibly the most common method when not using a fine step size or memory constraints are not a concern) compares to the proposed method. Finally, the authors mention in line 184 the linear assumption of the SDE is not a limiting factor, but the experiments do not seem to be extensive enough to validate such a claim. \n\nOther types of methods that are related and should be mentioned are [1, 2]. These methods consider similar strategies with [1] considering a gaussian process in the latent space and [2] considering a linear SDE approximating the latent evolution (using some of the same results mentioned in this paper based on casting linear SDEs in terms of their moments).\n\n[1] Fortuin et al 2020 AISTATS, GP-VAE: Deep Probabilistic Time Series Imputation\n[2] Duncker et al 2019 ICML, Learning Interpretable continuous-time models of latent stochastic dynamical systems", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "**The method only applies to SDEs where the marginal descriptions are known, \nas the authors make an assumption that the evolution of the latent space \nis given by a Markov Gaussian process.**\n\nThank you for your comments. Before addressing specific comments, we would like to \nmake an important clarification.  We agree that SDEs are useful when sampling from a \ncomplicated distribution.\nIn the present work, we only make the assumption that the posterior (i.e. the latent state *given* \na batch of observations) can be approximated by a Gaussian process. \nThe *generative model* we use in our work is characterized by a nonlinear SDE with time dependent diffusion meaning that\npredictions from the model will be a complicated non-Gaussian in general [L67-L74].\nIn order to help ensure that readers don’t leave with the same confusions, we have added the following clarifying section to the updated version of the manuscript.\n\n[L175] “**Summary of Assumptions.**\nIn the previous sections we introduced an ELBO which, when maximized,\nleaves us with a generative model in the form of a nonlinear latent SDE \nwith time-dependent diffusion and an approximation to the latent \nstate in the form of a Gaussian process.\nTo reiterate, we only assume that the approximating posterior, i.e. the \ndistribution over the latent state given a batch of observations, is a Gaussian \nprocess; this is an assumption that is commonly made in the context \nof nonlinear state estimation, for example [A,B].\nWhen making predictions, we sample from the nonlinear SDE which characterizes \nthe generative model.”\n\n**Other types of methods that are related and should be mentioned are [1, 2].**\n\nThank you for suggesting the papers from Fortuin et al and Duncker et al. We have added both to \nour list of references. \nIn contrast to [1], in our work the generative model is \ndefined by a nonlinear SDE. \nRegarding [2], the major difference between this approach and our work is that they still \nrely on solving ODEs to estimate gradients (see Section 4.1 in their paper). The main consequence of Theorem 1 is that we can eliminate the differential equality constraints required by this approach.\nAs we have argued throughout \nour work, solving differential equations as a part of an iterative optimization method \nis computationally challenging [L26 – L30]\n\n**How does the interpolation between the windows affect the quality of the sample path?**\n\nApologies if we have misunderstood your question.\nThe required data sampling frequency will be highly problem dependent. From a theoretical perspective our approach should perform similarly to other approaches for inferring latent SDEs \nas the sampling frequency is decreased.\n\n**Were there any comparisons to non-adjoint based back-propagation through the solver? \nFor example, using just an Euler solver, how does the method compare with backpropagation through that?**\n\nWhen backpropagating through a forward solver, we only expect to save evaluations of the model in estimating the adjoint (not in estimating the solution trajectory). This means that all of the challenges associated with methods for solving initial value problems as a part of an iterative optimization procedure remain [L26 – L30]. \nIn addition, it is worth noting that\ngradients computed via backpropagation of a forward solver are not consistent with the adjoint ODE in general [C]. \nIn our work we infer a latent SDE using unbiased gradients of a lower bound on the evidence.\nIt is worth reiterating that in example 4.1 we considered training a neural ODE with an adaptive stepping tolerance of $10^{-2}$, $10^{-4}$, and $10^{-6}$. When decreasing the tolerance we find that the standard NODEs struggled to achieve a comparable validation accuracy to our approach.\n\n**“While the previous sections have demonstrated how to replace a differential equation with \nan integral …” Can the authors please expand upon this? When solving a differential equation, \nthe integral is approximated as a sum, as the authors have done. \nIs this comment due to the fact that you are not using an adaptive solver but you are \nusing an expectation to compute the integral? Or is it more that you do not require the value \nof at some later time, only the expectation?**\n\nCasting an IVP problem in the form $x(t) = x(0) + \\int_0^t f(x(s),s) ds$ and invoking appropriate approximations for the integrand results in a sequential time-marching scheme. In Line 147, we are drawing the reader’s attention to the fact that equation (6) is a standard integral with respect to time that can be evaluated in parallel. This representation underpins the outstanding computational advantages offered by our approach. \n\n**”Is there an ablation on changing the window size/number of windows? How does this parameter affect the performance?”**\n\nIt is challenging to systematically select the partitions of the temporal axis. The \nclosest analogy to this hyperparameter in the standard supervised learning setting is the batch size; however, this isn’t the full story.\n\nIntuitively, there is a trade-off between the complexity of the encoder and having enough observations \nin each partition such that the assumption of a Gaussian process for the latent state \nis reasonable. For example, given two observations in each partition, the encoder only needs \nto interpolate between the latent states over a short time window (meaning the dependence \non $t$ could be relatively simple). \nHowever, using only \ntwo observations might make it challenging to approximate the latent state using \na Gaussian process (think of the pendulum problem, it is difficult to estimate the \nvelocity using only a few frames).\nFurthermore, like the batch size in the standard supervised learning setting, the parameter \n$M$ has an effect on the variance of the gradients.\n\nFor our work, we maximized the number of elements in each partition \nsuch that fast approximations of the gradients of the ELBO can be calculated on the hardware available to us."}]

Author rebuttal:
## General Comment

We would like to open by thanking the reviewers for their careful consideration of our work.
We understand that providing quality reviews is time-consuming, so we are thankful for your efforts.
Here we summarize the main changes made to the paper.
We address comments and questions from referees in separate comments below.

**Updates to the main text**

1. We recognize that some descriptions of technical details were too terse in the main text. We have used what space we have left to clarify details on the encoder, the Monte-Carlo based approximation to the evidence lower bound, and the assumptions made by our approach. Please see the detailed responses for clarifications on these updates.

2. We have edited the paper to ensure that the details of the numerical study design are described more clearly. In particular, we explain that section 4.1 shows the validation accuracy curves vs. the total number of function evaluations when training on the same dataset using 10 random seeds for each approach.

3. We have softened some claims following suggestions from reviewers. In particular, for section 4.3 we now claim that “our method performs similarly to other state-of-the-art methods.” In section 4.1 we now claim that our approach “requires more than one order of magnitude fewer evaluations of the model (NFE) than the neuralODE”. We believe these adjusted claims more precisely present our results without affecting the main message. 

3. We have also improved performance for the numerical study listed in section 4.1 after additional hyparameter tuning. The updated Figure 2 is provided in the attachment.

We have added the following references into the main text in order to support these new sections.

[A] Barfoot, TD, Forbes, JR, Yoon, DJ. “Exactly sparse Gaussian variational inference with applications to derivative-free batch nonlinear state estimation.” In: The Internationl Journal of Robotics Research 39(13), 1473 – 1502 (2020). 

[B] Barfoot, TD. “State Estimation for Robotics a Matrix Lie Group Approach. Cambridge University Press, UK. (2017)

[C] Alexe, M, Sandu, A. “On the discrete adjoints of adaptive time stepping algorithms”. In: Journal of Computational and Applied Mathematics 233, 1005-1020, 2009.

**Additional numerical study**

In response to reviewer questions regarding the effect of parameters in the Monte-Carlo approximation to the gradient (Eq. 9), we will provide a new numerical study in the supplementary material. In this numerical study we trained a latent SDE on data from a four-dimensional predator prey model. Keeping all other hyperparameters the same ($R=1$, $M=256$) we varied the number of nested Monte-Carlo samples $S=(10, 50, 100)$. The results are summarized in Figure 10 in the attachment. The left figure shows the validation RMSE vs. the iteration count. From this figure we see the model is insensitive to the choice of $S$ in terms of the validation RMSE for this problem. The right figure shows the validation RMSE vs. the total number of function evaluations of the model. We see that the while different values $S$ do not seem to impact validation RMSE at convergence, this hyperparameter can have a large impact on the total number of function evaluations. 

Meta-review: I am happy to convey my recommendation for the acceptance of your paper  to be presented at NeurIPS. 

This paper was very much on the line and could have gone either way. The work is interesting, important, and would have gotten accepted at some point. I sincerely hope that you, the authors, really actually do at least the edits you seem already to have done.  Based on this trust and the discussion I think the paper is worthy of publication at NeurIPS this year, however there is a risk to the authors and the work itself if the feedback from the reviewers is not carefully considered in the final, camera-ready version of the paper.

NeurIPS is a prestigious platform that attracts the brightest minds in the field, and your paper's acceptance adds to the conference's reputation for excellence. I am look forward to seeing your work presented and discussed among peers who share your passion for pushing the boundaries of knowledge.
Decision: Accept (poster)