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

AI summary

Overview

Research area: Machine learning for continuous-time stochastic dynamical systems — specifically generative modelling and inference for stochastic differential equations (SDEs), with connections to neural ODEs, normalising flows, diffusion models, and latent state-space models.

Technical level: Advanced. The paper builds on stochastic flow theory, Itô calculus, conditional normalising flows, and variational inference; readers need familiarity with SDEs, change-of-variables, and evidence lower bounds.

Scope: The paper proposes Neural Stochastic Flows (NSFs), a conditional normalising flow architecture that learns SDE transition densities directly, replacing numerical SDE solvers at both training and inference time.

Note: the supplied paper content is truncated, ending mid-sentence in the Experiments section. Numerical results for the CMU Motion Capture and Stochastic Moving MNIST benchmarks, and the descriptions of Appendices C–E, are not present in the available text and are therefore reported below as not available rather than summarised.

What This Paper Is About

Stochastic differential equations describe systems that evolve under a velocity field plus random noise, and they are widely used for noisy or irregularly sampled time series. The central quantity of interest — the transition law p(x_t | x_s), the distribution of a future state given an earlier one over an arbitrary time gap — is normally obtained by running a numerical solver with many small steps, and the cost grows with the size of the time gap. The goal of this paper is to learn that transition law directly with a neural network so that sampling between arbitrary time points requires a single step and no solver.

Key Contributions

1. Neural Stochastic Flows (NSFs). A conditional normalising flow architecture that directly learns the SDE transition distribution p(x_t | x_s) for any s < t, with architectural constraints that enforce the identity property (the map reduces to the identity when Δt = 0), the Markov property, and — for autonomous SDEs — the stationarity property. Because all transformations are bijective, the transition log-density is available in closed form, so both training and inference proceed without an SDE solver.

2. A flow-property regularisation loss. A bi-directional KL divergence loss that encourages the Chapman–Kolmogorov relation between one-step and two-step transitions, using an auxiliary bridge distribution b_ξ(x_tj | x_ti, x_tk) to obtain tractable variational upper bounds for both the forward and reverse KL directions. The total training objective combines the negative log-likelihood with this flow loss, weighted by λ.

3. Latent Neural Stochastic Flows. An extension to noisy and/or partially observed data as a variational state-space model whose transition kernels are NSFs, using a GRU-based amortised encoder, a β-weighted negative ELBO, and a newly proposed skip-ahead KL divergence loss that links non-adjacent time steps directly (rather than recursively, as in overshooting methods).

4. Empirical evaluation and positioning. Experiments on a synthetic stochastic Lorenz attractor, CMU Motion Capture, and Stochastic Moving MNIST, plus a taxonomy (Table 1) showing that NSF — unlike neural ODEs, neural flows, diffusion-model accelerateors, neural SDEs, ARCTA and SDE matching — is solver-free for both training and inference on general Itô SDEs.

Main Findings

• Distributional accuracy on the stochastic Lorenz attractor. Across time points t = 0.25, 0.5, 0.75 and 1.0, NSF achieves low KL divergences — for example at t = 1.0 the H_pred = 1.0 configuration reports 0.2 ± 0.6, compared with 1.5 ± 0.5 for Latent SDE, 53.1 ± 29.3 for Neural LSDE, 14.1 ± 8.4 for Neural GSDE, 44.6 ± 23.4 for Neural LNSDE, and 6.0 ± 0.3 for SDE matching at Δt = 0.01.

• Baselines degrade at longer horizons, NSF does not. Neural LSDE rises from 1.3 ± 0.4 at t = 0.25 to 74.5 ± 24.6 at t = 0.75; Neural GSDE from 1.2 ± 0.4 to 20.2 ± 7.6; SDE matching (Δt = 0.01) from 6.3 ± 0.4 to 11.7 ± 0.5 at t = 0.5. NSF's reported values stay in the same range across horizons.

• FLOPs are dramatically lower and constant in the time gap. NSF reports 53 kFLOPs at t = 0.25 with H_pred = 1.0 and 53 kFLOPs at t = 1.0; with H_pred = 0.5 the cost rises from 53 to 105 kFLOPs, and with H_pred = 0.25 from 53 to 208 kFLOPs as recursive steps are added. By contrast, Latent SDE scales from 959 to 3,760 kFLOPs, Neural GSDE from 1,925 to 7,548 kFLOPs, and SDE matching at Δt = 0.0001 from 184,394 to 737,354 kFLOPs.

• Runtime measurement. Average runtime per 100 samples is reported as 124–148 ms for latent SDE (JAX) versus 0.3 ms for NSF (JAX).

• Claimed speedup. The abstract states up to two orders of magnitude speedup for distant time points, with the largest gains on long-interval forecasts.

• Qualitative Lorenz results. Figure 3 compares 64 samples per panel across ground truth, latent SDE, SDE matching, and NSF (one-step); baseline methods are simulated step-by-step while NSF points are independent samples from the learnt conditional distribution p(x_t | x_s), originating from the same initial state and seed and connected visually.

• Results on CMU Motion Capture and Stochastic Moving MNIST. The abstract states that NSF maintains distributional accuracy comparable to or better than numerical solver methods on these benchmarks, but the truncated content does not include the corresponding tables or numerical values.

Methodology in Plain English

Rather than simulating an SDE step by step, the authors learn a map that jumps directly from one time point to another.

The building block. Sampling starts from a state-dependent Gaussian centred at x_ti + Δt · MLP_μ(c) with scale sqrt(Δt) · MLP_σ(c), where c = (x_ti, Δt, t_i) and ε is a single standard Gaussian vector. This mirrors the Euler–Maruyama form: drift scaled by Δt, diffusion scaled by sqrt(Δt). That Gaussian sample is then pushed through a stack of bijective affine coupling layers, each splitting the state into two partitions and updating one conditioned on the other. Every layer multiplies its scale and shift outputs by Δt, so when Δt = 0 every layer becomes the identity and the flow returns the input unchanged — satisfying the identity property by construction.

Enforcing the flow property. Training minimises the negative log-likelihood of observed transitions plus a regulariser. The regulariser compares the model's one-step distribution from t_i to t_k with the two-step marginal obtained through an intermediate t_j. Because that marginal is intractable, the authors introduce a variational bridge distribution and minimise upper bounds on the KL divergence in both directions — forward KL to encourage coverage, reverse KL to penalise mass placed in unsupported regions.

Handling partial observation. For real data where only noisy observations o are visible, the model becomes a variational state-space model: a GRU encoder produces Gaussian posteriors over latent states from observations and time gaps, the NSF serves as the latent transition kernel, and a decoder maps latents to observations. Training adds a skip-ahead KL term that directly connects a posterior at t_j to the prior transition from an earlier t_i, which is efficient precisely because the NSF can jump across arbitrary gaps without recursion.

Enforcing stationarity. For autonomous SDEs, the absolute time t_i is simply omitted from the conditioning vector, which gives the stationarity property.

Why This Matters

Impact on research. The paper reframes the SDE transition density as something to be learned as a conditional normalising flow rather than simulated, and it argues that this satisfies weak-solution conditions by construction plus regularisation. Table 1 positions it as the only entry in the compared family that is solver-free for both training and inference over general Itô SDEs, which is a different target from diffusion-model speedups such as consistency models, progressive distillation, or rectified flows — those are tied to prescribed boundary-conditioned diffusion processes rather than general SDEs. If it holds up, it also gives latent SDE-style sequence models closed-form transition densities and single-step sampling at arbitrary time gaps.

Real-world applications (as motivated in the paper):

• Robots, where transition queries must be answered in real time.
• Trading algorithms in finance, one of the domains the paper cites for noisy and irregularly sampled time series.
• Digital twins, which the introduction lists alongside robots and trading algorithms as settings needing the transition law p(x_t | x_s).
• Physics and biology, cited as domains where SDEs underpin models of noisy dynamical systems.

Industry relevance. The runtime figures reported (0.3 ms versus 124–148 ms per 100 samples, and up to two orders of magnitude fewer FLOPs for distant horizons) address the practical bottleneck that makes solver-based neural SDEs expensive in deployment. The paper's affiliation with Canon Inc. alongside Imperial College London suggests industrial interest in efficient stochastic sequence modelling. The cost pattern is especially relevant for long-horizon forecasting, where solver-based costs grow with the time gap while NSF's stay flat.

Future Directions

• Extending beyond the tested SDE classes. The framework is presented for Itô SDEs with smooth coefficients; whether the architectural constraints and flow loss extend to degenerate diffusions, non-Itô dynamics, or highly stiff systems is not addressed in the available content.

• Scaling behaviour in higher dimensions. All reported quantitative results are for state dimensions and model sizes not fully described in the truncated text; how the coupling-flow architecture and bridge-distribution regulariser scale to high-dimensional systems remains an open question.

• Tightening or replacing the flow-property regulariser. The Chapman–Kolmogorov relation is encouraged through a variational upper bound with an auxiliary bridge distribution and hyperparameter λ; how sensitive results are to λ, the bridge parameterisation, and the sampled time triplets (t_i, t_j, t_k) is not reported in the available content.

• Full evaluation of the latent model. The paper states that latent NSFs should model noisy and partially observed data, and lists CMU Motion Capture and Stochastic Moving MNIST as benchmarks, but the corresponding numerical results are not present in the truncated content, leaving their quantitative performance, and the effect of the skip-ahead KL weight β_skip, unverified here.

Target Audience

Researchers and advanced practitioners in probabilistic machine learning and continuous-time modelling: people working on neural SDEs, neural ODEs, normalising flows, variational state-space models, and generative models for time series. It is also relevant to applied scientists in finance, physics, robotics, and biology who need fast probabilistic forecasts over irregular time grids, and to engineers who must deploy stochastic sequence models under latency or compute budgets. Readers without a background in stochastic calculus and variational inference will find the theory sections demanding, though the core idea — learn the jump map instead of simulating it — is accessible to anyone familiar with normalising flows.

Authors’ 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.