# Neural Kolmogorov Equations: Parallelizable Learning of Stochastic Dynamics under General Noise

arXiv: https://arxiv.org/abs/2607.19173 · HTML: https://mlsanity.com/paper/2607.19173/

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Builds on 18 of its 39 references that are in the map.
- First posted on arXiv 21 Jul 2026.

## Builds on (most cited first, 18 of 18)

- [Denoising Diffusion Probabilistic Models](https://mlsanity.com/paper/2006.11239/) (2020)
- [Score-Based Generative Modeling through Stochastic Differential Equations](https://mlsanity.com/paper/2011.13456/) (2020)
- [Flow Matching for Generative Modeling](https://mlsanity.com/paper/2210.02747/) (2022)
- [Neural Ordinary Differential Equations](https://mlsanity.com/paper/1806.07366/) (2018)
- [Variational Inference with Normalizing Flows](https://mlsanity.com/paper/1505.05770/) (2015)
- [FFJORD: Free-form Continuous Dynamics for Scalable Reversible Generative Models](https://mlsanity.com/paper/1810.01367/) (2018)
- [Neural Controlled Differential Equations for Irregular Time Series](https://mlsanity.com/paper/2005.08926/) (2020)
- [Universal Differential Equations for Scientific Machine Learning](https://mlsanity.com/paper/2001.04385/) (2020)
- [Scalable Gradients for Stochastic Differential Equations](https://mlsanity.com/paper/2001.01328/) (2020)
- [Latent ODEs for Irregularly-Sampled Time Series](https://mlsanity.com/paper/1907.03907/) (2019)
- [Neural SDEs as Infinite-Dimensional GANs](https://mlsanity.com/paper/2102.03657/) (2021)
- [Neural Jump Stochastic Differential Equations](https://mlsanity.com/paper/1905.10403/) (2019)
- [Neural SDE: Stabilizing Neural ODE Networks with Stochastic Noise](https://mlsanity.com/paper/1906.02355/) (2019)
- [FP-Diffusion: Improving Score-based Diffusion Models by Enforcing the Underlying Score Fokker-Planck Equation](https://mlsanity.com/paper/2210.04296/) (2022)
- [SDE Matching: Scalable and Simulation-Free Training of Latent Stochastic Differential Equations](https://mlsanity.com/paper/2502.02472/) (2025)
- [Arbitrage-free neural-SDE market models](https://mlsanity.com/paper/2105.11053/) (2021)
- [Singularity structures and impacts on parameter estimation in finite mixtures of distributions](https://mlsanity.com/paper/1609.02655/) (2016)
- [Neural SDEs as a Unified Approach to Continuous-Domain Sequence Modeling](https://mlsanity.com/paper/2501.18871/) (2025)

## Abstract (arXiv)

Neural stochastic differential equations (SDEs) have emerged as powerful tools for learning noisy or stochastic dynamics directly from data; however, existing approaches largely assume uncoupled and continuous noise, limiting their applicability to realistic stochastic drivers, and often scale poorly in time, requiring expensive autoregressive training. To address these limitations, we propose Neural Kolmogorov Equations (NKEs), a deterministic, infinite-dimensional reformulation of Neural SDEs based on the Kolmogorov Forward equation, transforming the learning problem from modelling individual stochastic trajectories to modelling the evolution of probability densities. NKEs learn general L'evy-type stochastic forcing directly through the operator structure of the KFE, and enable parallel-in-time training via a Lagrangian Galerkin projection and operator splitting. We evaluate NKEs on several stochastic benchmarks, including systems with coupled noise and jump processes, and verify that NKEs provide flexible models that accurately recover deterministic and stochastic dynamics with competitive predictive accuracy and improved training efficiency. Code and pretrained models will be released.