One-sentence Summary

Researchers from Freie Universität Berlin, University of Cambridge, Microsoft Research AI4Science, and Rice University introduce Timewarp, a normalising-flow-based Markov chain Monte Carlo proposal distribution that learns time-coarsened dynamics with steps of $10^{5}-10^{6}$105−106 fs, enabling transferable acceleration of all-atom molecular dynamics and providing wall-clock speedups over standard MD when sampling unseen small peptides.

Key Contributions

• Timewarp is an enhanced sampling method that uses a normalising flow as a proposal distribution inside a Markov chain Monte Carlo framework targeting the Boltzmann distribution, with the flow trained offline on MD trajectories to propose large time-coarsened steps of 10^5 to 10^6 femtoseconds.
• Operating directly in all-atom Cartesian coordinates, Timewarp achieves transferability across molecular systems, allowing a single trained model to generalise to unseen small peptides (2 to 4 amino acids) without retraining.
• Results show that Timewarp with a Metropolis-Hastings correction accelerates asymptotically unbiased sampling on unseen dipeptides, while the version without the correction explores metastable states of both dipeptides and tetrapeptides faster than standard molecular dynamics, with the discovered metastable states verified as physically meaningful.

Introduction

Molecular dynamics (MD) is a widely used technique for simulating atomic-level physical systems, offering detailed insight into molecular motion without wet-lab experiments. A core goal is efficient sampling from the Boltzmann distribution, which enables computation of key biophysical quantities like free energies of protein folding or ligand binding. Standard MD implementations rely on Langevin dynamics discretized with a small timestep (around 1 fs) for stability, but many biomolecules feature metastable states separated by energy barriers that require millisecond-scale simulations, corresponding to roughly 10^12 integration steps, making direct simulation infeasible. Prior enhanced sampling methods, such as coarse graining and metadynamics, help but demand system-specific domain knowledge and do not transfer learned information between molecular systems, even when those systems exhibit closely related dynamics, as with proteins built from sequences of 20 amino acids.

The authors introduce Timewarp, a general and transferable enhanced sampling approach that uses a normalizing flow as a proposal distribution within a Markov chain Monte Carlo (MCMC) framework targeting the Boltzmann distribution. Their main contributions include the first machine learning algorithm operating in general Cartesian coordinates that demonstrates transferability to unseen small peptide systems, as well as the first wall-clock acceleration of asymptotically unbiased MCMC sampling of the Boltzmann distribution for such unseen systems. They also define an MCMC algorithm with a conditional normalizing flow proposal and a Metropolis-Hastings (MH) correction step to ensure detailed balance, design a permutation equivariant transformer-based flow, create a novel training dataset of MD trajectories for thousands of small peptides, and show that Timewarp without the MH correction can explore metastable states of new peptides much faster than standard MD.

Dataset

Dataset composition and sources

• The authors generate molecular dynamics (MD) trajectories by integrating Equation (2) using the OpenMM library.
• They simulate small proteins (peptides) in implicit water, meaning water molecule degrees of freedom are not explicitly modeled.

Key details and scale

• The dataset is denoted as D = {T_i} for i = 1 to P, where P is the total number of peptides.
• Each MD trajectory is temporally subsampled with a spacing of tau, so a single trajectory is represented as T_i = (x(0), x(tau), x(2tau), ...).

Usage in training

• During training, the authors randomly sample pairs x(t) and x(t + tau) from the dataset.
• Each pair serves as a sample from the conditional distribution mu(x(t + tau) | x(t)).
• Because the flow is trained on trajectory data from multiple peptides, it can generalize at test time to new peptides not present in the training data.
• Additional dataset details are provided in Appendix E.

Method

The authors introduce Timewarp, a deep probabilistic model designed to approximate the conditional distribution $\mu(x(t+\tau)|x(t))$μ(x(t+τ)∣x(t)) induced by molecular dynamics simulations over a time interval $\tau$τ. This approximation enables the generation of long-time-step proposals for sampling molecular configurations, effectively capturing the transition dynamics between metastable states. As shown in the figure below:

The core of the method is a conditional normalizing flow $p_\theta(x(t+\tau)|x(t))$pθ​(x(t+τ)∣x(t)). To enhance expressivity and simplify the construction of permutation-equivariant coupling layers, the authors employ an augmented normalizing flow framework. This involves introducing auxiliary latent variables, specifically treating latent velocities $z^v$zv as non-physical auxiliary variables drawn from a standard Gaussian, while the actual positions $z^p$zp are transformed. The architecture is built upon RealNVP coupling layers, as depicted in the schematic illustration of the flow architecture.

The flow consists of stacked coupling layers. In each layer, the position latent variables are affinely transformed based on the auxiliary latent variables, and vice versa. The scale and translation functions within these layers are implemented using a novel "atom transformer" module. This module processes atom-wise features by concatenating the atom's position, a learned type embedding, and the relevant latent variable. These features are passed through a multi-layer perceptron (MLP) and then through several transformer layers.

A key innovation within the atom transformer is the use of kernel self-attention instead of standard dot-product attention. This design choice is motivated by the physical observation that forces in molecular systems are local. The attention weight $w_{ij}$wij​ for atom $i$i attending to atom $j$j is defined using a Radial Basis Function (RBF) kernel that decays with the squared Euclidean distance between atoms:

where $\ell$ℓ is a learnable lengthscale parameter. A multihead version is employed where each attention head utilizes a different lengthscale, allowing the model to capture interactions at varying spatial scales.

The architecture is designed to respect physical symmetries. Permutation equivariance is enforced by processing atoms individually without positional encoding. Translation equivariance is achieved by centering the input positions, and rotation equivariance is handled via data augmentation during training, where random rotation matrices are applied to the training pairs.

The model undergoes a two-stage training procedure. In the first stage, likelihood training, the flow is optimized via maximum likelihood estimation on pairs of states sampled from molecular dynamics trajectories. This aligns the flow's density with the true conditional distribution. In the second stage, acceptance training, the model is fine-tuned to maximize the Metropolis-Hastings (MH) acceptance probability. The objective function is a weighted combination of the acceptance loss, the likelihood loss, and an entropy regularization term. The entropy term prevents the model from converging to trivial proposals that result in negligible state changes, ensuring sufficient exploration of the state space.

During inference, the trained flow serves as a proposal distribution in an MCMC algorithm. The process iteratively samples a proposal $\tilde{X}$X~ from the flow, accepts or rejects it based on the MH acceptance ratio targeting the Boltzmann distribution, and resamples the auxiliary variables via Gibbs sampling. Alternatively, for faster exploration, the authors propose a biased sampling mode that accepts proposals with energy changes below a certain cutoff, neglecting the MH correction to achieve significantly faster state space traversal.

Experiment

Timewarp is evaluated on small peptide systems, including alanine dipeptide, dipeptides, and tetrapeptides, using both MCMC and an exploration mode to accelerate sampling of metastable states. On alanine dipeptide and dipeptides, it matches MD distributions and achieves considerable speed-ups, with exploration mode showing much higher gains, while on tetrapeptides MCMC often misses states but exploration successfully uncovers and validates missing metastable states with short MD runs. Limitations include low acceptance rates, frequent state misses for larger peptides, and difficulty scaling to bigger systems or capturing transition states, though future integration with other enhanced sampling methods is suggested.