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

2026-07-21

AI Summary

This work addresses the challenge of efficiently modeling generalized Lévy-driven stochastic processes—characterized by coupled noise and jumps—with existing neural stochastic differential equation approaches, which suffer from non-parallelizable training and limited expressiveness. The authors propose a neural Kolmogorov equation framework that reformulates stochastic dynamics learning as a deterministic modeling task of probability density evolution governed by the Kolmogorov forward equation, naturally accommodating general Lévy noise. By combining Lagrangian Galerkin projection with operator splitting techniques and parameterizing the infinite-dimensional evolution operator via neural networks, the method enables fully parallelizable training across time. Empirical evaluations on multiple benchmark systems with jumps and coupled noise demonstrate substantial improvements over state-of-the-art methods in both predictive accuracy and computational efficiency.

Abstract

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évy-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.

Problem

• neural SDEs
• stochastic dynamics
• general noise
• training scalability
• Lévy processes

Innovation

• Neural Kolmogorov Equations
• Stochastic Differential Equations
• Kolmogorov Forward Equation
• Lévy processes
• Parallel-in-time training