Scalable Inference in SDEs by Direct Matching of the Fokker–Planck–Kolmogorov Equation
Year: 2021

Abstract:
Simulation-based techniques such as variants of stochastic Runge–Kutta are the de facto approach for inference with stochastic differential equations (SDEs) in machine learning. These methods are general-purpose and used with parametric and non-parametric models, and neural SDEs. Stochastic Runge–Kutta relies on the use of sampling schemes that can be inefficient in high dimensions. We address this issue by revisiting the classical SDE literature and derive direct approximations to the (typically intractable) Fokker–Planck–Kolmogorov equation by matching moments. We show how this workflow is fast, scales to high-dimensional latent spaces, and is applicable to scarce-data applications, where a non-parametric SDE with a driving Gaussian process velocity field specifies the model.

text: 1 Introduction
Differential equations are the standard method of modelling change over time. In deterministic systems the dynamics specifying how the system evolves, are typically written in the form of an ordinary differential equation (ODE). The dynamics act as prior knowledge and often stem from first-principles in application areas such as physics, control engineering, chemistry, or compartmental models in epidemiology and pharmacokinetics. Recently, learning ODE dynamics with modern automatic differentiation packages in machine learning has awakened an interest in black-box learning of continuous-time dynamics (e.g., [6, 37]) and enabled their more general use across time-series modelling applications.
A stochastic differential equation (SDE, [30, 40]) can be seen as a generalization of ODEs to stochastic dynamical settings, where the driving forces fluctuate or are uncertain. Stochastic dynamics appear naturally in applications where small (and typically unobserved) forces interact with the process, such as tracking applications, molecule motion, gene modelling, or stock markets. In machine learning, SDE models have received wide-spread attention due to their robustness and appealing properties for uncertainty quantification.
The concept of a ‘solution’ to an SDE is broader than that of an ODE. As the process is stochastic, the full solution entails a probability distribution, p(z, t), depending on time t and covering the space z (see, e.g., [36]). For Itô type SDEs, the evolution of the probability mass can be described in terms of the Fokker–Planck–Kolmogorov (FPK) partial differential equation (backward Kolmogorov equation). This equation is typically intractable, and instead the de facto approach for inference in SDEs in machine learning is sampling. The most common approaches in this space are based on
35th Conference on Neural Information Processing Systems (NeurIPS 2021).
Stochastic Runge–Kutta schemes (such as the Euler–Maruyama scheme) which are derived from the Itô–Taylor series. These schemes sample realization trajectories of the SDE by driving the dynamics with numerical simulation of Brownian motion. However, these schemes suffer from drawbacks both related to (ordinary) Runge–Kutta methods—such as step-size and sensitivity to stiffness—as well as problems associated with any sampling schemes, such as a high number of samples required for an accurate representation of the underlying distribution.
Despite these problems, few contemporary SDE approaches in machine learning explore SDE solutions beyond stochastic Runge–Kutta (or even the Euler–Maruyama scheme). Our aim is to try to broaden this view, and in Fig. 1 we sketch an example where we show three alternative solution perspectives to a Gaussian process prior SDE model (GP-SDE): the FPK probability density field, Euler–Maruyama samples, and a Gaussian assumed density approximation. We argue that a Gaussian approximation in latent space SDEs is reasonable, as Gaussian approximations are typically employed anyway in observation models, and allow for speeding-up learning by an order of magnitude.
The contributions of this paper are as follows. (i) We go through the workflow connecting ‘random ODE’ models with Itô SDEs driven by a Gaussian process prior over the velocity field, which allows for convenient specification of prior knowledge on the vector field and induces an implicit prior over the SDE trajectories; (ii) We revisit the classical SDE literature and derive direct approximations to the (typically intractable) Fokker–Planck–Kolmogorov equation in an assumed density Gaussian form that avoids sampling-based inference in the latent space, which makes inference fast and does not require sampling a high number of trajectories; (iii) We show how this workflow is fast, applicable to scarce-data applications, and how it also extends to previously presented latent SDE models.

heading: 1.1 Related Work
text: Neural ODEs [6] model ODE dynamics by a neural network. Such models were developed further in [37], where the encoder is an ODE-RNN that improves modelling of irregularly sampled time series. A latent Bayesian neural ODE model, ODE2VAE, was examined in Yıldız et al. [49], where an encoder is combined with an ODE model whose second order dynamics are given by a Bayesian neural network. The neural ODE paradigm of modelling latent dynamics has been expanded to neural SDEs [26, 12, 46, 19], where the typical workflow is that a variational autoencoder (VAE, [21, 35]) is combined with a latent neural SDE, whose drift and diffusion are modelled by neural networks. In addition to modelling time series, neural SDEs have been used in generative models [16, 43, 44], where the generation of images from noise is modelled as the reverse-time process of a diffusion SDE by using Langevin dynamics on score-based models. Continuous normalizing flows are another model family, which applies ODE dynamics in a generative model [6, 11].
These works leverage simulation/sampling for solving the SDE in the latent space. The model can be trained by a stochastic adjoint method [26, 18]. More recently, latent neural SDEs have been trained deterministically by moment matching [27]. However, they discretized the system before matching the moments, while we form a direct approximation to the solution of the FPK. Compared to optimizing the moment ODEs, as discussed in [27], by maximizing likelihood, we regularize during inference by Gaussian process priors, or prior stochastic processes as in [26]. Approximative solutions to non-linear SDEs have been applied earlier in filtering theory, where the optimal filter is approximated by a Gaussian assumed density filter [23]. In [39], such approximations are used for continuous-discrete state-space modelling. An alternative to assumed density filters are local linearization methods [31, 41], and simulation-based Itô–Taylor series solutions, stochastic Runge–Kutta methods, and leapfrog methods such as Verlet for second-order SDEs (see [22, 40]). The approximations presented in this work are also related to GP approximations [4, 2] of SDEs. The linearization approximations are related to statistical linearization [10, 42], and variational approximations [4].
Orthogonally to the SDE inference, we also consider SDE model specification in terms of GP priors. The seminal work by Ruttor et al. [38] considered GP-SDE models with unit diffusion. Yildiz et al. [48] built a model, where the drift and diffusion are sparse Gaussian processes with time-independent kernels. In Hegde et al. [13], a spatio-temporal SDE with GP priors for the drift and diffusion was combined with a GP as a continuous version of deep Gaussian processes. State-space models with a GP latent state transition function [8] train a non-parametric latent process to approximate unobserved dynamics. These are related to hierarchical GP dynamics [45, 24], where prior knowledge of the system can be encoded in multi-level hierarchies, for modelling, e.g., walking dynamics.

heading: 2 Methods
text: We draft the methodology based on the latent dynamics workflow components as presented in Fig. 2. The focus is first on specifying models for the latent space dynamics, starting from implicit priors in terms of random ODEs which we frame as Itô SDEs. Thereafter, the focus shifts from model specification to inference, where we show that the Fokker–Planck–Kolmogorov equation can be efficiently approximated in an assumed density form, and finally brings us to cover the likelihood structure of these models as well.

heading: 2.1 Random Field Ordinary Differential Equations as SDEs
text: The continuous dynamics of a latent (unobserved) z(t) ∈ Rd can conveniently be written in the form of a general first-order ordinary differential equation
d dt z(t) = vθ(z(t), t), (1)
where vθ(·) : Rd × R+ → Rd denotes the velocity field parametrized by θ. This is a general form of a non-linear ODE system, where the dynamics are deterministic and fully characterized by vθ(·). The methodology presented in this section directly extends to the case where the vector field vθ is time-dependent, but we omit the time dimension for simplicity of notation. Previously, the implicit prior on z(t) over t specified by Eq. (1) has been generalized to stochastic models by considering vθ(·) to be stochastic. These models are known as ‘random’ ODE models, and the random field vθ(·) is typically either characterized by a Gaussian random field or Gaussian process model (see, e.g., [38, 13]) or some parametric model (e.g., [26]).
We consider an unconventional ODE model (or actually no ODE model at all, to be precise), where we specify a GP prior [34] over the velocity field in form of a multi-output Gaussian process prior:
v(z, t) ∼ GP(µ(z),κ(z, z′)), (2)
where µ : Rd → Rd is a mean function and κ : Rd×Rd → Rd×d is a matrix-valued covariance function. The Gaussian process prior is completely specified by its mean and covariance function, which encapsulate the assumptions about the sample processes/fields v (such as continuity, differentiability, curl, divergence, etc.): µ(z) := E[v(z)] and κ(z, z′) := E[(v(z)−µ(z))(v(z′)−µ(z′))∗]. In Fig. 3, we will consider examples of useful vector-valued covariance functions that encode properties on the vector field. For inference, the GP is conditioned on input–output pair observations D = {(zi,∆zi)}ni=1 of the vector field, where ∆zi represents the observed derivative at zi. The conditioned vector field representation for an arbitrary point z∗ in the latent space can be given by
v(z∗) | D ∼ GP(E[v(z) | D],Cov[v(z) | D]), (3) where the E[·] and Cov[·] denote the marginal mean and (co)variance (for the multi-output GP, which means that the marginals are vector-valued). These take the form [34]: E[v(z∗)] = K∗K̂ −1y and
Cov[v(z∗)] = κ(z∗, z∗) − K∗K̂−1K⊤∗ , where y are the stacked observations of the derivatives such that y = (∆z⊤1 , . . . ,∆z ⊤ n )
⊤. The Gram matrix K corresponds to evaluations of the covariance function such that Kij = κ(zi, zj) are sub-blocks of K corresponding to the observation pairs
(i, j), and K̂ = K + γI, where γ is a nugget (observation noise/discrepancy) term, and K∗ is the cross-covariance between z and z∗. The prohibitive cubic computational scaling associated with GP models manifests in the inversion of K̂, and thus for large n, approximations based on inducing points or projections are used in practice to avoid this explicit inversion. In the light of Eq. (1) with v(z) ∼ GP(·, ·), the ODE is driven by a multi-dimensional Gaussian random field conditioned on D. A straightforward way of dealing with a model of this kind, is to do inference by sampling random draws of the velocity field from the GP, and then drive the ODE with those samples (can be viewed as an Monte Carlo approach for drawing ODE realizations).
However, a more convenient way is to specify the prior over the stochastic dynamics in a stochastic differential equation form. At its core, a lot of previous work in this space hinges on the realization that if everything is essentially Gaussian, an equivalent model can be specified in terms of an Itô SDE describing the stochastic evolution of trajectories affected by the GP velocity field [see 13, for discussion]. Informally, this takes the white noise form ddtz(t) = f(z) + L(z)w(t), where f(z) = E[v(z)] and L(z) denotes a square-root factor such that LL⊤ = Cov[v(z)] (in the scalar case, just the square-root, and in the multi-output case, e.g., the Cholesky factor). Here w(t) is a white noise process with unit spectral density. It is worth noting that we do not give guarantees for a direct link between the random ODE in Eq. (2) and the following SDE formulation (see App. A.3 for discussion). Yet, formally we write a similarly-behaving SDE in the standard Itô SDE form:
dz(t) = f(z, t) dt+ L(z, t) dβ(t), (4)
where dβ(t) is vector-valued unit Brownian motion (the spectral density Q is set to I). For a GP-SDE, the drift is driven by the GP mean, f(z, t) := E[v(z)] and the diffusion by the square-root factor of the marginal covariance at z(t), L(z, t) := √
Cov[v(z)]. To be precise, the GP-SDE drift and diffusion at a point z∗ are determined by the GP predicted mean and variance at z∗, which can be written as
f(z∗, t)=E[v(z∗)]=K∗K̂ −1y and L(z∗, t)=
√
Cov[v(z∗)]=
√
κ(z∗, z∗)−K∗K̂−1K⊤∗ . (5)

heading: 2.2 Fokker–Planck–Kolmogorov Equation
text: We are interested in solving SDE models of the form in Eq. (4), but without the restriction that the drift and diffusion are defined by a Gaussian process, and present the related theory with a model-agnostic view on the problem. Because the resulting solutions are stochastic processes, the full solution to the SDE can be characterized by its time-evolving probability density function. Let z(t0) ∼ p(z(t0)) be some initial condition which we assumed to be independent of the Brownian motion. The probability
density p(z, t) of the solution of the SDE in Eq. (4) solves the Fokker–Planck–Kolmogorov (FPK) partial differential equation (PDE):
∂p(z, t)
∂t = −
∑
i
∂
∂zi [fi(z, t) p(z, t)]+
1
2
∑
i,j
∂2
∂zi ∂zj
{ [L(z, t)QL⊤(z, t)]ij p(z, t) } . (6)
For a proof, see [40]. This PDE is also known as the Fokker–Planck equation (in physics) and the forward Kolmogorov equation (in stochastics). An appealing alternative form [40, Sec. 5.3] of the FPK equation can be given in terms of the following evolution equation with the adjoint operator A∗:
∂p ∂t = A∗ p, with A∗(•) = − ∑
i
∂
∂zi [fi(z, t) (•)] +
1
2
∑
i,j
∂2
∂zi ∂zj {[L(z, t)QL⊤(z, t)]ij (•)}.
(7) Eq. (7) allows for various kind of approaches for direct approximation of the FPK equation either by basis function approximations, finite differences, or other methods (Sec. 9.6 in [40] provides examples of using point collocation, Ritz–Galerkin, and FEM type of methods for approximating the solution). For example, the results in Fig. 1 and Fig. 4 are estimated by a grid discretization over z and solving the resulting (finite-dimensional) ODE corresponding to Eq. (7) by the matrix exponential: p(t) = exp(A(t − t0)). See App. A.2 for details. Even if the widely-used Euler–Maruyama, Milstein, and more general stochastic Runge–Kutta schemes are derived from the Itô–Taylor series, the resulting methods can still be viewed as an approximation of p(z, t).

heading: 2.3 Assumed Density Approximation of FPK
text: For the purpose of modelling latent space dynamics of systems of the kind in Fig. 2, we note that there the common practice of solving the latent SDE through costly simulation/sampling and then employing a variational (Gaussian) approximation in the encoder/decoder seems contradictory. That is, it might be unnecessary to sample realizations of the trajectory dynamics, if the interest is only in the time-marginals of the process. Thus, we seek to directly characterize the first two moments of the solution to the FPK equation in Sec. 2.2. We replace the FPK solution with a Gaussian approximation of form
p(z, t) ≈ N(z |m(t),P(t)), (8)
where m(t) and P(t) are interpreted as a mean and covariance of the state of the solution at time t. This kind of approximation is commonly referred to as a Gaussian assumed density approximation (see, e.g., [23, 39]), because the computations are done under the assumption that the state distribution is Gaussian. Assumed density approximations are common in signal processing–driven SDE methodology, and we refer the reader to Sec. 9.1 in [40] for a detailed overview. Following [39], we revisit the idea that a Gaussian process approximation to the SDE Eq. (4) can be obtained by integrating the following differential equations from the initial conditions m(t0) = E[z(t0)] and P(t0) = Cov[z(t0)] to the target time t:
dm
dt =
∫
f(z, t)N(z |m,P) dz and (9)
dP
dt =
∫
f(z, t) (z−m)⊤ N(z |m,P) dz
+
∫ (z−m) f⊤(z, t)N(z |m,P) dz+ ∫
L(z, t)QL⊤(z, t)N(z |m,P) dz. (10) These equations for the evolution of the first moments of the solution to the SDE can be interpreted as expectations over the drift and diffusion dynamics of the SDE, and can be derived from the FPK in Eq. (6). Conveniently, these expressions are not stochastic, but instead take the form of an ODE system that—given the integrals are tractable—can be solved with out-of-the-box ODE solvers. However, even if Eqs. (9) and (10) provide a generic Gaussian assumed density approximation framework for SDEs, an implementation of the method requires solving the following kind of d-dimensional Gaussian integrals:
EN[•] = ∫
[•] N(z |m,P) dz. (11) In the following sections we will consider two approaches (local linearization and moment matching with symmetric quadrature) which scale linearly in the number of latent dimensions d.

heading: 2.4 Linearizing the FPK Equation
text: Local linearization around the m (via a Taylor series approximation) is a classical approach widely used for this type of Gaussian integrals in machine learning and filtering theory [17, 29]. If the function f(z, t) is differentiable, the covariance differential equation can be simplified by using Stein’s lemma [32] such that
∫ f(z, t) (z−m)⊤ N(z |m,P) dz = [∫ Fz(z, t)N(z |m,P) dz ] P, (12)
where Fz(z, t) is the Jacobian of f(z, t) with respect to z. Linearizing around the mean m and approximating the diffusion as L(z, t) ≈ L(m, t) gives a linearized form of Eqs. (9) and (10):
dm
dt = f(m, t) and
dP
dt = PF⊤ z (m, t) + Fz(m, t)P+ L(m, t)QL ⊤(m, t), (13)
which provides a direct way of propagating the moments of the latent SDE through an ODE for the mean and covariance without the need of drawing multiple sample trajectories. The resulting ODE is (d+ d2)-dimensional, and only requires one evaluation of the drift, diffusion, and Jacobian per step.

heading: 2.5 Matching Moments of the FPK Equation
text: The local linearization approach given in the preceding section is efficient, but fully local. An alternative way of constructing an assumed density approximation to p(z, t) is to directly match the moments by solving the Gaussian integrals in Eqs. (9) and (10) by Gaussian quadrature methods. The approximation to Eq. (11) would take the form ∫
g(z, t)N(z |m,P) dz ≈ ∑i w(i) g(z(i), t), for an arbitrary integrand g(z, t), weights w(i), and so called sigma points z(i) = m+
√ P ξi. Here√
P denotes a square-root factor of P such as the Cholesky decomposition. The multi-dimensional Gaussian quadrature (or cubature, see [7]) rule is characterized by the evaluation points and their associated weights {(ξi, wi)}. We write Eqs. (9) and (10) in a Gaussian assumed density form which matches the moments by quadrature as follows [40]:
dm dt =
∑
i w (i) f(m+ √ P ξi, t) and (14)
dP dt =
∑
i w (i) f(m+ √ P ξi, t) ξ ⊤ i √ P ⊤
+ ∑ i w (i) √ P ξi f ⊤(m+ √ P ξi, t) + ∑ i w (i) L(m+ √ P ξi, t)QL ⊤(m+ √ P ξi, t). (15)
The computational complexity of this approach is highly dependent on the choice of quadrature method. A typical choice in ML applications would be Gauss–Hermite quadrature, which factorizes over the input dimensions leading to an exponential number (pd) of function evaluations/sigma points in the input dimensionality d for a desired order p. In order to guarantee scalability, we employ a symmetric 3rd order cubature rule [3] which similarly to Gauss–Hermite (p = 3) is exact for polynomials up to degree 3. The points are given by scaled unit coordinate vectors ei such that
ξi =
{ √ d ei, for i = 1, . . . , d,
− √ d ei, for i = d+ 1, . . . , 2d,
(16)
and the associated weights are wi = 1 2d . This approach provides a direct way of propagating the ‘true’ moments of the latent SDE through an ODE for the mean and covariance and without the need of drawing multiple sample trajectories. The resulting ODE is (d+ d2)-dimensional, and requires only 2d evaluations of the drift and diffusion per step.
2.6 Analysis of the Computational Complexity
In terms of the asymptotic computational complexity, the linearization approach in Sec. 2.4 requires O(1) evaluations of the drift, diffusion, and Jacobian per step. The moment matching approach in Sec. 2.5 requires O(d) evaluations of the drift and diffusion as well as an O(d3) Cholesky decomposition per step. The simplest Monte Carlo simulation method with p samples requires O(p) evaluations of the drift and diffusion per step. Additionally, the naïve requirement for p grows exponentially in d. On the other hand, the simulation approach is fully parallelizable over p, while the moment matching approach the number of nonparallelizable operations is O(d2) with the Cholesky decomposition being the bottleneck, and the linearization approach is nonparallelizable. While the linearization approach has the lowest number of function evaluations with respect to d, the cost of computing the Jacobian can be prohibitively large for arbitrarily complex models. Nevertheless, the Jacobian is available in closed-form for GPs and may be evaluated reasonably fast for neural network based drifts, see App. B.2 for empirical computational costs of evaluating the Jacobian when growing the network size. Thus we expect the FPK approximation schemes to always be beneficial in CPU-only cases (incl. CPU multi-threading and embedded devices). In multicore GPU use, for low-dimensional d, sampling remains appealing if GPU memory does not become a bottleneck.

heading: 3 Experiments
text: The goals of the experiments are three-fold: We first provide a study of the computational complexity. Then, we look into properties of the GP-SDE model from Sec. 2.1, where the experiments are concerned with showcasing model specification rather than inference. Finally, we consider two benchmark problems with high-dimensional inputs for learning a latent SDE model, where we test the performance of the approximations presented when the model is not defined by GPs, as the SDE methods presented in Sec. 2 are model-agnostic.
Timing Experiments To confirm the analysis in Sec. 2.6 and provide a practical insight, we run numerical experiments with the error of final marginal mean/covariance controlled. We use a highdimensional model of d independent Beneš SDEs (dz(t) = tanh(z) dt + dβ(t), see [40]) with different z0 per dimension. The model is non-linear and solution-space multi-modal, but both p(z, t) and the marginal moments (m(t),P(t)) are available in closed form (see App. B.1). In comparison
(a) Latent space trajectories
Input Predictions →
B as
el in e (1 0 0 0 ×
E M
)
M o m
en t
m at
ch in g L in ea ri za ti o n
(b) Forward prediction
Figure 6: Results on rotating MNIST. (a) shows the latent space prediction mean trajectories for one test image. Evolution of the true trajectory is shown in HSV colour and the predicted trajectory by Euler–Maruyama, moment matching, and linearization scheme is shown in green, blue, and orange (all overlapping one another). (b) shows the progression of the prediction (mean and std dev) of the test image, when it traverses the learned dynamics in the latent space. Both the moment matching and linearization schemes match the baseline (with 1000 Euler–Maruyama trajectories).
to Euler–Maruyama, we control the number of trajectories n by bounding the error (in terms of KL divergence) between Euler–Maruyama and the ground-truth to match the error in our moment matching approach, and consider the methods equivalent in terms of the quality of the solution. See App. B.1 for the required number of trajectories to match the KL divergence for each dimensionality plotted. Fig. 5 shows GPU/CPU wall-clock times (GPU: NVIDIA Tesla V100 32 GB with Intel Xeon Gold 6134 3.2 GHz; CPU: Xeon Gold 6248 2.50GHz). We implement the models in PyTorch [33] and report means of 10 repetitions (std negligible). In low-dimensional cases, both approximation methods outperform sampling, whereas in high dimensions GPU parallelisation becomes dominant and only the linearized approximation remains highly competitive. This example should favour sampling: In the Beneš model, the number of trajectories, n, per d in Euler–Maruyama remains low (n is linear in d), which is due to the diagonal (independent) diffusion matrix. In a correlated latent space n would grow super-linearly (even exponentially), which would further push the difference between methods.
GP-SDE Model Specification We consider an GP-SDE model with just 8 observations of the dynamics, where the lack of data can be compensated for with encoding prior knowledge into the model. We use the model formulation given in Sec. 2.1, and study the effect of GP priors, the first of which is an independent squared exponential (RBF) prior for each dimension which encode continuity and smoothness in the velocity field. The second GP prior is the multi-dimensional curl-free kernel [47] (see App. B.3) which encodes the assumption of a curl-free random vector field. This property can be interpreted as ‘loop aversion’ in the GP-SDE context. The third prior, is a multi-dimensional divergence-free kernel [47] which encodes the assumption of no divergence in the random vector field. This property can be interpreted as ‘energy preservation’ or source-freeness. These properties are visible in Fig. 3, where the hyperparameters are fixed to same values for all models.
Assumed Density Approximation of the FPK We provide an illustrative example of the moment evolution methods in a GP-SDE model with 8 observations along a bean curve and independent squared exponential GP priors per dimension. As a baseline, we solve the FPK in Eq. (7) by finitedifferences discretization in z (see App. A.2 for details). Fig. 4 shows the evolution of the SDE solution over the time-course of 5 s. The probability mass dissolves quickly, which is hard to interpret from the top-row figure alone. Comparison between the point clouds and the top row shows that even with 1000 trajectories and just a two-dimensional space, it is hard to capture detailed structure in the SDE solution. The bottom row compares the local linearization (Sec. 2.4) and the moment matching (Sec. 2.5) assumed density approximations to the exact Gaussian approximation of the FPK solution. The linearized approach is mode-seeking (matches local curvature), while the moment matching approach captures the overall structure of the optimal Gaussian approximation.
Rotating MNIST In the spirit of Fig. 2, we run the proposed methods on Rotating MNIST ([25], available under CC BY-SA 3.0), similar to [49, 5]. The data set consists of various handwritten digit ‘3’s rotated uniformly in 64 angles. We train a VAE [21] first by freezing the latent space dynamics, allowing us to generate the latent samples for learning the dynamics by applying the
Table 1: Rotating MNIST results.
INFERENCE SCHEME MSE NLPD (t = 64)
Euler–Maruyama 0.046 ± 0.006 33.0 ± 5.4 Moment matching 0.051 ± 0.007 52.7 ± 9.5 Linearization 0.052 ± 0.007 54.5 ± 9.9
Table 2: Wall-clock timings for MOCAP.
TIME/S±STD NUMBER OF E-M PATHS LIN. MOM. MAT. 1 25 200
GPU 2.1±.1 6.0±.1 37.1±.1 39.9±.1 40.4±.7 CPU 1.8±.1 1.8±.1 27.7±.7 38.6±1.5 94.2±3.5
Table 3: Test MSE on 297 future MOCAP points averaged over 50 samples. 95% confidence interval reported based on t-statistic. †results from [49], ‡results from [26]
METHOD TEST MSE
DTSBN-S [9] 34.86 ± 0.02† npODE [14] 22.96† NeuralODE [6] 22.49 ± 0.88† ODE2VAE [49] 10.06 ± 1.4† ODE2VAE-KL [49] 8.09 ± 1.95† Latent ODE [37] 5.98 ± 0.28‡ Latent SDE [26] 4.03 ± 0.20‡ Latent SDE (assumed density) 7.55 ± 0.05
trained VAE. Then freezing the VAE encoder/decoder and training a 16-dimensional GP-SDE model in the latent space with independent squared exponential GP priors (see App. B.5). In Fig. 6 we feed in one observation and let it follow the learned dynamics of rotation. As the baseline, we use 1000 trajectories computed using Euler–Maruyama. Fig. 6a demonstrates the model’s capability to learn the latent trajectory, and we show the trajectories for all the methods in three latent dimensions with most variation. The trajectories for the three methods overlap, and qualitatively the results are identical in Fig. 6b. Quantitatively the brute-force sampling baseline gives slightly better MSEs over images and final-step mean negative log predictive densities (NLPD, see Table 1).
Motion Capture Data The CMU walking data set ([1], CMU MoCap available under CC BYND 4.0) is a real-world noisy data set with 50 sensors that track a human subject’s walking. As in Yıldız et al. [49] and Li et al. [26], we model the sequences of a single subject, 35, for which there are 16 train set, three validation set and four test set sequences. The task is to predict the state of the system in the future given three initial points. In this experiment, we demonstrate that replacing SDE solver–based methods by an assumed density approximation, a latent neural SDE system can be learned efficiently without sampling trajectories. For this purpose, and for better comparability to earlier work, the latent SDE drift and diffusion are neural networks. As in Li et al. [26], we regularize the learned posterior process by a prior process. The loss function consists of three terms: reconstruction loss, VAE encoded initial position KL-divergence, and the KL-divergence between posterior/prior processes. The moments of the posterior SDE approximation are denoted by m(t), P(t), those of the prior process m∗(t), P∗(t), and the observation times by {tj}mj=0. The loss becomes
L = − ∑m j=0 log p(x(tj) | z(tj)) + DKL [q(z(t0) |x(t0)) ‖ p(z)] + ∑m
j=1 γDKL [N(m(tj),P(tj)) ‖N(m∗(tj),P∗(tj))] , (17)
where z(tj) and x(tj) are the latent codes and observations, respectively, q denotes the conditional encoder distribution, p(z) is the normal distribution, and log p(x(tj) | z(tj)) is the model likelihood of the observations, given latent codes. While training, the parameters to optimize include those of a latent neural network which defines the prior and posterior dynamics, and of the VAE which encodes an initial point and decodes at each discrete time step corresponding to a train set frame. As the VAE is trained simultaneously to the latent dynamics, the approach does not provide a set of true latent samples to compare to, in contrast to the MNIST experiments. The VAE encoder prior is the normal distribution, whereas the encoder posterior is acting as the initial distribution for the latent SDEs. For the encoder design, we use a fully-connected neural network, which encodes the three first data points to create latent state and context vectors. As the prior process was a SDE with zero drift and σI diffusion. The context is passed through the dynamics, similar to the treatment of velocity in Yıldız et al. [49] (see App. B.6).
The result in Table 3 is competitive considering that solving the latent SDE with the linearization approach roughly matches the required computation budget for one stochastic Runge–Kutta sample in the other methods. This is highlighted in Table 2 which shows the wall-clock times for the model used in this experiment both in a CPU and a GPU setting (see hardware description in previous experiment).

heading: 4 Discussion and Conclusions
text: In this paper our interest has been in both SDE model specification and approximative inference. We considered GP-SDE models for data-scarce applications that need injection of prior knowledge such as in Fig. 3. For inference, we built upon the established methodology of assumed density approximations in signal processing for directly approximating the solution distribution of an Itô SDE, where the methods apply to any latent space SDE models.
We put interest in weak solution concepts for SDEs. We considered both linearization and moment matching based methods for capturing the first two moments of the SDE solutions, which respectively require only O(1) and O(d) evaluations of the model functions per solver step in latent space dimensionality d. Furthermore, they only require solving of one ODE rather than simulating multiple trajectories. This makes them orders of magnitudes lighter than the current state-of-the-art in neural SDEs, which we analyzed both through theoretical bounds in Sec. 2.6, ran numerical experiments with the final error controlled (Fig. 5), and highlighted the practical wall-clock time in a MOCAP experiment. We further argued that a Gaussian assumption makes sense in applications of the form in Fig. 2 as one is typically employed anyway in the encoder–decoder.
There are some key differences between our assumed density approach and commonly used sampling approaches to solving neural SDE models. While stochastic Runge–Kutta methods are typically either concerned with strong (pathwise) or weak (in distribution) solutions, we leverage the even weaker solution concept of only tracking the first two moments of the solution. This simplification of turning solving the SDE into a deterministic ODE problem of its moments, comes with some remarkable computational savings, and often suffice in practical modelling. Then again, the proposed method does not lend itself well to cases where pathwise sample trajectories are required.
We recognize that we cannot give guarantees for the approximated Gaussian integrals to capture the evolution of the true moments outside particular special cases (e.g., 3rd order cubature is exact for polynomials up to order 3). Yet the performance on practical applications is considered reliable, and these kinds of approached are commonly employed across assumed density filtering in signal processing. The use of the approximations schemes we present inherently makes the assumption that the time-marginals of the process are of higher interest compared to the pathwise solutions to the SDE. While the sampling-based methods are less efficient than assumed density approximations, their use is well-justified in applications where sampling trajectories is the purpose of the application.
Codes for the methods and experiments in this paper are available at http://github.com/AaltoML/ scalable-inference-in-SDEs.
Acknowledgments and Disclosure of Funding
Authors acknowledge funding from Academy of Finland (grant numbers 324345 and 339730). We also wish to thank the anonymous reviewers for their comments on our manuscript, and Çağatay Yıldız and Xuechen Li for providing useful details on the MOCAP experiment. We acknowledge the computational resources provided by the Aalto Science-IT project. ET has been employed part-time at Sellforte Oy during the project. The MOCAP data was obtained from mocap.cs.cmu.edu (created with funding from NSF EIA-0196217).

Review summaries:
The paper presents a practical methodology for SDE-based inference in machine learning that avoids the use of Monte-Carlo sampling. The methodology invokes a Gaussian density assumption (often used in latent space modelling) to directly approximate the first two moments of the SDE solution. This is achieved by solving a system of ODEs obtained from the Fokker-Planck-Kolmogorov equation. Empirical evidence shows that this scales well with dimension and the approach is demonstrated on Rotating MNIST and motion capture datasets. The motion capture experiment shows the proposed approach achieves competitive accuracy to previous ODE and sampling-based SDE methods, whilst being substantially faster than the latter.

This paper uses random ordinary differential equations (ODEs) to specify stochastic differential equations (SDEs), where Gaussian processes (GPs) are used to define the vector fields. It is well-known that the dynamics of probability density of states driven by SDEs is given in form of Fokker-Planck equation which is a partial differential equation. In general, there is no closed-form solution of Fokker-Planck equation and the solution is derived by SDE solvers such as Euler–Maruyama (EM), Milstein, and stochastic Runge–Kutta (SRK). These solvers are expensive and inefficient to use in high-dimensional settings. This paper proposes to use an approximate solution by assuming a Gaussian density and tracking the first two moments, mean and covariance matrix. Then the authors derive two ODEs defining the dynamics of mean and covariance matrix. To simplify the ODEs, the authors suggest two schemes: 1- linearizing the Fokker-Planck equation and 2- matching the moments of Fokker-Planck equation. The first scheme is a local approximation and so is less accurate than the second one. However, in terms of runtime, the linearized method is faster than the moment matching one. Overall, the proposed method in this paper can be used to 1- apply prior knowledge to the vector field (trajectories) by using different GP priors and 2- do an efficient inference of probability density function of an SDE by assumed density (Gaussian) approximation of Fokker-Planck equation and deriving two ODEs for the mean and covariance of the distribution which can be solved using ODE solvers and without any stochastic solvers like SRK. 

In this paper, a fast, scalable inference method in stochastic differential equations is introduced, which proceeds by approximating the solutions of the Fokker-Planck-Kolmogorov (FPK) equation instead of relying on sampling schemes. In particular, the solutions to the FPK system is approximated by a Gaussian, whose mean and covariance are shown to satisfy a system of ODEs, which are then further approximated by linearisation or moment matching. The authors demonstrate that solving these approximate systems have significant computational advantage to sampling directly from the SDEs, especially in high dimensional settings and yield qualitatively similar results on experiments on the rotating MNIST dataset and motion capture data.