Deep learning speeds molecular dynamics without full trajectories

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

Machine Learning

Summary

Simulating molecules moving over time usually needs millions of tiny steps, which takes a long time on computers. The authors propose a new deep learning method called Kolmogorov-Chapman Maps that can predict these changes directly without taking all the tiny steps. Their method uses mathematical properties and known rules about the system to learn how molecules evolve. This approach runs faster on GPUs and works well on peptide molecules.

What this means in practice

Authors

Majdi Hassan, Marta Skreta, Dominique Beaini, Kirill Neklyudov

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.