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

• URL: http://arxiv.org/abs/2607.19173v1
• Date: Tue, 21 Jul 2026 15:07:37 GMT
• Title: Neural Kolmogorov Equations: Parallelizable Learning of Stochastic Dynamics under General Noise
• Abstract summary: Existing approaches largely assume uncoupled and continuous noise, limiting their applicability to realistic data.<n>We propose a deterministic, infinite-dimensional reformulation of Neural SDEs based on the Kolmogorov Forward equation.<n>NKEs learn general Lévy-type forcing directly through the operator structure of the KFE, and enable parallel-in-time training.
• Score: 17.664845869094925
• License: http://creativecommons.org/licenses/by/4.0/
• 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.

Related papers

• Stochastic Deep Learning: A Probabilistic Framework for Modeling Uncertainty in Structured Temporal Data [0.0]
  I propose a novel framework that integrates differential equations (SDEs) with deep generative models to improve uncertainty in machine learning applications involving structured and temporal data.<n>This approach, termed Latent Differential Inference (SLDI), embeds an It SDE in the latent space of a variational autoencoder.<n>The drift and diffusion terms of the SDE are parameterized by neural networks, enabling data-driven inference and generalizing classical time series models to handle irregular sampling and complex dynamic structure.
  arXiv  Detail & Related papers  (2026-01-08T18:53:59Z)
• Expanding the Chaos: Neural Operator for Stochastic (Partial) Differential Equations [65.80144621950981]
  We build on Wiener chaos expansions (WCE) to design neural operator (NO) architectures for SPDEs and SDEs.<n>We show that WCE-based neural operators provide a practical and scalable way to learn SDE/SPDE solution operators.
  arXiv  Detail & Related papers  (2026-01-03T00:59:25Z)
• Deep Hierarchical Learning with Nested Subspace Networks [53.71337604556311]
  We propose Nested Subspace Networks (NSNs) for large neural networks.<n>NSNs enable a single model to be dynamically and granularly adjusted across a continuous spectrum of compute budgets.<n>We show that NSNs can be surgically applied to pre-trained LLMs and unlock a smooth and predictable compute-performance frontier.
  arXiv  Detail & Related papers  (2025-09-22T15:13:14Z)
• Improving the Noise Estimation of Latent Neural Stochastic Differential Equations [4.64982780843177]
  Latent neural differential equations (SDEs) have recently emerged as a promising approach for learning generative models from time series data.<n>We investigate this underestimation in detail and propose a straightforward solution: by including an explicit additional noise regularization in the loss function.<n>We are able to learn a model that accurately captures the diffusion component of the data.
  arXiv  Detail & Related papers  (2024-12-23T11:56:35Z)
• Neural Ideal Large Eddy Simulation: Modeling Turbulence with Neural Stochastic Differential Equations [22.707574194338132]
  We introduce a data-driven learning framework that assimilates two powerful ideas: ideal eddy simulation (LES) from turbulence closure modeling and neural differential equations (SDE) for large modeling. We show the effectiveness of our approach on a challenging chaotic dynamical system: Kolmogorov flow at a Reynolds number of 20,000.
  arXiv  Detail & Related papers  (2023-06-01T22:16:28Z)
• Stabilizing Machine Learning Prediction of Dynamics: Noise and Noise-inspired Regularization [58.720142291102135]
  Recent has shown that machine learning (ML) models can be trained to accurately forecast the dynamics of chaotic dynamical systems. In the absence of mitigating techniques, this technique can result in artificially rapid error growth, leading to inaccurate predictions and/or climate instability. We introduce Linearized Multi-Noise Training (LMNT), a regularization technique that deterministically approximates the effect of many small, independent noise realizations added to the model input during training.
  arXiv  Detail & Related papers  (2022-11-09T23:40:52Z)
• Homotopy-based training of NeuralODEs for accurate dynamics discovery [0.0]
  We develop a new training method for NeuralODEs, based on synchronization and homotopy optimization. We show that synchronizing the model dynamics and the training data tames the originally irregular loss landscape. Our method achieves competitive or better training loss while often requiring less than half the number of training epochs.
  arXiv  Detail & Related papers  (2022-10-04T06:32:45Z)
• Neural Operator with Regularity Structure for Modeling Dynamics Driven by SPDEs [70.51212431290611]
  Partial differential equations (SPDEs) are significant tools for modeling dynamics in many areas including atmospheric sciences and physics. We propose the Neural Operator with Regularity Structure (NORS) which incorporates the feature vectors for modeling dynamics driven by SPDEs. We conduct experiments on various of SPDEs including the dynamic Phi41 model and the 2d Navier-Stokes equation.
  arXiv  Detail & Related papers  (2022-04-13T08:53:41Z)
• Incorporating NODE with Pre-trained Neural Differential Operator for Learning Dynamics [73.77459272878025]
  We propose to enhance the supervised signal in learning dynamics by pre-training a neural differential operator (NDO) NDO is pre-trained on a class of symbolic functions, and it learns the mapping between the trajectory samples of these functions to their derivatives. We provide theoretical guarantee on that the output of NDO can well approximate the ground truth derivatives by proper tuning the complexity of the library.
  arXiv  Detail & Related papers  (2021-06-08T08:04:47Z)
• Stochasticity in Neural ODEs: An Empirical Study [68.8204255655161]
  Regularization of neural networks (e.g. dropout) is a widespread technique in deep learning that allows for better generalization. We show that data augmentation during the training improves the performance of both deterministic and versions of the same model. However, the improvements obtained by the data augmentation completely eliminate the empirical regularization gains, making the performance of neural ODE and neural SDE negligible.
  arXiv  Detail & Related papers  (2020-02-22T22:12:56Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.