Kolmogorov-Chapman Maps: Trajectory-Free Learning for Simulating Molecular Dynamics

Majdi Hassan; Marta Skreta; Dominique Beaini; Kirill Neklyudov

cs.LGNov 9, 2026

Kolmogorov-Chapman Maps: Trajectory-Free Learning for Simulating Molecular Dynamics

Authors: Majdi Hassan, Marta Skreta, Dominique Beaini, Kirill Neklyudov

Organizations: Mila – Quebec AI Institute · Université de Montréal · Valence Labs · Institut Courtois

Abstract

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 $X_{0}$ (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 ( $\Phi_{2}$ , $\Phi_{3}$ ).

Table 1: Ablation study on AL2. NLL on samples from MD simulation and energy

\mathcal{E}

\mathcal{W}_{2}

(

k_{\mathrm{B}}T

) 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

\tau=10,20,\dots,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 .

Source temperature $T_{S}=4$
	LJ-13		LJ-55
Method	$\mathcal{E}$ - $\mathcal{W}_{2}$ $\downarrow$	Geo- $\mathcal{W}_{2}$ $\downarrow$	$\mathcal{E}$ - $\mathcal{W}_{2}$ $\downarrow$	Geo- $\mathcal{W}_{2}$ $\downarrow$
MALA	$0.325$	$\mathbf{1.51}$	$3.40$	$\mathbf{3.80}$
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 ( ${\mathcal{L}}_{\mathrm{KC}}$ )	$\mathbf{0.212}$	$\mathbf{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 (

\mathcal{E}

\mathcal{W}_{2}

) and the geometric 2-Wasserstein distance (Geo-

\mathcal{W}_{2}

). 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 $\times$ 2 $\mu$ 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.

\mathcal{W}_{2}

metrics are computed between

10^{4}

reference frames and

10^{4}

frames drawn from the run after a 10 % burn-in, and report the mean

\pm

std. over 3 draws. Best of HFM and KC-maps in bold .

Panel of Figure 3 — Figure 3: One GPU-hour on single L40S. (a) TICA densities of the reference (16 chains; $dt=1$ fs; 2 $\mu$ s each), MD (16 parallel chains; $dt=1$ fs; one GPU-hour), HFM (16 chains; $\tau=5$ fs; one GPU-hour) and KC-maps (16 chains; $\tau=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 ( $\tau=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 ( $\alpha$ , $\beta$ , ppII, $\alpha_{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 $\mu$ s each; KC-maps: 16 chains at $\tau=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
$X_{\mathrm{std}}$	nm	0.167	0.228	0.300
$T_{\mathrm{std}}$	fs	10	10	10
$P_{\mathrm{std},i}$ (H)	amu nm ps -1	1.585	1.585	1.611
$P_{\mathrm{std},i}$ (C)	amu nm ps -1	5.473	5.473	5.564
$P_{\mathrm{std},i}$ (N)	amu nm ps -1	5.910	5.910	6.008
$P_{\mathrm{std},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 $\times$ attention layers	$9\times 1$	$9\times 1$	$8\times 2$	$8\times 2$	$8\times 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 $(n_{a},n_{b})$	$(1,2)$	$(1,2)$	$x$ : $(3,3)$ , $p$ : $(2,3)$
Fixed $\omega$ (all gates)	$40$		$333\,\mathrm{ps}^{-1}$
Reference cap $\tau_{0}=\delta$	0.01	0.01	0.75 fs

Appendix

Table A3: Training hyperparameters. LJ times are in reduced units. For LJ,

T_{\max}

is given for

T_{S}=1

T_{S}=4

System	Method	Parameters (M)	Batch size	Updates	FLOPs ( $10^{18}$ )
AL4	HFM	9.09	512	$2\,061\,097$	15.513
	KC-maps	46.58	512	$1\,500\,000$	15.513
AL6	HFM	9.09	512	$1\,957\,876$	31.299
	KC-maps	46.58	1024	$1\,000\,000$	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 $400\,\mathrm{ns}$ 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

\tau

for

\texttt{n\_chains}\in{1,16,32}

. All runs on a single L40S GPU.

Figure A5: Equilibrium accuracy on AL2 across step sizes: trajectory MLE vs.

{\mathcal{L}}_{KC}

. 2-Wasserstein distances to reference MD for 16 chains run for one hour on a single L40S GPU at each

\tau

. 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.

P_{ij}(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

\mathtt{ema\_decay}

for the two-step target

\eta=\operatorname{EMA}(\theta)

in Eq. 12 . All models are trained for 500K steps with batch size 512. We report NLL on MD transition samples at

\tau\in\{10,20,50,100\}

fs. A decay of zero corresponds to

\eta=\operatorname{stopgrad}(\theta)

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:

\mathcal{E}

\mathcal{W}_{2}

in units of

k_{\mathrm{B}}T

, averaged over source configurations.

Stochastic Differential Equations

Stochastic Differential Equations describe continuous-time systems whose dynamics combine deterministic drift with random noise.

See all 117 papers

Neural Kolmogorov Equations: Parallelizable Learning of Stochastic Dynamics under General Noise
Jul 21, 2026Arthur Bizzi, Olga FinkDynamical Systems Stochastic Differential Equations
EPFL, Imperial College · EPFL
Strong Stochastic Flow Maps
May 31, 2026Sam McCallum, Zander W. Blasingame, Timothy Herschell +3Stochastic Differential Equations Few-Step Diffusion Sampling
University of Bath · 2AITHYRA
Langevin Flow Maps: Efficient Molecular Dynamics and Transition Path Sampling
Oct 5, 2026Sam McCallum, Niklas Rindtorff, Alexander Tong +1Stochastic Differential Equations Machine Learning Interatomic Potentials
University of Bath · AITHYRA

Jul 21, 2026J/K move · Enter open · S save

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

Arthur Bizzi, Olga Fink

EPFL, Imperial College · EPFL

Dynamical Systems Stochastic Differential Equations Neural PDE Solvers

Open paper AskPDF arXiv

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.

Abstract

Stochastic Differential Equations

Explore similar work

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

Strong Stochastic Flow Maps

Langevin Flow Maps: Efficient Molecular Dynamics and Transition Path Sampling