Neural Kolmogorov Equations: Parallelizable Learning of Stochastic Dynamics under General Noise
Arthur Bizzi EPFL, Imperial College [email protected]
Olga Fink EPFL [email protected]
arXiv:2607.19173v1 [cs.LG] 21 Jul 2026
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.
1
Introduction
Stochastic differential equations (SDEs) are a canonical framework for modelling continuous-time Markov systems, with applications spanning science, engineering, and finance. In this formulation, system trajectories evolve under the combined influence of a deterministic drift, representing forces and systemic trends, and a stochastic driver, accounting for uncertainty and unresolved dynamics. While classical approaches often assume continuous perturbations, such as Brownian motion, real systems frequently exhibit heavy-tailed fluctuations, abrupt discontinuities, and intermittent shocks, as naturally captured by the family of Lévy processes. Neural stochastic differential equations extend this framework by parameterizing deterministic and stochastic terms with neural networks. Introduced as stochastic generalizations of Neural Ordinary Differential Equations (Neural ODEs [1]), Neural SDEs enable the learning of stochastic dynamics directly from data, allowing flexible modelling in the absence of governing equations. They provide a principled way to incorporate uncertainty into continuous-time models, naturally capture irregularly sampled time series, and enable generative modelling of stochastic trajectories. These properties have led to growing applications across scientific [2] and financial [3] modelling. However, existing implementations are still limited in their noise diversity and scalability in time. Neural SDEs are generally limited to scalar or uncoupled Gaussian noise, due to the limitations of maximum-likelihood estimation, preventing them from modelling non-isotropic and discontinuous stochasticity. Moreover, existing Neural SDE models typically assume that the structure of the driving noise process is known, i.e., that the stochastic forcing follows a fixed parametric form such as Brownian motion, or jumps with a predefined distribution. Preprint.
Moreover, unlike Neural ODEs, Neural SDEs cannot be trained by directly differentiating observed trajectories, due to the presence of noise. As a result, most approaches rely on repeatedly simulating trajectories using stochastic integration schemes. This procedure is computationally expensive, difficult to parallelize in time, and often numerically inaccurate. Recent approaches attempt to alleviate these issues by leveraging Euler-Maruyama (EM) discretizations to enable simulation-free or partially parallelized training. However, these discretization-based strategies may suffer from ill-posedness, ultimately degrading predictive performance. In this paper, we address these limitations by leveraging a deterministic representation of stochastic dynamics via the Kolmogorov Forward Equation (KFE). While SDEs describe the microscopic evolution of individual sample paths, the KFE provides a macroscopic description of the system in terms of the evolution of a path’s probability density (see Fig. 1), given by a partial differential equation (PDE). The KFE therefore provides a deterministic representation of SDEs, enabling the direct application of numerical and theoretical tools from the analysis of PDEs. Most importantly, it is noise-agnostic, which will allow us to represent substantially more diverse stochastic drivers.
Figure 1: A stochastic harmonic oscillator. The Kolmogorov Equations treat individual particle realizations (left) from the macroscopic perspective of mass distributions moving (right).
Inspired by this perspective, we propose Neural Kolmogorov Equations (NKEs), a framework for Neural SDE modeling based on directly learning the advection and diffusion dynamics of the KFE. By modeling these operators with neural networks, NKEs provide a principled and accurate alternative to classical Euler–Maruyama–based training schemes, while allowing the model to infer the structure of the driving noise process itself. While the KFE has been used implicitly for SDE learning [4], to the best of our knowledge NKEs are the first framework to treat it as the primary learning object and parameterize both drift and stochastic forcing through its operator structure with neural networks. 1. General-form noise learning. We introduce a noise-agnostic learning strategy for stochastic dynamics based on semigroup representations and generative modelling. NKEs directly model the stochastic forcing of an SDE, enabling a unified treatment of Gaussian/Brownian noise with arbitrary coupling and jump processes with general jump distributions. Consequently, the framework can not only estimate the intensity of pre-specified noise models but also learn the structure of the driving noise itself when no parametric assumptions are available. 2. Scalability in time and accuracy. We propose a novel learning scheme for drift and noise estimation based on a Neural PDE representation of the Kolmogorov equation. This formulation enables training procedures that avoid the expensive autoregressive simulation typically required in Neural SDE training while retaining accuracy. We interpret Euler–Maruyama schemes as Lie-Trotter operator splitting, and introduce a deterministic generalization based on Strang splitting, which results in more accurate parallelizable learning. We evaluate the proposed framework on multiple benchmarks in the field of scientific machine learning, including systems driven by non-diagonal and discontinuous noise. We further compare NKEs with existing autoregressive and simulation-free frameworks for SDE learning. 2
2
Related Work
Neural SDEs. Neural ODEs [1] and Neural Controlled Differential Equations [5] learn continuoustime deterministic dynamics in terms of neural vector fields, which can be efficiently trained by directly regressing on estimated derivatives of time-series data. Neural SDEs extend this paradigm to stochastic dynamics by jointly modelling deterministic and stochastic forcing. Early formulations such as latent SDEs [6] and Neural SDE GANs [7] relied on autoregressive simulation and stochastic adjoint methods, making training computationally expensive and difficult to parallelize in time. Diffusion and Flow-Matching. In the generative setting, the limitations of autoregressive NeuralSDE-based Continuous Normalizing Flows [8] motivated the development of simulation-free approaches such as diffusion models [9], score-based generative models [10], and Flow Matching [11]. While highly successful for generative sampling, these approaches primarily learn mappings between static distributions and are not naturally designed for continuous-time stochastic system identification from sequential observations. Simulation-free learning and Euler–Maruyama methods. The scalability limitations of conventional Neural SDEs have motivated multiple simulation-free approaches based on Euler–Maruyama discretizations of Itô SDEs [12]. These methods directly regress on finite differences between consecutive observations, similarly to Neural ODE training. Examples include Trajectory Flow Matching [13], which adapts flow-matching ideas to stochastic sequential processes, SDE Matching [14], which proposes simulation-free latent SDE training via matching objectives, and Unified Neural SDE frameworks [15], which leverage discretized stochastic dynamics for scalable learning. These methods typically rely on separable drift and diffusion estimation under Gaussian assumptions. Jump-SDEs. Many systems in science, engineering and finance exhibit stochastic dynamics beyond continuous Gaussian perturbations. Neural Jump SDEs [16] and Neural Stochastic Temporal Jump Processes [17] study discontinuous dynamics in the context of temporal point processes and event modelling. Neural Jump ODEs [18] instead model discontinuities arising from irregular observations and filtering problems. More recently, Neural SDEs with jumps [19] have incorporated Poisson jump processes into neural stochastic dynamics. However, existing approaches generally assume fixed jump families or focus on event-driven dynamics, rather than learning general jump-diffusion operators directly from density evolution. Neural Fokker–Planck models. Neural approximations to Fokker–Planck equations have been proposed using normalizing flows [20] and Physics-Informed Neural Networks [21]. To the best of our knowledge, these approaches do not learn stochastic operators directly from trajectory data and are generally limited to continuous diffusion dynamics. Generator learning. The classical machinery of kernel methods and composition operators has also been used for stochastic generator learning. In particular, Extended Dynamic Mode Decomposition [22] and related Koopman-based methods approximate backward generators and transfer operators from data. These approaches may be viewed as adjoint to ours, but typically learn combined operators implicitly and do not directly parameterize advection, diffusion and jump terms with neural representations. Projection filtering and particle methods. Projecting stochastic dynamics onto Gaussian families is a classical idea in nonlinear filtering and stochastic simulation. Projection filtering methods [23, 24, 25] provide the theoretical background for finite-dimensional approximations of infinite-dimensional stochastic dynamics. Related moment-closure methods [26] and Smoothed Particle Hydrodynamics (SPH) [27, 28] similarly approximate evolving densities through moving localized kernels. Our work builds on these ideas to construct neural operator models on projected Gaussian-mixture representations.
3
Figure 2: The forward Kolmogorov semigroup eL on a Gaussian. Advection transports its mean and stretches its covariance, diffusion spreads its covariance, and the jump semigroup splits it.
3
Problem formulation
We consider fully general continuous-time Markov dynamics with infinitely divisible increments; by the Lévy–Khintchine theorem, they take the following canonical representation as a Lévy-Itô SDE: Z dxt = f (xt ) dt + g(xt ) dwt + z J(dt, dz), (1) Rd d
given in terms of a time-indexed state variable xt ∈ R and a Lévy triplet (f, g, h): • The drift f : Rd → Rd , which models smooth, deterministic dynamics in terms of forces; ′ • The noise coefficient g : Rd → Rd×d , which models continuous, small-scale random ′ fluctuations driven by a Brownian motion wt ∈ Rd ; • The jump term z J(dt, dz), that captures discontinuous dynamics induced by a compound jump process J with a state-dependent Lévy measure h : Rd → M(Rd ). (m)
M Our objective is to solve the following inverse problem: Given {xn }N n=1 m=1 , a sequence of M independent realizations of (1) observed over N uniformly distributed discrete time steps 0, s, 2s, . . . , N s, learn a drift fθ : Rd → Rd , a diffusion coefficient ggϑ⊤ : Rd → Rd×d , and a conditional generative model hΘ : Rd → M for the jump distribution, where M := M(Rd ) is the space of measures in Rd . We make the following assumptions: (A1) all terms are continuously differentiable and (A2) uniformly Lipschitz continuous and (A3) all jump measures have bounded moments.
4
The Kolmogorov Forward Equation
The Kolmogorov Forward Equation, also known as the Fokker-Planck Equation, is a linear PDE that describes the evolution of a Lévy SDE at the level of moving probability masses, with individual particles represented as Dirac delta measures and ensembles represented as smooth densities. Take the d−dimensional SDE in eq. 1. The KFE models the probability pt (x) of finding a particle at a given position x ∈ Rd at time t. Given a probability measure p0 ∈ M for the initial position of the particles, the KFE describes the evolution pt ∈ M in terms of a PDE in d spatial dimensions, given by the Kolmogorov Forward generator L : M 7→ M : d pt (x) = Lpt = (A + D + J )pt . (2) dt L describes the infinitesimal time variation of pt ; again by the Lévy-Khintchine Theorem, it may be decomposed similarly to 1 into an advection term A, a diffusion term D, and a jump term J : A pt (x) = −∇ · (f (x)pt (x)) (advection) (3) 1 2 (diffusion) (4) D pt (x) = ∇ : (gg ⊤ (x)pt (x)) Z2 J pt (x) =
pt (x − y)h(x − y, y) − pt (x)h(x, y) dy
(jump)
(5)
Rd
Here, f : Rd → Rd is the drift vector field driving deterministic transport, gg ⊤ := gg ⊤ ∈ Rd×d is 2 the diffusion tensor, and the jump term is governed by h. The operators ∇ and P∇ denote the gradient ⊤ and Hessian operators, with inner product A · B := Tr(A B) and a : b = i,j aij bij . The density at time t is then given by the action of the semigroup eLt : M → M via pt = eLt p0 ; likewise, the fixed-step increments in the data are given by the short-time iterate pt+s = eLs pt . The generator L, on the other hand, is never directly observed and must be estimated. 4
4.1
Lagrangian Galerkin Projection and Gaussian Mixtures
The forward Kolmogorov generators and semigroups are infinite-dimensional objects; still, local probability masses generated by short-time stochastic dynamics are often approximately Gaussian. Inspired by Smoothed Particle Hydrodynamics, we will use a Lagrangian Galerkin discretization to extract a tractable, finite-dimensional projection of the KFE which is amenable to learning. Lagrangian Galerkin methods approximate L via its action on a set of moving basis functions; more specifically, we will use a basis of meshless Gaussian mixtures to represent our evolving densities: K X pt (x) = π (k) Gµ(k) ,Σ(k) (x), (6) t
t
k=1
where each k−indexed component πGµ,Σ is a normal distribution parametrized by its weight π ∈ R, its mean µ ∈ Rd and its covariance Σ ∈ Rd×d . In practice, this representation will reduce the original infinite-dimensional evolution to a corresponding dynamics in the Gaussian parameters [π, µ, Σ], which NKEs will take as inputs to describe the KFE as moving, spreading, and splitting Gaussians. The main tool for this will be approximate closure of Gaussians under L. For sufficiently small s, the image of a localized Gaussian under the advection–diffusion semigroup is well approximated by a Gaussian with transported mean and covariance (see the Appendix for a rigorous statement): eAs+Ds Gµ,Σ ≈ Gµ+s∆µ,Σ+s∆Σ . (7) In contrast, the image of a Gaussian under the jump semigroup is given by the convolution with the jump measure, and cannot generally be described with a single Gaussian. Still, if the jump semigroup eJ s is given by a mixture H s , its action leads to another mixture, as per closure under convolution: eJ s Gµ,Σ = H s ∗ Gµ,Σ , (8) We may formalize this intuition by defining the projection of the dynamics onto the space of Gaussian Mixtures. Let G (k) be the space of K-element Gaussian mixtures, with g ∈ G (k) ⇐⇒ g = K (π (k) , µ(k) , Σ(k) )K and Π†K : G K → M be the projection and lifting k=1 . Let also ΠK : M → G operators between the space of measures and their corresponding mixture representation: K X ΠK : p → (π (k) , µ(k) , Σ(k) )k , Π†K : (π (k) , µ(k) , Σ(k) )k → π (k) Gµ(k) ,Σ(k) . (9) k=1
For K = 1, this is equivalent to moment matching and can be directly calculated: Π1 pt = (1, ⟨x, pt ⟩, ⟨(x − µ)(x − µ)⊤ , pt ⟩), (10) where µ is the mean of pt and 1 denotes unit mass. For larger K, regularization is required to ensure uniqueness and continuity (see Appendix). NKEs will learn the projected generator Lb : G (k) → G (k) : Lb = Π ◦ L ◦ Π† . (11) 4.2
Advection Generator
The action of the advection operator takes a simple form: it transports the mass under the Gaussian according to the drift f . The action of the projected advection operator Ab = Π1 ◦ A ◦ Π†1 follows classically from its adjoint A⋆ , which can be obtained via integration by parts: Z Z ⟨ϕ, Aφ⟩ = − ϕ(x) [∇ · (f (x)φ(x))] dx = [f · ∇ϕ] φ(x) dx = ⟨A⋆ ϕ, φ⟩ (12) Rd
Rd
For the mean, we may take the inner product with the test function ϕ = x: ⟨x, A Gµ,Σ ⟩ = ⟨A⋆ x, Gµ,Σ (x)⟩ = ⟨f (x) , Gµ,Σ ⟩, (13) ⊤ while for the covariance, we take the inner product with the test function ϕ = (x − µ)(x − µ) : (x − µ)(x − µ)⊤ , A Gµ,Σ = A⋆ (x − µ)(x − µ)⊤ , Gµ,Σ = f (x)(x − µ)⊤ + (x − µ)f (x)⊤ , Gµ,Σ . (14) Combined, these expressions lead to full action of the projected advection on a Gaussian’s parameters: Ab (π, µ, Σ) = 0, ⟨f, Gµ,Σ ⟩, (x − µ)f (x)⊤ + f (x)(x − µ)⊤ , Gµ,Σ (15) Intuitively, we have that the mean is transported according to a smoothed vector field and the covariance is stretched and rotated by its divergence. Its total mass does not change. 5
4.3
Diffusion Generator
The action of the diffusion operator likewise admits a simple interpretation: it spreads mass towards b = Π{1} ◦ D ◦ Π† follows local equilibrium, possibly at space-dependent rates. Its projection D {1} analogously from its adjoint, obtained via integration by parts: Z Z 1 2 1 2 ∇ : gg ⊤ (x)φ(x) dx = ∇ ϕ(x) : gg ⊤ (x) φ(x) dx = ⟨D⋆ ϕ, φ⟩. ⟨ϕ, Dφ⟩ = ϕ(x) 2 Rd 2 Rd (16) We may calculate its action on the mean by again testing against ϕ = x: ⟨x, DGµ,Σ ⟩ = ⟨D⋆ x, Gµ,Σ ⟩ = 0, (17) 2 ⊤ since ∇ x = 0. For the covariance, analogous derivation with ϕ = (x − µ)(x − µ) leads to: (x − µ)(x − µ)⊤ , DGµ,Σ = D⋆ (x − µ)(x − µ)⊤ , Gµ,Σ = gg ⊤ (x), Gµ,Σ , (18) leading to the interpretation of the projected diffusion operator as spreading out a Gaussian’s support by adding to its covariance, without any effect on its mean or total mass: b µ, Σ) = 0, 0, gg ⊤ , Gµ,Σ . D(π, (19) 4.4
Jump Semigroup
Unlike advection and diffusion, the action of the jump generator on a Gaussian cannot be generally described in terms of a new Gaussian with perturbed mean and covariance. For small jumps y ≤ ϵ with smooth distributions, a Taylor expansion yields that the cumulative effect of J<ϵ is well approximated by an state-dependent advection-diffusion given by the first two moments of h(see the appendix): 1 (20) J<ϵ pt ≈ −∇ · (β(x) pt (x)) + ∇2 : γ(x) pt (x) , 2 where β and γ stand for the mean and covariance of h. For large jumps, however, this no longer holds. Instead, we will model the full jump semigroup over the time span s. For state-independent h, this may be written in terms of the finite-time Poisson distribution H s : ∞ X si h∗i eJ s pt = H s ∗ pt , where H s := e−λs (21) i! i=0 where λ is the total mass of h, and h∗n denotes n-fold iterated convolution, with h∗0 = δ0 corresponding to trajectories that do not jump during the interval [t, t + s]. For H s also a K-Gaussian mixture (α(k) , β (k) , γ (k) )k , the projected action of the jump semigroup on a Gaussian may be succinctly described in terms of theconvolution between two mixtures: eJ s (π, µ, Σ) = α(k) , β (k) , γ (k)
∗ (π, µ, Σ) = α(k) π, β (k) + µ, γ (k) + Σ
b
k
.
k
(22)
This may be visualized in a simple context as follows: If particles have, say, a 10% chance to jump right by one unit, the finite-time distribution will be modelled as two Gaussians, one with 90% mass at the origin, and another with 10% mass one unit to the right. If the jump distribution is instead space-dependent, we may instead consider H s (x) to have spatially varying parameters [α(k) (x), β (k) (x), γ (k) (x)]k . However, this breaks the convolution structure from above: different points underneath the Gaussian will be exposed to different jump distributions, leading to a largely intractable object. We may recover tractability by leveraging the rapid decay of concentrated Gaussians and the smoothness of the space-varying jump parameters. We assume that the jump measures under the integral are well approximated by their distribution at the mean, around which most of the integral is concentrated: Z Z Gµ,Σ (x − y)h(x − y, y) − Gµ,Σ (x)h(x, dy) ≈ Gµ,Σ (x − y) − Gµ,Σ (x) h(µ, dy) (23) Rd
Rd
This will allow us to approximate the projected semigroup as the convolution with the Gaussian mixture that represents the jumps at the mean µ, H s (µ) = [α(k) (µ), β (k) (µ), γ (k) (µ)]k : b eJ s (π, µ, Σ) ≈ α(k) (µ), β (k) (µ), γ (k) (µ) ∗ (π, µ, Σ) (24) k
For sharp (Σ → 0) Gaussians pt and slowly varying parameters α, β, γ, this approximation is provably sharp (see the Appendix). 6
Figure 3: Neural Kolmogorov Equations. NKEs operate as deterministic Neural ODEs on the finitedimensional (projected) space of Gaussian mixtures and model flows on their mean and covariance.
5
Neural Kolmogorov Equations
Neural Kolmogorov Equations describe SDEs by modelling the associated advection, diffusion, and jump operators in the Kolmogorov formalism; more specifically, NKEs learn the projection of the forward operators on the space of Gaussians, based on the dynamics imposed by the projected Kolmogorov equation on the Gaussian-projected density pb: d pb = Lb pb (25) dt NKEs are composed of three parts, advection, diffusion and jump, which are trained in four steps: 1. Cluster particles into approximate Gaussians and separate jumps. The trajectory data, given in terms of particles, is clustered into empirical/weak approximations for Gaussians. Large jumps are filtered out via a thresholding procedure to allow for separate treatments for the advection-diffusion generator and the jump semigroup. 2. Learn the projected advection generator Abθ : G → G from mean evolution. NKEs approximate the lifted advection generator in terms of a neural drift fθ : Rd → Rd , which plays an analogous role to the drift in Neural SDEs. Abθ transports single localized Gaussians, and is trained by tracking the movement of their means µ over time. bϑ : G → G from covariance evolution. NKEs 3. Learn the projected diffusion generator D approximate the lifted diffusion generator as parametrized in terms of a (squared) noise coefficient bϑ acts on single localized ggϑ⊤ : Rd → Rd×d , analogous to the noise coefficient in Neural SDEs. D Gaussians via spreading, and is trained by tracking the time variation of their covariances. 4. Learn the projected jump semigroup exp(Jbs)Θ : G → G from filtered-out jump increments. b s eJ s is approximated in terms of a conditional generative model HΘ : Rd → (Rd+d×d )k that outputs the parameters for a Gaussian mixture modelling the short-time jump distribution at each Gaussian center µ. It splits Gaussians via convolution, and is trained on filtered-out jump residues. 5.1
Operator splitting
The basis of our strategy will be the interpretation of separable learning as operator splitting. We observe that Euler-Maruyama approximations – which consist on a drift step, followed by the addition of Gaussian noise – correspond to first-order Lie-Trotter [29] splitting for L = A + D + J : (Lie-Trotter) eLs ≈ eJ s eDs eAs . (26) In contrast, NKEs are based on a Strang-type splitting [30], second-order in the absence of jumps: (Strang) eLs ≈ eAs/2 eJ s eDs eAs/2 . (27) This scheme instead resembles higher-order numerical schemes, such as Itô-Heun, composed of half-advection steps before and after the addition of noise. This will allow us to conceive more accurate loss functions by leveraging the invertibility of the advection semigroup, leading to the use of second-order trapezoidal differences in the data. These can then be combined with inner products to obtain the dynamics of the moments of Gaussians: ⟨ϕ(x), e−As/2 pt+1 (x)⟩ ≈ ⟨ϕ(x), eJ s eDs eAs/2 pt ⟩. In fact, operator splitting will allow us to apply and train each component separately. 7
(28)
5.2
Gaussian approximation and jump thresholding
5.2.1
Gaussian approximation
In practice, we do not directly observe the evolution of Gaussian densities during training. Instead, (m) they will be represented empirically as particle clusters: we cluster the particles {xn } into C (c) (c) (c) clusters {G̃n }c,n with {µ̃n } and {Σ̃n } as means and covariances respectively. The clusters can be non-exclusive, and are built from the κnn nearest neighbors of selected particles, leading to weak approximations to an element in our basis. We then track the changes in mean and covariance of these clusters between observations to estimate their time derivative: X X 1 1 (m) ⊤ µ̃(c) x(m) (x(m) − µ̃(c) Σ̃(c) − µ̃(c) (29) n = n , n = n )(xn n ) κnn (m) (c) κnn (m) (c) n xn
∈G̃n
xn
∈G̃n
These clusters act as Dirac delta ensembles, which will appear as discrete sums under inner products. X X G̃(c) δx(m) , =⇒ ⟨ϕ(x), G̃(c) ϕ(x(m) (30) n = n (x)⟩ = n ). n
(m)
xn
(c)
(m)
∈G̃n
xn
(c)
∈G̃n
Intuitively, this strategy may be seen as using these Gaussian-like clusters as local denoised estimators for the dynamics, so that groups of particles move in an approximately deterministic manner. 5.2.2
Jump thresholding
We will treat jumps separately via thresholding. Under sufficiently small timesteps, diffusion-driven increments concentrate around a characteristic local scale, while jump events induce comparatively large displacements. (m)
(m)
(m)
Let ∆xn := ∥xn+1 − xn ∥ denote the magnitude of the observed increment. We flag each increment as jump-contaminated when it exceeds the typical scale of its peers at the same time step, as measured by the median absolute deviation (MAD) [31]. Writing medm for the median across m-indexed particles, we define n o ′′ m′ χ(m) := 1 ∆x(m) − medm ∆x(m) > κsens κGaussian · medm′′ ∆xm , (31) n n n n − medm′ ∆xn with κsens > 0 controlling the detector sensitivity, with κGaussian = 1/Φ−1 (0.75) ≈ 1.4826 the standard Gaussian calibration. Detected jumps (χ = 1) are excluded when learning the advection– diffusion generator and retained when learning the jump semigroup. In the Kolmogorov picture, this jump thresholding implies a decomposition of the Lévy measure h. For any cutoff ε > 0, the Lévy measure admits the canonical decomposition h = h<ε + h≥ε , separating small and large jumps and separating the generator accordingly with J = J<ε + J≥ε . In practice, any jumps too small to be detected will be modelled as effective diffusion, as discussed. 5.3
Advection-diffusion Learning
b defines a deterministic flow on G{1} . From In the absence of jumps, the lifted generator Ab + D the action of the lifted generators on Gaussians (Appendix D), the parameters (µ, Σ) evolve under advection–diffusion as the ODE d b (π, µ, Σ). (πt , µt , Σt ) = (Ab + D) (32) dt This means we can now extract the projected generators L from our data by estimating the time derivative of the mean and covariance or the surrogate Gaussians with finite differences; we may inherit the second-order accuracy of Strang splitting by using a trapezoidal symmetrical scheme. We observe that the means evolve solely due to advection. Strang-splitting training then consists of enforcing a "meet-me-halfway" condition: we advect clusters forward and backward in time, then minimize the mismatch between their means. This leads to the advection loss Ladv for the drift fθ : 2 X (c) s (c) Ladv (θ) = µn+1 − µ(c) ⟨fθ , G̃n+1 ⟩ + ⟨fθ , G̃(c) . (33) n − n ⟩ 2 2 c,n 8
An analogous calculation may be carried out for diffusion. Crucially, having already obtained an estimate fθ from the first-moment dynamics, we may subtract the drift contribution to isolate the contribution of ggϑ⊤ . Again, we approximate these quantities empirically from clusters and match trapezoidal differences: 2 X (c) s A (c) A (c) D (c) D (c) ∆ Σ̃ + ∆ Σ̃ + ∆ Σ̃ + ∆ , (34) Ldiff (ϑ) = Σ̃n+1 − Σ̃(c) − Σ̃ n n n n+1 n+1 2 F c,n where E D (c) ⊤ (c) ⊤ (c) , ∆A Σ̃(c) = (x − µ )f (x) + f (x)(x − µ ) , G̃ θ θ n n n n
D E ⊤ (c) ∆D Σ̃(c) = gg (x), G̃ (x) . n ϑ n
(35) Note that both training objectives are parallelizable over both particles and time-steps; by precomputing clusters, the cost of evaluating this loss becomes not substantially larger than that of computing a corresponding EM loss. Note also that this procedure allows for general diffusion terms with off-diagonal and cross-effect terms, unlike EM-based methods. In particular, because no maximum-likelihood estimation is necessary, learning coupled and non-diagonal covariances becomes simple. 5.4
Jump learning
We learn the finite-time jump law H s directly from thresholded jump increments. For sufficiently small s, retained increments √ are assumed jump-dominated, since diffusion and drift contributions scale respectively as O( s) and O(s). We define the time-s increment for the projected KFE: b k πt+s , µkt+s , Σkt+s k = eJ s (πt , µt , Σt ). (36) We may then fit these increments with Gaussian mixtures to learn the projected semigroup. Note that approximating the entire increment as a jump is only accurate to leading order. We start by defining the jump targets y, which consist of a particle’s future position if it was flagged as a jump, or its current position if otherwise. By doing this, we include both jump and not jump increments in order to model the full semigroup, implicitly modelling also the jump probability: (m)
(m) yn(m) = x(m) + χ(m) n n (xn+1 − xn )
e (c) We then cluster particles into localized Gaussian packets G n .
(37)
For each cluster, we define the
associated target cloud n o e (c) Ỹn(c) := yn(m) : x(m) ∈G . n n
(38)
This target cloud is interpreted as an empirical sample from the action of the jump semigroup on the source cluster: s (c) e (c) Ỹn(c) ∼ eJ s G (39) (c) e (c) n ≈ H (·|µn ) ∗ Gµ e ,Σ n
H s (·|µ(c) (c) e (c) = n ) ∗ Gµ e ,Σ n
KJ X
n
n
αj (µ(c) n ) Gµ(c) +βj (µ(c) ), Σ(c) +γj (µ(c) ) . n
n
n
n
(40)
j=1
Since Gaussian convolution adds means and covariances, recovering the jump law reduces to local s deconvolution. We parameterize HΘ : Rd → (Rd+d×d )KJ with a Mixture Density Network (MDN) [32], which outputs the weights, means and covariances of a conditional Gaussian mixture: s HΘ (·|x) =
KJ X
αj (x) Gβj (x),Γj (x) .
(41)
j=1
The network is trained by maximizing the conditional log-likelihood of the target clouds at their respective cluster centers: h i X X s Ljump (Θ) = − log HΘ (·|µ(c) (42) (c) e (c) (y). n ) ∗ Gµ e ,Σ n
n
n,c y∈Y (c) n
The convolution with the source cluster can be evaluated analytically by addition of means and covariances. Unlike classical expectation-maximization methods for Gaussian-mixture fitting, this objective is fully differentiable and naturally accommodates state-dependent conditioning within end-to-end gradient-based training. 9
Figure 4: Advection-diffusion learning. Euler-Maruyama methods attempt to directly match the movement of particles between observations as drift; any remaining residue is modelled as noise. In contrast, Kolmogorov methods model the movement of particle clusters, modelling the difference in their means as advection and the remaining difference in their covariances as diffusion.
Figure 5: Jump learning. NKEs cluster jump origins, then model their destinations with a Gaussian mixture. The result should compose the jump distribution with the cluster’s own covariance.
6
Experiments
To validate NKEs as a conceptual framework for SDE learning, we evaluate them on a set of small and large-scale SDE-learning benchmarks in the field of scientific machine learning. For the small synthetic benchmarks, we combine three types of noise – Geometric Brownian motion, coupled Brownian motion, and jump diffusion – with three canonical multidimensional ODE systems – the chaotic Lorenz oscillator, the Black-Scholes model, and the multidimensional double-well model. For the large-scale experiments, we tackle two coarse-graining problems from molecular dynamics, in which a heavy tracked particle moves in a 3D particle bath. Each model is trained on a set of 1024 trajectory observations. After training, 256 trajectories are generated then evaluated against a test set of 256 additional observations via three metrics: the mean squared error in the mean vector at each time, mean squared error in the covariance matrix at each time, and mean (dimensionwise) Wasserstein distance at each time. We compare NKEs to three baselines: • Parallelizable Euler-Maruyama models, based on the classical EM scheme as implemented in [33]; these are equipped with a drift network and a diagonal noise network; • Parallelizable Trajectory Flow Matching models, based on the implementation in [13]; they are equipped with a drift network and a scalar uncertainty network, used for the noise; • Autoregressive SDE-GAN models, inspired by the implementation in [7]; these are equipped with a generator Neural SDE a Neural CDE discriminator. Training is performed using the AdamW optimizer [34] for all baselines. Each experiment was run on five seeds, and the mean and standard deviation of each metric was reported. The appendix contains detailed specifications for each experiment. Due to the presence of jumps in some of the benchmarks, the baseline models had to be adapted from their original contexts; we provide implementation-specific details in the Appendix. 10
7
Analysis and limitations
The results show that NKEs achieve consistently strong performance across the experiments, displaying generally superior accuracy when compared to the baseline methods, while maintaining competitive training times. In particular, NKEs deal with general forms of driving noise, including geometric and coupled Brownian motions, as well as diverse jump processes, and do not require parametric assumptions on its form. This provides a unified treatment of Gaussian jumps and Poisson processes within the same generative framework. NKEs also provide a more accurate framework for simulation-free Neural SDEs and EM methods: Instead of tracking noisy individual trajectories, NKEs model entire clusters of neighboring particles as probability masses, whose approximately deterministic dynamics may be learned more easily. While NKEs offer compelling advantages when learning drift and noise from data, it has limitations. First is the requirement for constant sampling rates in the presence of jumps: NKEs’ noise learning step learns a fixed iteration of the semigroup eJ s , restricting the method to fixed sampling rates. Extending jump-learning to general non-uniform sampling would require a tractable representation of the jump generator J , which seems out of reach for our Lagrangian-GMM representations, but could possible be achieve in other bases. Still, in the absence of jumps, NKEs may be trivially extended to non-constant sampling rates. The second limitation lies in the Gaussian mixture modelling of distributions. The dimensional limitations of GMMs make Kolmogorov models as implemented here not adequate for systems in extremely high dimensions, as is the case with generative models. Likewise, while GMMs intuitively extend the range of operators we can model in the projected manifold, in practice the projection operator for K > 1 is often ill-posed, leading to non-uniqueness in representation. When coupled with state-dependent distributions, we observe that this leads to discontinuities and poor convergence when learning complex jump patterns. Overcoming this issue would require regularization of the GMM models across the state space. Table 1: Lorenz oscillator with diagonal geometric Brownian noise. Method TFM SDE-GAN Euler-Maruyama Kolmogorov (ours)
Mean MSE ↓
Covariance MSE ↓
Wasserstein ↓
TT ↓
1.8111 ± 0.0401 2.6020 ± 0.2698 1.8027 ± 0.0113 0.0471 ± 0.0016
0.7516 ± 0.0580 15.7637 ± 0.0787 0.9494 ± 0.0451 0.0755 ± 0.0089
0.3294 ± 0.0454 0.511 ± 0.0049 0.3065 ± 0.0066 0.0391 ± 0.0032
39.1 ± 1.3 2584 ± 65 33.9 ± 0.9 42.4 ± 0.7
Table 2: Multivariate Black-Scholes with coupled Geometric Brownian Motion. Method TFM SDE-GAN Euler-Maruyama Kolmogorov (ours)
Mean MSE ↓
Covariance MSE ↓
Wasserstein ↓
TT ↓
0.0824 ± 0.0123 0.4873 ± 0.0211 0.1951 ± 0.0108 0.0685 ± 0.0104
1.8421 ± 0.0612 1.9923 ± 0.1021 0.9732 ± 0.0489 0.0523 ± 0.0061
0.2145 ± 0.0187 0.412 ± 0.0038 0.2261 ± 0.0105 0.0997 ± 0.025
41.3 ± 1.5 2490 ± 70 35.7 ± 1.1 44.1 ± 0.9
Table 3: Multivariate double-well system with constant diffusion and Gaussian distributed jumps. Method TFM SDE-GAN Euler-Maruyama Kolmogorov (ours)
Mean MSE ↓
Covariance MSE ↓
Wasserstein ↓
TT (s) ↓
0.0063 ± 0.0034 0.8334 ± 0.8334 0.0113 ± 0.0057 0.0062 ± 0.0021
0.0276 ± 0.0036 0.6720 ± 0.0360 0.0160 ± 0.0021 0.0085 ± 0.0057
0.1809 ± 0.0083 0.4091 ± 0.1244 0.1907 ± 0.0155 0.109 ± 0.0084
1525 ± 63 – 1152 ± 47 1545 ± 11
Table 4: Rarefied coarse-graining benchmark. ‘–’ indicates no meaningful convergence. Method TFM SDE-GAN Euler-Maruyama Kolmogorov (ours)
Mean MSE ↓
Covariance MSE ↓
Wasserstein ↓
TT (s) ↓
1.0277 ± 0.0916 – 0.0423 ± 0.0036 0.0343 ± 0.0132
1.1150 ± 2.6132 – 1.1155 ± 0.2612 0.1365 ± 0.0021
0.5121 ± 0.0024 – 0.3121 ± 0.0024 0.2321 ± 0.0154
270.5 ± 4.7 – 206.3 ± 3.4 323.1 ± 25.3
11
(a) Ground truth.
(b) Euler-Maruyama.
(c) Kolmogorov.
Figure 6: Results for the Lorenz system.
(a) Ground truth.
(b) Euler-Maruyama.
(c) Kolmogorov.
Figure 7: Results for the Black-Scholes system.
(a) Ground truth.
(b) Kolmogorov.
(c) Euler-Maruyama.
Figure 8: Results for the Double well system.
(a) Ground truth.
(b) Euler-Maruyama.
(c) Kolmogorov.
Figure 9: Results for the rarefied coarse-graining system (positions only).
12
References [1] Ricky T. Q. Chen, Yulia Rubanova, Jesse Bettencourt, and David Duvenaud. Neural ordinary differential equations. In Advances in Neural Information Processing Systems, volume 31, 2018. [2] Christopher Rackauckas, Yingbo Ma, Julius Martensen, Collin Warner, Kirill Zubov, Rohit Supekar, Dominic Skinner, Ali Ramadhan, and Alan Edelman. Universal differential equations for scientific machine learning. arXiv preprint arXiv:2001.04385, 2020. [3] Samuel N. Cohen, Christoph Reisinger, and Sheng Wang. Arbitrage-free neural-SDE market models. Applied Mathematical Finance, 30(1):1–46, 2023. [4] Chieh-Hsin Lai, Yuhta Takida, Naoki Murata, Toshimitsu Uesaka, Yuki Mitsufuji, and Stefano Ermon. FP-Diffusion: Improving score-based diffusion models by enforcing the underlying score Fokker–Planck equation. In Proceedings of the 40th International Conference on Machine Learning, 2023. [5] Patrick Kidger, James Morrill, James Foster, and Terry Lyons. Neural controlled differential equations for irregular time series. In Advances in Neural Information Processing Systems, volume 33, 2020. [6] Xuechen Li, Ting-Kam Leonard Wong, Ricky T. Q. Chen, and David Duvenaud. Scalable gradients for stochastic differential equations. In Proceedings of the 23rd International Conference on Artificial Intelligence and Statistics, volume 108 of Proceedings of Machine Learning Research, 2020. [7] Patrick Kidger, James Foster, Xuechen Li, Harald Oberhauser, and Terry Lyons. Neural SDEs as infinite-dimensional GANs. In Proceedings of the 38th International Conference on Machine Learning, volume 139 of Proceedings of Machine Learning Research, 2021. [8] Will Grathwohl, Ricky T. Q. Chen, Jesse Bettencourt, Ilya Sutskever, and David Duvenaud. FFJORD: Free-form continuous dynamics for scalable reversible generative models. In International Conference on Learning Representations, 2019. [9] Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models. In Advances in Neural Information Processing Systems, volume 33, pages 6840–6851, 2020. [10] Yang Song, Jascha Sohl-Dickstein, Diederik P. Kingma, Abhishek Kumar, Stefano Ermon, and Ben Poole. Score-based generative modeling through stochastic differential equations. In International Conference on Learning Representations, 2021. [11] Yaron Lipman, Ricky T. Q. Chen, Heli Ben-Hamu, Maximilian Nickel, and Matt Le. Flow matching for generative modeling. In International Conference on Learning Representations, 2023. [12] Peter E. Kloeden and Eckhard Platen. Numerical Solution of Stochastic Differential Equations. Springer, 1992. [13] Xi Zhang, Yuan Pu, Yuki Kawamura, Andrew Loza, Yoshua Bengio, Dennis L Shung, and Alexander Tong. Trajectory flow matching with applications to clinical time series modelling. Advances in Neural Information Processing Systems, 37:107198–107224, 2024. [14] Grigory Bartosh, Dmitry Vetrov, and Christian A. Naesseth. SDE matching: Scalable and simulation-free training of latent stochastic differential equations. In Proceedings of the 42nd International Conference on Machine Learning, volume 267 of Proceedings of Machine Learning Research, pages 3054–3070, 2025. [15] Xuanqing Liu, Tesi Xiao, Si Si, Qin Cao, Sanjiv Kumar, and Cho-Jui Hsieh. Neural SDE: Stabilizing neural ODE networks with stochastic noise. arXiv preprint arXiv:1906.02355, 2019. [16] Junteng Jia and Austin R. Benson. Neural jump stochastic differential equations. In Advances in Neural Information Processing Systems, volume 32, 2019. 13
[17] Shuai Zhang, Chuan Zhou, Yang Aron Liu, Peng Zhang, Xixun Lin, and Zhi-Ming Ma. Neural jump-diffusion temporal point processes. In Proceedings of the 41st International Conference on Machine Learning, volume 235 of Proceedings of Machine Learning Research, pages 60541–60557, 2024. [18] Yulia Rubanova, Ricky T. Q. Chen, and David Duvenaud. Latent ODEs for irregularly-sampled time series. In Advances in Neural Information Processing Systems, volume 32, 2019. [19] Qi Rong, Liang Yan, and Xiaojun Duan. Convergence and stability analysis of neural stochastic differential equations with poisson jumps. Chaos, Solitons & Fractals, 2026. [20] Danilo Rezende and Shakir Mohamed. Variational inference with normalizing flows. In Proceedings of the 32nd International Conference on Machine Learning, volume 37 of Proceedings of Machine Learning Research, 2015. [21] Maziar Raissi, Paris Perdikaris, and George E. Karniadakis. Physics-informed neural networks: A deep learning framework for solving forward and inverse problems involving nonlinear partial differential equations. Journal of Computational Physics, 378:686–707, 2019. [22] Matthew O. Williams, Ioannis G. Kevrekidis, and Clarence W. Rowley. A data-driven approximation of the koopman operator: Extending dynamic mode decomposition. Journal of Nonlinear Science, 25(6):1307–1346, 2015. [23] Damiano Brigo, Bernard Hanzon, and François Le Gland. A differential geometric approach to nonlinear filtering: The projection filter. IEEE Transactions on Automatic Control, 43(2):247– 252, 1998. [24] Damiano Brigo, Bernard Hanzon, and François Le Gland. Approximate nonlinear filtering by projection on exponential manifolds of densities. Bernoulli, 5(3):495–534, 1999. [25] Bernard Hanzon and Remi L. M. Hut. New results on the projection filter. In Proceedings of the European Control Conference, pages 623–628, Grenoble, France, 1991. [26] Mario Di Paola and Alba Sofi. Approximate solution of the Fokker–Planck–Kolmogorov equation. Probabilistic Engineering Mechanics, 17(4):369–384, 2002. [27] Robert A. Gingold and Joseph J. Monaghan. Smoothed particle hydrodynamics: Theory and application to non-spherical stars. Monthly Notices of the Royal Astronomical Society, 181(3):375–389, 1977. [28] Leon B. Lucy. A numerical approach to the testing of the fission hypothesis. Astronomical Journal, 82:1013–1024, 1977. [29] H. F. Trotter. On the product of semi-groups of operators. Proceedings of the American Mathematical Society, 10(4):545–551, 1959. [30] Gilbert Strang. On the construction and comparison of difference schemes. SIAM Journal on Numerical Analysis, 5(3):506–517, 1968. [31] Peter J. Rousseeuw and Christophe Croux. Alternatives to the median absolute deviation. Journal of the American Statistical Association, 88(424):1273–1283, 1993. [32] Christopher M. Bishop. Mixture density networks. Technical Report NCRG/94/004, Aston University, 1994. [33] Macheng Shen and Chen Cheng. Neural sdes as a unified approach to continuous-domain sequence modeling. arXiv preprint arXiv:2501.18871, 2025. [34] Ilya Loshchilov and Frank Hutter. Decoupled weight decay regularization. arXiv preprint arXiv:1711.05101, 2017. [35] Shun-ichi Amari and Hiroshi Nagaoka. Methods of Information Geometry, volume 191 of Translations of Mathematical Monographs. American Mathematical Society and Oxford University Press, 2000. 14
[36] Nhat Ho and XuanLong Nguyen. Singularity structures and impacts on parameter estimation in finite mixtures of distributions. SIAM Journal on Mathematics of Data Science, 1(4):730–758, 2019. [37] Arthur P. Dempster, Nan M. Laird, and Donald B. Rubin. Maximum likelihood from incomplete data via the EM algorithm. Journal of the Royal Statistical Society: Series B, 39(1):1–38, 1977. [38] Radek Erban. From molecular dynamics to brownian dynamics. Proceedings. Mathematical, Physical, and Engineering Sciences/The Royal Society, 470(2167):20140036, 2014.
15
A
Symbols and notation
The paper alternates between the microscopic SDE picture, given in terms of particles and their trajectories, and the macroscopic Kolmogorov picture, given in terms of densities and operators. We provie a table describing the notation and the corresponding objects across pictures. Whenever possible, we have attempted to separate the notation as follows: • Lower-case latin letters refer to particles and terms in the SDE picture; • Calligraphic letters refer to abstract spaces and operators in the Kolmogorov picture; • Greek characters refer to neural-network or Gaussian mixture parameters; • Tildes refer to sample-based/empirical objects; • Hats refer to projections onto the space of Gaussian Mixtures. Table 5: Correspondence between the particle-level and density-level formulations. Particle / Path Level Density / Operator Level State space Rd State trajectory xt Discretely observed trajectory xn Drift f (x) Noise coefficient g(x) Jump distribution h(x) Particle increment xn+1 − xn (m) Single particle xn
Space of measures M(Rd ) Density pt (x) Discretely observed density pns (x) Advection generator A Diffusion generator D Jump generator J Infinitesimal generator action (eLs −1)pns Delta measure δx(m) n
{c}
Local particle cluster G̃n {c} {c} Cluster mean/covariance µ̃n , Σn s Jump finite-time distribution H Jump thresholding ∥xn+1 − xn ∥ > τ Learned drift fθ Learned diffusion gϑ2 s Learned finite-time jump distribution HΘ
Gaussian packet Gµ,Σ Gaussian parameters µ, Σ Jump semigroup eJ s Lévy decomposition J<ε + J≥ε Learned lifted advection generator Abθ bϑ Learned lifted diffusion generator D ˆ s Learned lifted jump semigroup eJ Θ
Additionally, we have the following operations: • The divergence ∇· refers to the sum: ∇ · a(x) =
d X d ai (x), dx i i=1
where ai and xi loop over the vector positions; • The operator ∇ : refers to the sum: ∇ : a(x) =
d X
d2 aij (x), dxi dxj i,j=1
where aij loops over matrix entries; • The inner product ⟨ϕ, φ⟩ refers to the classical inner product in integral form: Z ⟨ϕ, φ⟩ = ϕ(x)φ(x)dx; Rd ⋆
• The superscript refers to an operator’s classical adjoint ⟨ϕ, Lφ⟩ = ⟨L⋆ ϕ, φ⟩.
16
B
Projection onto the Gaussian Mixture Manifold
We study the projection operator ΠK : M(Rd ) → GK , which maps probability densities onto the family of K-component Gaussian mixtures. We characterize the geometry of this projection, discuss uniqueness and non-uniqueness properties, and connect them to the lifted dynamics used throughout NKEs. B.1
The Gaussian Mixture Family
A Gaussian mixture with K components takes the form pλ (x) =
K X
K X
π (k) Gµ(k) ,Σ(k) (x),
k=1
π (k) = 1,
π (k) > 0,
(43)
k=1
with parameters K λ = π k , µk , Σk k=1 . The corresponding family is GK := {pλ : λ ∈ ΛK }, where
(44)
d ΛK = ∆K−1 × (Rd × S++ )K ,
d with ∆K−1 the K − 1 simplex and S++ the space of positive definite matrices. For K = 1, G1 forms a smooth statistical manifold under the Fisher–Rao metric [35]. In particular, the Fisher information is strictly positive definite: −1 Σ 0 I(µ, Σ) = . (45) 1 −1 0 ⊗ Σ−1 2Σ
For K > 1, the geometry becomes singular whenever: 1. two Gaussian components coincide, 2. or a mixture weight vanishes. Consequently, GK becomes a stratified space rather than a smooth manifold [36]. B.2
Projection Operator
We will define Π in terms a divergence D on M(Rd ): ΠK (p) := arg min D(p∥q). q∈GK
Choosing D = KL yields the population maximum-likelihood projection: Z ΠKL (p) = arg max p(x) log q(x) dx. K q∈GK
(46)
(47)
In practice, this corresponds to Gaussian-mixture fitting through expectation-maximization (EM) [37]. B.2.1
Uniqueness for K = 1.
Proposition 1. For any density p with finite second moments, the KL projection onto G1 is uniquely given by moment matching: ΠKL (48) 1 (p) = Gµ∗ ,Σ∗ , where Z Z µ∗ = x p(x) dx, Σ∗ = (x − µ∗ )(x − µ∗ )⊤ p(x) dx. (49) Proof sketch. The KL divergence KL(p∥Gµ,Σ ) is strictly convex in (µ, Σ) because the Fisher information on G1 is positive definite. Differentiating with respect to µ and Σ yields the classical moment-matching conditions, thus clearly unique. 17
This uniqueness is fundamental for the advection and diffusion losses, which only rely on singleGaussian projections. B.3
Non-Uniqueness for K > 1
For K > 1, the projection is generally not unique at the level of mixture parameters. This nonuniqueness is intrinsic to the Gaussian-mixture representation itself and does not depend on the particular optimization algorithm used to compute the projection. Proposition 2 (Non-uniqueness of Gaussian-mixture projections). Let K > 1 and define the projection ΠK (p) := arg min D(p∥q), q∈GK
where D is any divergence that depends only on the represented density q. Then the projection is generally non-unique in parameter space. In particular: 1. Permutation symmetry. If λ⋆ is an optimal set of mixture parameters, then any permutation of its Gaussian components yields the same density and is therefore also optimal. 2. Redundant representations. If p ∈ GK ⋆ for some K ⋆ < K, then infinitely many distinct parameter vectors in GK represent the same density and attain the same projection error. Proof sketch. Permuting Gaussian components leaves the mixture density unchanged. Therefore, if p
λ⋆
=
K X
π (k) Gµ(k) ,Σ(k) ,
k=1
is optimal, then any relabeling of the components produces the same density and achieves the same objective value. For the second claim, suppose a component πGµ,Σ appears in the mixture. For any α ∈ (0, 1), πGµ,Σ = απGµ,Σ + (1 − α)πGµ,Σ . Thus a single component may be replaced by two identical components whose weights sum to the original weight without changing the represented density. Repeating this construction produces infinitely many distinct parameter vectors corresponding to the same density. The non-uniqueness described above is structural and reflects the non-identifiability of Gaussianmixture parameterizations. Additional optimization-related difficulties may arise because the projection objective is generally non-convex and may possess singular or degenerate stationary points. B.3.1
Practical Projection in NKEs
The non-uniqueness of these projection will only be an issue when estimating state-depedent jump distributions, as in this case the non-uniqueness of the projection may lead to discontinuity. NKEs use two distinct projection mechanisms. e {c} Advection–diffusion projection (K = 1). Empirical clusters G n are projected through moment matching onto Gaussian packets Gµe{c} ,Σ e {c} . n
n
This projection is unique and fully explicit. Jump projection (K ≥ 1). Finite-time jump laws H s (·|x) are approximated through conditional s Gaussian mixtures parameterized by HΘ (·|x). We expect the network’s smoothness to act as implicit regularization. To stabilize training, we furthermore: 1. use small mixture counts, 2. regularize covariances, 3. prune low-weight components, 4. enforce spatial smoothness. 18
C
Lifted Dynamics
This appendix provides theoretical justification for the approximations underlying Neural Kolmogorov Equations. We show that: 1. Localized Gaussian packets remain approximately Gaussian under short-time advection– diffusion evolution; 2. Sufficiently small jumps behave as an effective diffusion process; 3. State-dependent jump distributions may be frozen at the packet center with controlled error. Throughout this section, we assume f ∈ Cb3 (Rd , Rd ),
a := gg ⊤ ∈ Cb2 (Rd , Rd×d ),
and that all jump measures possess finite third moments. We denote by L=A+D+J the Kolmogorov generator, 1 Ap = −∇ · (f p), Dp = ∇2 : (ap), 2 Z J p(x) = p(x − y)h(x − y, y) − p(x)h(x, y) dy. Rd
C.1
Short-time Gaussianity
We now justify the Gaussian lifting procedure used throughout the paper, in particular the claim that Gaussianity is preserved to second order in time. Let p0 = Gµ,Σ be a Gaussian packet and define ps = e(A+D)s p0 . The projected mean and covariance increments are ∆µ = ⟨f, p0 ⟩, and
D E ∆Σ = (x − µ)f ⊤ + f (x − µ)⊤ + a, p0 .
Let qs = Gµ+s∆µ,Σ+s∆Σ denote the Gaussian obtained by evolving only the first two moments. Lemma 3 (Growth of higher cumulants). Let pt = e(A+D)t Gµ,Σ , with Ap = −∇ · (f p),
Dp =
1 2 ∇ : (ap), 2
a := gg ⊤ .
Assume f ∈ Cb3 (Rd , Rd ) and a ∈ Cb2 (Rd , Rd×d ). Let κα (t) denote the cumulant of pt associated with a multi-index α with |α| ≥ 3. Then, for sufficiently small t, κα (t) = O(t). Proof sketch. We work entirely in weak form. For any smooth test function φ, d ⟨φ, pt ⟩ = ⟨φ, (A + D)pt ⟩ = ⟨(A⋆ + D⋆ )φ, pt ⟩, dt 19
where A⋆ φ = f · ∇φ,
D⋆ φ =
1 a : ∇2 φ. 2
Let mα (t) := ⟨xα , pt ⟩ be the raw moment associated with the multi-index α. Taking φ(x) = xα gives d 1 α 2 α mα (t) = f · ∇x + a : ∇ x , pt . dt 2 Since xα is a polynomial, ∇xα and ∇2 xα are polynomials of degrees |α| −1 and |α| −2, respectively. Under the assumed boundedness of f and a, the right-hand side is bounded on a short time interval by a constant depending only on finitely many moments of pt . Since p0 = Gµ,Σ has finite moments of all orders, standard moment estimates imply that d mα (t) = O(1), dt and therefore mα (t) = mα (0) + O(t). Now let κα (t) be the corresponding cumulant. Cumulants are polynomial functions of raw moments up to order |α|: κα (t) = Pα {mβ (t) : |β| ≤ |α|} , for a universal polynomial Pα . Since each raw moment satisfies mβ (t) = mβ (0) + O(t), it follows that κα (t) = κα (0) + O(t). Because the initial packet p0 = Gµ,Σ is Gaussian, all cumulants of order |α| ≥ 3 vanish: κα (0) = 0. Hence κα (t) = O(t), as claimed. Therefore a localized Gaussian packet remains approximately Gaussian under short-time advection– diffusion evolution, with second-order error in relative entropy. C.2
Small-jump approximation
We now justify the approximation of sufficiently small jumps by an effective advection–diffusion process. For a state-dependent jump kernel h(x, dy), the backward (adjoint) jump generator acting on test functions is: Z J ⋆ ϕ(x) = ϕ(x + y) − ϕ(x) h(x, dy). (50) Rd
Equivalently, J admits the gain–loss representation Z
Z p(x − y) h(x − y, dy) − p(x)
(J p)(x) = Rd
h(x, dy).
(51)
Rd
The first term represents mass arriving at x from points x − y, while the second represents mass leaving x through jumps to other locations. Restricting attention to jumps smaller than a threshold ε > 0, define Z ⋆ J<ε ϕ(x) = ϕ(x + y) − ϕ(x) h(x, dy). (52) |y|<ε
20
Proposition 4 (Small-jump diffusion approximation). Suppose ϕ ∈ C 3 (Rd ). Define Z β(x) = y h(x, dy),
(53)
|y|<ε
and
Z
yy ⊤ h(x, dy).
γ(x) =
(54)
|y|<ε
Then 1 ⋆ J<ε ϕ = β · ∇ϕ + Γ : ∇2 ϕ + R⋆ε ϕ, 2
(55)
where |R⋆ε ϕ(x)| ≤ C∥∇3 ϕ∥∞
Z
|y|3 h(x, dy).
(56)
|y|<ε
Proof. Applying Taylor’s theorem, 1 ϕ(x + y) = ϕ(x) + y ⊤ ∇ϕ(x) + y ⊤ ∇2 ϕ(x)y + R3 (x, y), 2
(57)
with
|R3 (x, y)| ≤ C|y|3 ∥∇3 ϕ∥∞ . ⋆ and integrating term-by-term yields Substituting into J<ε 1 ⋆ J<ε ϕ = β · ∇ϕ + γ : ∇2 ϕ + R⋆ε ϕ. 2
(58)
(59)
The remainder estimate follows immediately. Passing to the forward operator by duality we get the usual advection-diffusion form: 1 J<ε p = −∇ · (βp) + ∇2 : (γp) 2
(60)
Consequently, sufficiently small jumps behave, to second order in jump size, as an effective advection– diffusion process with drift coefficient β and diffusion tensor Γ. This recovers the classical diffusion approximation underlying the Lévy–Khintchine representation and the construction of Lévy-type diffusions. C.3
Jump freezing
Proposition 5 (Freezing the finite-time jump operator). Fix s > 0. Let Hs (x, y) denote the density of the finite-time jump increment y ∈ Rd associated with a particle located at x ∈ Rd . Define the state-dependent jump operator Z (Hs p)(z) := Hs (x, z − x)p(x) dx, Rd
and the operator obtained by freezing the jump law at µ, Z (Hsµ p)(z) := Hs (µ, z − x)p(x) dx = Hs (µ, ·) ∗ p (z). Rd
Assume that Hs is Lipschitz continuous in its spatial argument in L1 , namely that there exists Ls > 0 such that Z |Hs (x, y) − Hs (x′ , y)| dy ≤ Ls ∥x − x′ ∥ Rd
21
for every x, x′ ∈ Rd . Then
√ ∥Hs Gµ,Σ − Hsµ Gµ,Σ ∥L1 ≤ Ls tr Σ.
Consequently, in L1 (Rd )
Hs Gµ,Σ − Hs (µ, ·) ∗ Gµ,Σ −→ 0 as tr Σ → 0.
Proof. By the definitions of Hs and Hsµ , ∥Hs Gµ,Σ − Hsµ Gµ,Σ ∥L1 Z Z [Hs (x, z − x) − Hs (µ, z − x)] Gµ,Σ (x) dx dz. = Rd
Rd
Using the triangle inequality and exchanging the order of integration, ∥Hs Gµ,Σ − Hsµ Gµ,Σ ∥L1 Z Z ≤ Gµ,Σ (x) |Hs (x, z − x) − Hs (µ, z − x)| dz dx. Rd
Rd
Under the change of variables y = z − x, this becomes ∥Hs Gµ,Σ − Hsµ Gµ,Σ ∥L1 Z Z ≤ Gµ,Σ (x) |Hs (x, y) − Hs (µ, y)| dy dx. Rd
Rd
The Lipschitz assumption therefore gives ∥Hs Gµ,Σ − Hsµ Gµ,Σ ∥L1 ≤ Ls By the Cauchy–Schwarz inequality, Z Z ∥x − µ∥Gµ,Σ (x) dx ≤ Rd
Z ∥x − µ∥Gµ,Σ (x) dx. Rd
1/2 ∥x − µ∥ Gµ,Σ (x) dx 2
Rd
Z ×
1/2 Gµ,Σ (x) dx .
Rd
Since Gµ,Σ has unit mass and Z
∥x − µ∥2 Gµ,Σ (x) dx = tr Σ,
Rd
we obtain
D
√ ∥Hs Gµ,Σ − Hsµ Gµ,Σ ∥L1 ≤ Ls tr Σ.
Experimental setup
Across all experiments we test the ability of neural networks to recover the deterministic and stochastic components of the Lévy-Itô SDE given in 1. Every experiment follows the same pipeline: we generate a training and a test ensemble of the reference process on a fine time grid (Euler– Maruyama integration with step dt, states saved every ∆t); subsample in time; then apply a z-score normalization X̃ = (X − µ)/σ. Unless noted otherwise we use Ntrain = 1024 and Ntest = 256 trajectories, integration step dt = 10−3 , and each model is trained across 5 random seeds. D.1
Lorenz system
This benchmark probes chaotic, strongly nonlinear drift under multiplicative isotropic noise. The state X = (x, y, z) ∈ R3 evolves as ẋ = σL (y − x), ẏ = x(ρ − z) − y, ż = xy − βz, g = σ diag([x, y, z]), (61) 8 with the classical parameters (σL , ρ, β) = (10, 28, 3 ) and noise scale σ = 3. Trajectories are simulated on t ∈ [0, 2] with save interval ∆t = 0.02 from standard-normal initial conditions X0 ∼ N (0, I3 ). 22
D.2
Multivariate Black–Scholes
This benchmark exercises state-dependent, correlated (multiplicative) noise for n = 3 coupled assets S ∈ R3>0 : X dSi = µi Si dt + Si σi Lij dWj , aij (S) = σi ρij σj Si Sj , (62) j
where a = gg ⊤ is the state-dependent diffusion matrix and L is the Cholesky factor of the asset correlation matrix ρ. We use drift rates µ = (0.05, 0.06, 0.04), volatilities σ = (0.20, 0.15, 0.25), and correlations ρ12 = ρ23 = 0.66, ρ13 = 0.33. Log-normal initial prices S0 = exp(0.1 N (0, I3 )) are integrated on t ∈ [0, 2] with ∆t = 0.02; learning is performed in raw (unnormalized) coordinates. D.3
Three-dimensional double-well jump-diffusion
As a test problem combining metastability with discontinuous forcing, we consider a threedimensional gradient system whose first coordinate is bistable and whose remaining coordinates are linearly confined. The dynamics follow the jump-diffusion stochastic differential equation dXt = f (Xt ) dt + σ dWt + J dNt , f1 (x) = x1 1 − x21 , fi (x) = −c xi (i = 2, 3), (63) 2 P which is the gradient of the potential V (x) = 41 x21 − 1 + 2c i≥2 x2i . The drift f has two stable fixed points at (±1, 0, 0) separated by an unstable saddle at the origin, so the first coordinate x1 constitutes the “axis” of the double well while x2 , x3 relax toward it with confinement strength c. The state is excited continuously by isotropic Gaussian noise, with Wt a three-dimensional Brownian motion of per-component strength σ. Superimposed on this continuous diffusion is a compound-Poisson jump term J dNt that acts only along the double-well axis: Nt is a Poisson process of rate λ, and at each event the first coordinate receives a symmetric Gaussian kick J ∼ N (0, σJ2 ), i.e. x1 7→ x1 + J, with the transverse coordinates left unchanged. These kicks intermittently drive the particle over the potential barrier, inducing well-to-well transitions that the pure-diffusion component would produce only rarely, and endow the observed increments with a non-Gaussian, heavy-tailed character. In our experiments we use c = 1, σ = 0.5, jump rate λ = 1, and jump scale σJ = 1 (comparable to the inter-well distance), integrating with the Euler–Maruyama scheme at step dt = 10−3 over t ∈ [0, 5] and recording snapshots every ∆t = 0.05. We simulate 1024 training and 256 test trajectories from Gaussian-distributed initial conditions and benchmark the neural-SDE architectures on their ability to recover both the bistable drift and the jump-augmented noise structure from data alone. This setting is a deliberately stringent test: a method that models the increments as purely Gaussian must either underestimate the barrier-crossing frequency or inflate the diffusion coefficient, whereas the jump events appear as anomalously large displacements localized to a single coordinate. We report drift reconstruction error and (sliced) Wasserstein distance between simulated and learned rollouts, averaged over independent training runs of each model. D.4
Coarse-graining (Molecular dynamics)
Finally, we include a data-driven benchmark with no analytical ground-truth drift, based on [38]. A heavy tagged particle (mass M , radius R, mass ratio µ = M/m) collides elastically with an ideal gas of light point particles; the bath is held in equilibrium by exact Poisson boundary influx and collisions are resolved analytically within each step. It is known that as the particle density and relative mass go to infinity, the dynamics converge to the limiting Langevin process √ dXi = Vi dt, dVi = −γVi dt + γ 2D dWi , as µ → ∞ (64) with friction γ and diffusion D = 1. However, there is analytical model for rarefied gases or for only moderately heavier tagged particles; in particular, dynamics in this rarefied regime are distinctly non-Brownian and involve jump-like phenomena. 23
We run an experiment in this rarefied non-Brownian regime with µ = 10. We run 512 high-fidelity microscopic simulations of a three-dimensional particle bath with 20000 small particles in the bath and examine their effect on the tagged particle. The full 6D phase-space dynamics of the tagged particle is then exported so it can be fed to the models for training. Further details may be found in the repository.
E
Implementation details
The code for simulations and training was written in the Julia programming language. Julia is a highlevel, high-performance language made for scientific computing. In particular, it has best-in-class stochastic simulation capabilities, which include accurate resolution of jumps, which are not available in mainstream ML languages such as Python. This has led to the baseline models (Euler-Maruyama, SDE-GAN and Trajecotry Flow Matching) being carefully ported and adapted to Julia. The ports may be found in the code repo. Any mistakes are our own.
24