Research map / Probabilistic and scientific modeling / Differential equations and PDEs
Neural Stochastic Flows: Solver-Free Modelling and Inference for SDE Solutions
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Proposes Neural Stochastic Flows (NSFs), conditional normalising flows that learn SDE transition laws directly with architectural constraints and a bidirectional KL flow (Chapman-Kolmogorov) regulariser, plus latent variants, enabling solver-free one-step sampling between arbitrary times.
Problem
Sampling the transition law p(x_t|x_s) of SDEs over arbitrary time gaps normally requires costly numerical solvers, and existing neural flow or diffusion-acceleration methods do not handle general stochastic dynamics with the required flow properties.
Approach
Parameterise the SDE weak solution directly as a conditional normalising flow whose Gaussian base and affine coupling layers scale with the time gap (identity at Δt=0, Euler-Maruyama-like), and regularise with a bridge-based bidirectional KL loss enforcing Chapman-Kolmogorov consistency; embed this as transition kernel of a latent variational state-space model.
Key results
- On the stochastic Lorenz attractor, one-step NSF achieves lowest KL at most time points (e.g. 0.2 at t=1.0) with 53 kFLOPs versus 3,760-7,548 kFLOPs for latent SDE and neural SDE baselines, and about 0.3 ms vs >100 ms per 100 samples.
- On CMU Motion Capture Setup 2 (beyond-horizon extrapolation), Latent NSF achieves MSE 3.41 ± 0.27, the best reported, versus 4.26 for SDE matching and 9.52 for latent SDE.
- Latent NSF is reported as more than two orders of magnitude faster than latent SDE on Motion Capture (3.5 ms vs 75 ms per 100 samples) and on Moving MNIST one-step NSF takes 358 ms vs 751 ms for latent SDE for full-sequence prediction.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| use NSF instead of Latent SDE for one-step prediction | Stochastic Lorenz attractor · t = 0.25 | KL: 2.1 ± 0.9 -> 0.8 ± 0.7 | Helped |
| use NSF instead of Latent SDE for one-step prediction | Stochastic Lorenz attractor · t = 0.5 | KL: 1.8 ± 0.1 -> 1.3 ± 0.1 | Helped |
| use NSF instead of Latent SDE for one-step prediction | Stochastic Lorenz attractor · t = 0.75 | KL: 0.9 ± 0.3 -> 0.6 ± 0.3 | Helped |
| use NSF instead of Latent SDE for one-step prediction | Stochastic Lorenz attractor · t = 1.0 | KL: 1.5 ± 0.5 -> 0.2 ± 0.6 | Helped |
| use Latent NSF instead of SDE matching | CMU Motion Capture · Setup 1 | MSE: 5.20 ± 0.43 -> 8.62 ± 0.32 | Hurt |
| use Latent NSF instead of SDE matching | CMU Motion Capture · Setup 2 | MSE: 4.26 ± 0.35 -> 3.41 ± 0.27 | Helped |
| reduce maximum single-step horizon from 1.0 to 0.5 (ablation) | NSF on Stochastic Lorenz attractor · t = 0.25 | KL: 0.8 ± 0.7 -> 2.4 ± 1.9 | Hurt |
| reduce maximum single-step horizon from 1.0 to 0.5 (ablation) | NSF on Stochastic Lorenz attractor · t = 0.5 | KL: 1.3 ± 0.1 -> 1.3 ± 0.1 | No clear effect |
| reduce maximum single-step horizon from 1.0 to 0.5 (ablation) | NSF on Stochastic Lorenz attractor · t = 0.75 | KL: 0.6 ± 0.3 -> 1.0 ± 0.4 | Hurt |
| reduce maximum single-step horizon from 1.0 to 0.5 (ablation) | NSF on Stochastic Lorenz attractor · t = 1.0 | KL: 0.2 ± 0.6 -> 1.7 ± 1.1 | Hurt |
| reduce maximum single-step horizon from 1.0 to 0.25 (ablation) | NSF on Stochastic Lorenz attractor · t = 0.25 | KL: 0.8 ± 0.7 -> 1.2 ± 0.7 | Mixed |
| reduce maximum single-step horizon from 1.0 to 0.25 (ablation) | NSF on Stochastic Lorenz attractor · t = 0.5 | KL: 1.3 ± 0.1 -> 1.2 ± 0.1 | Helped |
Surprising findings
- Latent NSF is best on CMU Motion Capture Setup 2 but SDE matching is best on Setup 1 (5.20 vs 8.62), and the authors' own latent SDE reproduction differs from the originally reported numbers.
- SDE-GAN training was unstable on the Lorenz data and was excluded from comparison.
Limitations
- Chapman-Kolmogorov consistency is only approximately enforced via variational upper bounds, so discrepancies may arise outside the training regime.
- The maximum one-shot horizon must be chosen to balance computational efficiency and accuracy.
- The affine coupling architecture constrains expressivity although it keeps Jacobians analytic.
Experimental setup
Datasets: Stochastic Lorenz attractor (synthetic) CMU Motion Capture walking subset Stochastic Moving MNIST (modified bouncing variant)
Benchmarks: Stochastic Lorenz attractor CMU Motion Capture (Setup 1 and Setup 2) Stochastic Moving MNIST
Metrics: KL divergence (KDE-estimated) FLOPs runtime test MSE Frechet distance (static, dynamics, frame-wise)
Baselines: Latent SDE Neural LSDE Neural GSDE Neural LNSDE SDE matching ARCTA Latent ODE NCDSSM
Builds on: latent SDE neural flows conditional normalising flows variational state-space models
Closest papers by meaning
- In-Context Learning of Stochastic Differential Equations with Foundation Inference Models (NeurIPS 2025 · similarity 0.91)
- Stable Neural Stochastic Differential Equations in Analyzing Irregular Time Series Data (ICLR 2024 · similarity 0.90)
- SDE Matching: Scalable and Simulation-Free Training of Latent Stochastic Differential Equations (ICML 2025 · similarity 0.90)
- Efficient Training of Neural Stochastic Differential Equations by Matching Finite Dimensional Distributions (ICLR 2025 · similarity 0.90)
- Trajectory Flow Matching with Applications to Clinical Time Series Modelling (NeurIPS 2024 · similarity 0.88)
- Closing the Approximation Gap in Simulation-free Latent SDEs (NeurIPS 2026 · similarity 0.88)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.