Stochastic differential equations (SDEs) describe stochastic processes across a wide range of scientific domains, including physics, biology, and finance. Although individual numerical integration steps are inexpensive, observing meaningful events or transitions can require millions of sequential steps, resulting in substantial wall-clock simulation time and poor hardware utilization. Deep learning offers an alternative by amortizing many integration steps into a single neural network forward pass, potentially enabling much faster simulation. However, existing supervised approaches typically require long simulated trajectory data, which restricts their applicability. In this work, we introduce Kolmogorov-Chapman Maps (KC-Maps), a method for learning the transition kernel of an SDE that relies only on the trajectory-free marginal samples together with the known drift and diffusion coefficients. Given an initial time and state of a process, KC-Maps models its distribution integrated to any given time, corresponding to the weak solution of the SDE. To recover the correct dynamics, KC-Maps satisfies two key properties: (i) it satisfies the boundary conditions, which we incorporate directly into the neural network architecture, and (ii) it respects the Kolmogorov--Chapman equation, which provides an efficient training objective. We empirically demonstrate that KC-Maps learns to integrate the dynamics of peptide molecular systems and outperforms classical molecular dynamics in terms of GPU-hours.
Figures & tables
Figure 1: A single KC-maps rollout explores the landscape of metastable alanine tetrapeptide conformers. Starting from a single conformer X0 (left), KC-maps is rolled out for 200 ns. Orange arrows trace the first visit to each newly discovered basin. Representative conformers from each basin are shown on the right, colored by basin identity, along with their backbone torsion angles ( Φ2 , Φ3 ).
Table 1: Ablation study on AL2. NLL on samples from MD simulation and energy E - W2 ( kBT ) over 64 source states with 256 samples per state compared to 256 independent MD trajectories with a 0.1 fs timestep. Results are averaged over τ=10,20,…,100 fs, with lower values being better. Each subtable compares the full model with variants that change only the indicated component, keeping all other settings fixed. All models were trained with a reduced compute budget of 250K steps. Best value in each column within each subtable is bold .
LJ-13
LJ-55
Method
E - W2↓
Geo- W2↓
E - W2↓
Geo- W2↓
MALA
0.325
1.51
3.40
3.80
Source temperature TS=4
PITA ( Akhound-Sadegh et al., 2025 )
2.260
1.65
−
−
ITM ( Potaptchik et al., 2025 )
0.879
1.54
26.37
3.98
KC-maps ( LKC )
0.212
1.51
1.51
3.82
Table 2: Performance of KC-maps on LJ-13 and LJ-55. Published baseline results are taken from Potaptchik et al. (2025) . Metrics are the 2-Wasserstein distance between energy distributions ( E - W2 ) and the geometric 2-Wasserstein distance (Geo- W2 ). In all settings, the target temperature is T=1 ; best values are bolded and the second best values are underlined . Dashes indicate unavailable results.
Figure 2: Equilibrium sampling of alanine peptides. Density panels show the reference MD (16 parallel chains × 2 μ s per chain) beside KC-maps (Ours). Potential energy histograms of KC-maps against each system’s MD reference. (a) AL2: Ramachandran and TICA densities, energy histogram, and TIC0 against simulated MD trajectory (top) and ours (bottom). (b) AL4: TICA densities and energy histogram. (c) AL6: TICA densities and energy histogram. For each molecule, samples generated from KC-maps were run with varying compute budgets (1 GPU-hour for AL2; 22 GPU-hour for AL4; 40 GPU-hours for AL6).
Table 3: Equilibrium sampling of alanine peptides per method with a fixed budget of one GPU-hour on single L40S, scored against each molecule’s long MD. W2 metrics are computed between 104 reference frames and 104 frames drawn from the run after a 10 % burn-in, and report the mean ± std. over 3 draws. Best of HFM and KC-maps in bold .
Figure 3: One GPU-hour on single L40S. (a) TICA densities of the reference (16 chains; dt=1 fs; 2 μ s each), MD (16 parallel chains; dt=1 fs; one GPU-hour), HFM (16 chains; τ=5 fs; one GPU-hour) and KC-maps (16 chains; τ=30 fs; one GPU-hour). KC-maps consistently visits most if not all of the modes within the same GPU-hour. (b) Left: simulated aggregate time per GPU-hour against the number of parallel chains for KC-maps ( τ=30 fs for AL4 and AL6), MD and HFM (7 fs for AL4, 5 fs for AL6). Right: metastable TICA modes discovered against GPU time. Dotted: modes of the reference MD.
Figure 4: AL2 kinetics. (a) The four metastable modes ( α , β , ppII, αL ) on the reference MD Ramachandran plot, with their populations. (b) Mean first-passage time between modes. (c) Autocorrelation C(t) of mode occupancy. MD: 16 chains of 2 μ s each; KC-maps: 16 chains at τ=30 fs run for one GPU-hour on L40S. Both metrics are measured on trajectories with frames saved every 1 ps.
Appendix figures & tables11 assets
Supplementary material from the paper’s appendix.
Appendix
Scale
Unit
AL2
AL4
AL6
Xstd
nm
0.167
0.228
0.300
Tstd
fs
10
10
10
Pstd,i (H)
amu nm ps -1
1.585
1.585
1.611
Pstd,i (C)
amu nm ps -1
5.473
5.473
5.564
Pstd,i (N)
amu nm ps -1
5.910
5.910
6.008
Pstd,i (O)
amu nm ps -1
6.317
6.317
6.421
Appendix
Table A1: Standardization scales. Position scales are fitted to the training split. Momentum scales are listed by element at 300 K for AL2/AL4 and 310 K for AL6.
AL2
AL4
AL6
Sequence
Ace-A-Nme
Ace-AAA-Nme
AAAAAA
Atoms (dimension d )
22 (66)
42 (126)
63 (189)
Force field
Amber ff99SB-ILDN
Amber ff99SB-ILDN
Amber ff14SB
Implicit solvent
OBC2
OBC2
OBC1
Nonbonded cutoff
none
none
2 nm
Temperature
300 K
300 K
310 K
Appendix
Table A2: Peptide data and physical settings. Training uses random configurations from these datasets. The learned dynamics use the listed force fields, temperatures and friction coefficients. The source trajectories use a 1 fs integration step with flexible bonds.
LJ-13
LJ-55
AL2
AL4
AL6
Blocks × attention layers
9×1
9×1
8×2
8×2
8×2
Channels / heads
256 / 4
256 / 4
384 / 6
384 / 6
384 / 6
Parameters
9.4M
9.4M
46.6M
46.6M
46.6M
Gate orders (na,nb)
(1,2)
(1,2)
x : (3,3) , p : (2,3)
Fixed ω (all gates)
40
333ps−1
Reference cap τ0=δ
0.01
0.01
0.75 fs
Appendix
Table A3: Training hyperparameters. LJ times are in reduced units. For LJ, Tmax is given for TS=1 / TS=4 .
System
Method
Parameters (M)
Batch size
Updates
FLOPs ( 1018 )
AL4
HFM
9.09
512
2061097
15.513
KC-maps
46.58
512
1500000
15.513
AL6
HFM
9.09
512
1957876
31.299
KC-maps
46.58
1024
1000000
31.299
Appendix
Table A4: Training compute on AL4 and AL6. Parameter counts, completed training updates and estimated FLOPs for the loss and parameter gradients. FLOP totals follow the accounting described above.
Figure A1: Equilibrium results across step size. Equilibrium metrics for KC-maps with and without detailed-balance loss against reference long MD.
Figure A2: AL6 sampling beyond the training distribution. TICA projections of reference long MD, training data, and KC-maps samples. The KC-maps samples are pooled from 16 chains of approximately 400ns each. The lower-left mode is absent from the training configurations but visited by KC-maps.
Figure A3: Aggregate physical time simulated across chains across τ for n_chains∈1,16,32 . All runs on a single L40S GPU.
Figure A5: Equilibrium accuracy on AL2 across step sizes: trajectory MLE vs. LKC . 2-Wasserstein distances to reference MD for 16 chains run for one hour on a single L40S GPU at each τ . Error bars are the standard deviation over three subsamples of the trajectory. The dashed line is MD with 16 chains for one hour.
Figure A6: AL2 transition probabilities. Pij(t) from mode i (rows) to mode j (columns) against the lag t . Shaded bands are 95% bootstrap intervals over chains.
Figure A7: Target EMA decay on AL2. Models differ only in ema_decay for the two-step target η=EMA(θ) in Eq. 12 . All models are trained for 500K steps with batch size 512. We report NLL on MD transition samples at τ∈{10,20,50,100} fs. A decay of zero corresponds to η=stopgrad(θ) .
Figure A8: Flow parameterization on AL2. Comparison of standard, displacement, and whitened flows, each trained for 500K steps with batch size 512 and evaluated using EMA parameters with decay 0.999 . Lines and shaded bands show the mean and standard deviation across training seeds: 42 , 43 , and 44 for all three parameterizations. Left: NLL per degree of freedom (DoF) on molecular dynamics samples. Right: E - W2 in units of kBT , averaged over source configurations.
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.
Flow and diffusion models generate high-quality samples in many modalities; however, many network evaluations are required during inference due to numerical integration of an underlying differential equation. Flow maps alleviate this problem by learning the solution map of the differential equation directly, enabling few-step sampling. Yet, current methods are restricted to approximating the solution map of ODEs. These methods can be used to learn the transition kernel of an SDE, thereby obtaining a solution map that recovers the marginal distributions of the process (weak convergence) rather than the solution path (strong convergence). We propose Strong Stochastic Flow Maps (SSFMs) as a novel framework for learning the strong solution map of additive-noise SDEs, directly generalizing deterministic flow maps to the stochastic setting. Further, a polynomial approximation to Brownian motion is introduced and shown to converge pathwise. These results enable a simulation-free training objective for the solution map of diffusion models. We demonstrate that SSFMs outperform previous stochastic flow map methods on image generation and enable few-step sampling of molecular systems.
Sam McCallum, Zander W. Blasingame, Timothy Herschell +3
Molecular dynamics simulations proceed by integrating the Langevin equations over many small femtosecond timesteps. This poses a challenge for estimating ensemble properties and transition dynamics that occur on much longer timescales. We introduce Langevin Flow Maps, which extend machine-learned force-fields to additionally learn the stochastic Langevin integrator. We show that Langevin Flow Maps enable large-timestep molecular dynamics and recover accurate dynamical properties of the system, while running an order of magnitude faster than current machine-learned force fields. Further, by training on a diverse molecular dataset, we demonstrate a path towards transferable Langevin Flow Maps.