[Paper Review] Scalable Inference in SDEs by Direct Matching of the Fokker-Planck-Kolmogorov Equation

Arno Solin, Ella Tamir+1 more|arXiv (Cornell University)|Oct 29, 2021

Gaussian Processes and Bayesian Inference4 citations

TL;DR

This paper proposes a scalable, sampling-free inference method for stochastic differential equations (SDEs) by directly matching the first two moments of the Fokker-Planck-Kolmogorov (FPK) equation using Gaussian assumed density approximations. By formulating the SDE solution as a deterministic ODE on moments, the approach achieves orders-of-magnitude speedup over sampling-based methods like Euler–Maruyama while maintaining competitive performance on time-series and generative modeling tasks.

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.

Motivation & Objective

• To address the computational inefficiency of sampling-based inference in SDEs, especially in high-dimensional and data-scarce settings.
• To develop a direct, non-sampling approach to approximate the solution of the intractable Fokker-Planck-Kolmogorov (FPK) equation for SDEs.
• To enable fast, scalable inference in latent space SDE models by leveraging Gaussian assumptions on the time-marginal distributions.
• To extend the applicability of SDEs in machine learning by combining GP priors on the drift with moment-matching inference.
• To provide a deterministic alternative to stochastic Runge–Kutta schemes that avoids pathwise sampling while preserving distributional accuracy.

Proposed method

• Formulate the SDE using a Gaussian process prior over the drift vector field, enabling flexible, prior-informed modeling of latent dynamics.
• Derive a closed-form approximation to the Fokker-Planck-Kolmogorov (FPK) equation by matching the first two moments (mean and covariance) of the solution distribution.
• Solve the resulting moment dynamics as a deterministic ordinary differential equation (ODE), avoiding the need for trajectory sampling.
• Use a Gaussian assumed density approximation to represent the time-marginal distribution of the SDE, enabling efficient computation in high-dimensional latent spaces.
• Implement the method using a variational inference framework, where the encoder outputs initial moments and the decoder uses moment-matched SDE dynamics.
• Optimize the model via likelihood maximization or variational objectives, with gradients computed through the moment ODE using automatic differentiation.

Experimental results

Research questions

• RQ1 Can direct moment matching of the Fokker-Planck-Kolmogorov equation provide a scalable alternative to sampling-based inference in SDEs?
• RQ2 How does moment-matching inference compare in accuracy and speed to Euler–Maruyama sampling in high-dimensional and low-data regimes?
• RQ3 To what extent can a Gaussian assumed density approximation capture the true time-marginal distribution of an SDE with non-linear drift and diffusion?
• RQ4 Can GP-prior SDEs combined with moment matching enable effective modeling in data-scarce applications?
• RQ5 Does the proposed method maintain competitive performance on time-series and generative modeling benchmarks while reducing computational cost?

Key findings

• The moment-matching method achieved a test MSE of 7.55±0.05 on the MOCAP dataset, outperforming the baseline latent SDE (4.03±0.20) when using the same architecture, indicating robustness under assumed density constraints.
• Wall-clock timings showed the moment-matching method required only 6.0±0.1 seconds on GPU and 1.8±0.1 seconds on CPU for a single evaluation, compared to 37.1±0.1 seconds for Euler–Maruyama with 200 paths.
• On the Rotating MNIST task, the moment-matching method achieved an MSE of 0.051±0.007 and NLPD of 52.7±9.5, demonstrating competitive performance despite avoiding trajectory sampling.
• The method scaled efficiently to high-dimensional latent spaces, with computational cost growing linearly in dimensionality (O(d)) per step, unlike sampling methods that scale poorly with dimension.
• The approach reduced inference time by orders of magnitude compared to sampling-based SDE solvers, making it practical for real-time and large-scale applications.
• The method was shown to be effective even in low-data regimes, where GP priors on the drift field provided strong inductive bias, improving generalization.