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

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 . 1.1 Related Work . 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 . 2 Methods . 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 . 2.1 Random Field Ordinary Differential Equations as SDEs . 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 )

Summary: 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.