🤖 machine learning

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

This paper introduces Neural Kolmogorov Equations (NKEs), a deterministic, infinite-dimensional framework that reformulates Neural SDEs to model probability density evolution, thereby enabling the learning of general Lévy-type noise and achieving parallel-in-time training for improved efficiency and accuracy compared to existing autoregressive approaches.

Original authors: Arthur Bizzi, Olga Fink

Published 2026-07-22

📖 7 min read🧠 Deep dive

CC BY 4.0

Original authors: Arthur Bizzi, Olga Fink

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are trying to predict the future path of a leaf floating down a river. In a perfect, calm world, you could draw a single, smooth line showing exactly where it will go. But real rivers are messy. There are whirlpools, sudden gusts of wind, and random bumps against rocks. This is the world of "stochastic dynamics"—systems that move with a mix of predictable rules and chaotic, random noise. Scientists and engineers use these models to understand everything from how stock markets crash to how molecules bounce around inside your body.

For a long time, computers tried to learn these messy paths by simulating thousands of individual "leaves" (or particles) one by one, step by step. It's like trying to predict the weather by tracking every single raindrop individually. It's slow, it's hard to do in parallel (like having many computers work together), and it often gets stuck when the noise is weird or jumps suddenly. The big question has been: Can we teach a computer to understand the shape of the chaos itself, rather than just chasing individual drops?

This paper introduces a clever new way to do just that, called Neural Kolmogorov Equations (NKEs). Instead of trying to guess where a single particle will land next, the authors teach a computer to watch how a whole "cloud" of particles spreads, stretches, and splits over time. Think of it like switching from tracking a single runner in a race to watching the entire pack of runners as a single, shifting blob of energy.

The team, led by Arthur Bizzi and Olga Fink, found that by looking at the "cloud" (the probability density) instead of the individual dots, they could learn the rules of the game much faster and more accurately. Their method, tested on various chaotic systems, showed that this "cloud-watching" approach handles weird, sudden jumps and complex, tangled noise better than previous methods. It allows computers to learn the underlying rules of randomness without needing to simulate every single step of every single particle, making it a powerful new tool for understanding the unpredictable world.

The Old Way: Chasing the Chaos

To understand why this new method is a big deal, let's look at how things used to work. Imagine you are trying to learn how a ball bounces in a room filled with invisible, jittery air currents. The old way, used by "Neural Stochastic Differential Equations" (Neural SDEs), was to simulate the ball's journey over and over again.

The computer would say, "Okay, the ball is here. Now, let's guess the wind, move the ball, guess the wind again, move it again." It does this thousands of times to build a picture of where the ball might end up. This is like trying to learn how to drive a car by crashing into a wall, resetting, and trying again, over and over. It's expensive, slow, and if the wind suddenly changes direction (a "jump" or a "shock"), the computer often gets confused. It also struggles if the wind blows differently in different directions (coupled noise), because it usually assumes the chaos is simple and uniform.

The New Way: Watching the Cloud

The authors of this paper realized that while individual particles are chaotic, the group of particles often follows a smoother, more predictable pattern. They borrowed a concept from physics called the Kolmogorov Forward Equation.

Here is the analogy: Instead of tracking a single leaf, imagine the leaf is part of a giant, glowing fog.

Drift: If there is a current in the river, the whole fog moves in one direction.
Diffusion: If the water is turbulent, the fog spreads out, getting wider and flatter.
Jumps: If a sudden gust of wind hits, the fog might split into two separate blobs or get stretched into a weird shape.

The "Neural Kolmogov Equations" (NKEs) treat the fog as the main character. The computer learns a set of rules that describe how this fog moves, spreads, and splits.

How the Magic Happens

The team broke this down into three parts, using a technique called Operator Splitting. Think of it like a chef preparing a complex dish by handling the ingredients one by one, then mixing them together perfectly.

The Drift (The Push): The computer learns how the fog moves on average. It looks at where the center of the fog is moving and learns the "push" force.
The Diffusion (The Spread): The computer learns how the fog gets wider. It looks at how the edges of the fog expand and learns the "spread" force.
The Jumps (The Split): This is the tricky part. Sometimes the fog doesn't just spread; it suddenly breaks apart. The computer learns to recognize these sudden breaks and models them as a new, smaller cloud appearing somewhere else.

To make this work, they used a trick called Gaussian Mixture Models. Imagine the fog isn't just one big blob, but a collection of smaller, round, glowing orbs (like a swarm of fireflies). The computer tracks how these orbs move, stretch, and multiply. If an orb gets hit by a "jump," it might split into two new orbs. By tracking these orbs, the computer can learn the rules of the chaos without ever needing to simulate a single, lonely particle.

What They Found

The authors tested this new method on several challenging scenarios, including:

The Lorenz Oscillator: A famous chaotic system that looks like a butterfly shape.
The Black-Scholes Model: A system used to predict stock prices, which often has "coupled" noise (where the randomness in one price affects another).
Systems with Jumps: Scenarios where the system suddenly jumps to a new state, like a stock market crash or a particle hitting a wall.

In these tests, the NKE method was a clear winner.

Accuracy: It predicted the future shape of the fog much better than the old methods. For example, in the Lorenz system test, the old "Euler-Maruyama" method had an error score of about 0.95 for the spread of the cloud, while the new NKE method got it down to 0.075. That's a huge difference.
Speed: Because the computer doesn't have to simulate thousands of individual paths one by one, it can learn much faster. In some tests, the old methods took thousands of seconds to train, while the new method did it in a fraction of the time.
Flexibility: The new method could handle "coupled" noise (where things are tangled together) and sudden jumps, which the old methods struggled with or couldn't do at all.

The Catch (Limitations)

Of course, no magic trick is perfect. The authors point out a few things to keep in mind:

Sampling Rate: The method works best when the data comes in at regular intervals (like a clock ticking every second). If the data is messy and comes at random times, especially when jumps are involved, the method gets a bit stuck.
Complexity: The method uses "Gaussian mixtures" (the glowing orbs). If the system is incredibly complex or has too many dimensions (like a fog with millions of variables), the math gets too heavy, and the "orbs" might not be enough to describe the shape accurately.

Why It Matters

This paper doesn't just offer a slightly faster way to do the same thing; it changes the perspective. Instead of fighting the chaos by chasing individual particles, it embraces the chaos by modeling the shape of the probability cloud.

By turning a messy, random problem into a deterministic (predictable) problem about how shapes move and change, the authors have opened the door to learning complex, real-world systems that were previously too difficult to model. Whether it's predicting how a virus spreads, how a financial market reacts to a shock, or how a robot moves through a bumpy environment, this "cloud-watching" approach gives us a sharper, faster, and more flexible lens to see the future.

Technical Summary: Neural Kolmogorov Equations (NKEs)

1. Problem Formulation

The paper addresses the limitations of existing Neural Stochastic Differential Equations (SDEs) in learning stochastic dynamics from data. While Neural SDEs have emerged as powerful tools for modeling continuous-time Markov systems, current approaches face three primary constraints:

Noise Diversity: Existing methods largely assume uncoupled, continuous Gaussian noise (e.g., Brownian motion). They struggle with non-isotropic noise, coupled noise, and discontinuous dynamics such as heavy-tailed fluctuations or jump processes (Lévy processes).
Scalability and Training Efficiency: Most Neural SDEs rely on autoregressive training via stochastic integration schemes (e.g., Euler-Maruyama). This requires repeated trajectory simulation, which is computationally expensive, difficult to parallelize in time, and often numerically inaccurate.
Modeling Assumptions: Many approaches assume the structure of the driving noise is known (fixed parametric forms), preventing the learning of general jump-diffusion operators directly from data.

The authors frame the problem as an inverse problem: given a sequence of independent realizations of a Lévy-Itô SDE observed at discrete time steps, the goal is to learn the drift $f$ , the diffusion coefficient $gg^\top$ , and a conditional generative model for the jump distribution $h$ , without assuming a specific parametric form for the noise.

2. Methodology: Neural Kolmogorov Equations (NKEs)

The core proposal is Neural Kolmogorov Equations (NKEs), a deterministic, infinite-dimensional reformulation of Neural SDEs. Instead of modeling individual stochastic trajectories, NKEs model the evolution of the probability density function (PDF) governed by the Kolmogorov Forward Equation (KFE), also known as the Fokker-Planck equation.

2.1 Theoretical Foundation

The KFE describes the macroscopic evolution of a system's probability mass $p_t(x)$ via a linear Partial Differential Equation (PDE):
$\frac{d}{dt}p_t = \mathcal{L}p_t = (\mathcal{A} + \mathcal{D} + \mathcal{J})p_t$
where $\mathcal{L}$ is the Kolmogorov generator decomposed into:

Advection ( $\mathcal{A}$ ): Driven by the drift vector field $f$ .
Diffusion ( $\mathcal{D}$ ): Driven by the diffusion tensor $gg^\top$ .
Jump ( $\mathcal{J}$ ): Driven by the Lévy measure $h$ .

This formulation transforms the learning problem from stochastic trajectory simulation to deterministic PDE learning, enabling the use of tools from PDE analysis and allowing for noise-agnostic modeling.

2.2 Lagrangian Galerkin Projection

Since the KFE operates in an infinite-dimensional space of measures, the authors propose a tractable finite-dimensional projection using Lagrangian Galerkin discretization.

Representation: The evolving density $p_t$ is approximated as a mixture of $K$ Gaussians: $p_t(x) \approx \sum \pi_k \mathcal{N}(\mu_k, \Sigma_k)$ .
Projection: The infinite-dimensional generator $\mathcal{L}$ is projected onto the space of Gaussian mixtures ( $\mathcal{G}$ ) via operators $\Pi$ and $\Pi^\dagger$ , yielding a projected generator $\hat{\mathcal{L}} = \Pi \circ \mathcal{L} \circ \Pi^\dagger$ .
Operator Actions:
- Advection: Transports the mean and stretches/rotates the covariance based on the drift $f$ .
- Diffusion: Spreads the covariance without affecting the mean.
- Jumps: Modeled as a convolution with a jump measure. For tractability, small jumps are approximated as state-dependent advection-diffusion, while large jumps are modeled via a conditional Gaussian mixture that splits the original Gaussian components.

2.3 Learning Strategy and Operator Splitting

NKEs employ a Strang splitting scheme (second-order accuracy) rather than the first-order Lie-Trotter splitting used in standard Euler-Maruyama (EM) methods. The training proceeds in four distinct steps:

Data Preprocessing: Trajectory data is clustered into empirical Gaussian packets. Large jumps are identified and filtered out using a threshold based on the Median Absolute Deviation (MAD) to separate the advection-diffusion learning from the jump learning.
Advection Learning: A neural network $f_\theta$ learns the drift by minimizing the mismatch between the observed mean evolution and the predicted transport of Gaussian means. A "meet-me-halfway" condition (trapezoidal rule) is used to enforce second-order accuracy.
Diffusion Learning: A neural network parameterizing $gg^\top_\vartheta$ learns the diffusion coefficient by tracking the time variation of the Gaussian covariances, after subtracting the drift contribution.
Jump Learning: A conditional generative model (Mixture Density Network) learns the jump semigroup $H_s$ by modeling the distribution of filtered jump increments as a convolution of the source Gaussian with a learned jump mixture.

Crucially, this approach allows for parallel-in-time training because the loss functions for drift and diffusion can be computed independently across time steps and particles, avoiding the autoregressive dependency of standard Neural SDEs.

3. Key Contributions

The paper claims the following contributions:

General-Form Noise Learning: NKEs introduce a noise-agnostic strategy capable of handling arbitrary coupling in Gaussian noise and general jump processes (including state-dependent jump distributions) without predefined parametric assumptions.
Scalability and Accuracy: By reformulating the problem as a deterministic PDE on Gaussian mixtures and utilizing Strang splitting, NKEs enable parallel-in-time training. This avoids the computational cost of autoregressive simulation while retaining (and often improving) predictive accuracy compared to EM-based methods.
Unified Framework: The method provides a unified treatment for continuous diffusion and discontinuous jump dynamics within a single generative framework, learning the structure of the driving noise directly from density evolution.

4. Experimental Results

The authors evaluated NKEs on several benchmarks, including systems with coupled noise and jump processes:

Benchmarks: Lorenz oscillator (geometric Brownian noise), Multivariate Black-Scholes (coupled geometric Brownian motion), Multivariate double-well system (Gaussian jumps), and a rarefied coarse-graining molecular dynamics problem.
Baselines: Compared against Parallelizable Euler-Maruyama, Trajectory Flow Matching (TFM), and Autoregressive SDE-GANs.
Performance:
- Accuracy: NKEs consistently achieved lower Mean Squared Error (MSE) in both mean vectors and covariance matrices, as well as lower Wasserstein distances, compared to all baselines. For instance, on the Lorenz oscillator, NKEs reduced Mean MSE from ~1.8 (EM/TFM) to ~0.047.
- Training Time: NKEs maintained competitive training times. While slightly slower than the simplest EM baselines in some cases, they were significantly faster than SDE-GANs (which required thousands of seconds) and comparable to TFM.
- Robustness: NKEs successfully modeled systems with coupled, non-diagonal noise and jump processes where baselines either failed to converge or exhibited poor accuracy.

5. Significance and Limitations

Significance:
The paper argues that NKEs represent a principled alternative to classical simulation-based Neural SDE training. By shifting the focus from individual noisy trajectories to the deterministic evolution of probability densities, the method overcomes the scalability bottlenecks of autoregressive training and the noise diversity limitations of maximum-likelihood estimation. It enables the learning of complex, general Lévy-type stochastic drivers directly from data.

Limitations:
The authors explicitly acknowledge two main limitations:

Sampling Rate Constraints: In the presence of jumps, the method currently requires constant sampling rates because the jump semigroup is learned as a fixed iteration $e^{\mathcal{J}s}$ . Extending this to non-uniform sampling remains an open challenge for the current Lagrangian-GMM representation.
Dimensionality and Representation: The reliance on Gaussian Mixture Models (GMMs) limits applicability to extremely high-dimensional systems (e.g., complex generative models). Furthermore, for $K > 1$ , the projection operator can be ill-posed, leading to non-uniqueness and convergence issues when learning complex, state-dependent jump patterns without additional regularization.