Neural Stochastic Flows: Solver-Free Modelling and Inference for SDE Solutions

TL;DR

The paper introduces Neural Stochastic Flows (NSFs) to learn the weak solutions of SDEs directly as conditional transition densities, enabling solver-free sampling between arbitrary time points. By using conditional normalising flows and a flow-consistency regularisation based on forward and reverse KL bounds with a bridge distribution, NSFs preserve key stochastic-flow properties while providing tractable log-densities. The latent extension (Latent NSFs) embeds NSFs within variational state-space models to handle irregular sampling and partial observations, incorporating skip-ahead KLs to mitigate long-horizon error accumulation. Across stochastic Lorenz, CMU Motion Capture, and stochastic Moving MNIST, NSFs achieve distributional accuracy comparable to solver-based approaches but with substantial computational speedups, including state-of-the-art extrapolation in latent settings. This solver-free framework promises real-time stochastic modelling for irregular time series and digital twins, with potential extensions to action-conditioned control and diffusion-model hybrids.

Abstract

Stochastic differential equations (SDEs) are well suited to modelling noisy and irregularly sampled time series found in finance, physics, and machine learning. Traditional approaches require costly numerical solvers to sample between arbitrary time points. We introduce Neural Stochastic Flows (NSFs) and their latent variants, which directly learn (latent) SDE transition laws using conditional normalising flows with architectural constraints that preserve properties inherited from stochastic flows. This enables one-shot sampling between arbitrary states and yields up to two orders of magnitude speed-ups at large time gaps. Experiments on synthetic SDE simulations and on real-world tracking and video data show that NSFs maintain distributional accuracy comparable to numerical approaches while dramatically reducing computation for arbitrary time-point sampling.

Figures (10)

• Figure 1: Comparison between (a) traditional neural SDE methods requiring numerical integration and (b) our NSF, where blue and red contours represent conditional distributions for one-step sampling and recursive application, respectively. Panels (c) and (d) illustrate our flow loss concept, ensuring distributional consistency between one-step and two-step transitions through intermediate states.
• Figure 2: Graphical models. (a) Standard state-space model with discrete-time transitions between states. (b) NSF model with continuous-time transitions parameterised by time intervals ($\Delta t_i \coloneqq t_i - t_{i-1}$). (c) Example of skip KL divergence over two time steps. The upper row shows the two-step-ahead prediction from the posterior at $t_0$ via NSF, while the lower row shows the posterior at $t_2$ conditioned on all observations $(\boldsymbol{o}_{t_0}, \boldsymbol{o}_{t_1}, \boldsymbol{o}_{t_2})$. Across all models: solid arrows represent generative processes, dashed arrows represent inference processes, wavy arrows in (c) represent KL divergences.
• Figure 3: Comparison of generated samples (64 samples per panel) on the stochastic Lorenz attractor. Baseline methods are simulated step-by-step. For NSF, each point is an independent sample from the learnt conditional distribution $p(\boldsymbol{x}_t \mid \boldsymbol{x}_s)$ originating from the same initial state and seed, connected visually for comparison. This visualisation choice aids in assessing distributional accuracy over time.
• Figure 4: Trade-off between prediction accuracy (MSE) and computational cost (runtime, using JAX) for latent SDE and Latent NSF on CMU Motion Capture dataset. (a) Setup 1: Within-horizon forecasting. (b) Setup 2: Beyond-horizon extrapolation.
• Figure 4: Fréchet distance comparison on Stochastic Moving MNIST using pre-trained SRVP embeddings. The lower the better. Note that dynamics FD is ill-defined for one-step simulation of latent NSF, which only predicts the marginal distribution of the future frames conditioned on the last observation.