KC-maps drops trajectory data: learning a transition kernel from drift and diffusion lets peptide sampling outrun classical MD
Synopsis
KC-maps learns the finite-time transition kernel of a stochastic differential equation directly as a neural network map, trained only on uncorrelated marginal samples plus the known drift and diffusion coefficients rather than on paired simulated trajectories. Boundary conditions are built into the network architecture, and a training objective derived from the Kolmogorov–Chapman equation keeps transitions composed over different time scales consistent in distribution. On alanine dipeptide, tetrapeptide and hexapeptide it replaces up to 70 fs of sequential integration with a single learned transition, reaches lower TICA-2 distances at matched GPU-hour budgets, and on AL6 covers conformational modes absent from the training data.
Interpretation
Long-time evolution of a molecular simulation can now be crossed in one learned transition instead of step-by-step integration. Earlier learned accelerators for molecular simulation required pairs of states from reference trajectories as supervision, and producing those pairs still needed numerical integration at femtosecond steps, making training data expensive. On alanine dipeptide, tetrapeptide and hexapeptide, a transition kernel trained only on equilibrium configurations and force-field evaluations covers intervals up to 70 fs, replacing up to 70 sequential integration steps.
The transition kernel is represented by a conditional normalizing flow whose short-time behavior is structurally constrained to match the given stochastic differential equation. A normalizing flow is an invertible density model with a tractable Jacobian; here it starts from a Gaussian reference kernel and applies corrections through time-gated affine coupling layers. The time gates drive corrections to zero as elapsed time goes to zero, satisfying the no-jump condition and the first- and second-moment boundary conditions set by drift and diffusion, while the training objective requires a two-step composition to agree in distribution with a single direct transition.
At an equal one-GPU-hour budget, the learned transition kernel covers conformational space better than classical molecular dynamics and the HFM baseline. Classical molecular dynamics is limited by sequential integration at femtosecond steps, so slow conformational transitions demand very long wall-clock time. On AL2, AL4 and AL6, KC-maps has the lowest TICA-2 distance while remaining competitive on energy and torsion metrics; on AL6 it reaches three of four TICA modes within an hour, while the MD baseline still misses two modes.
Perspective
The demonstrated scope is the LJ-13 and LJ-55 Lennard-Jones systems and the AL2, AL4 and AL6 peptide systems, with transition intervals up to 70 fs and comparisons run under a one-GPU-hour budget on a single L40S; the method addresses readers working on long-time conformational sampling in molecular simulation.
The authors note that the method depends on training-data coverage: kinetic metrics deviate most for conformational modes sparsely represented in training, and in AL2 that mode accounts for only a small fraction of training configurations while its share in the reference molecular dynamics is clearly higher. The method is currently trained separately for each molecular system and does not support transfer to unseen systems.