Track papers across labs, topics, and trends.
Search papers, labs, and topics across Lattice.
This paper introduces Neural Kolmogorov Equations (NKEs), a novel approach to learning stochastic dynamics that overcomes the limitations of traditional neural stochastic differential equations (SDEs) by reformulating the problem to model the evolution of probability densities rather than individual trajectories. By leveraging the Kolmogorov Forward equation, NKEs facilitate parallel-in-time training and can handle complex stochastic drivers, including coupled noise and jump processes. The results demonstrate that NKEs achieve competitive predictive accuracy and enhanced training efficiency on various stochastic benchmarks, marking a significant advancement in the modeling of noisy dynamics.
NKEs transform the landscape of stochastic dynamics modeling, enabling efficient learning from complex noise structures while achieving superior predictive performance.
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.