[论文解读] Timewarp: Transferable Acceleration of Molecular Dynamics by Learning Time-Coarsened Dynamics

Leon Klein, Andrew Y. K. Foong+6 more|arXiv (Cornell University)|Feb 2, 2023

Machine Learning in Materials Science被引用 20

一句话总结

Timewarp 使用条件正则化流作为 MCMC 提案来对 Boltzmann 分布进行采样，通过学习时间粗化的动力学，对未知的小肽实现可转移的、实时钟加速。

ABSTRACT

Molecular dynamics (MD) simulation is a widely used technique to simulate molecular systems, most commonly at the all-atom resolution where equations of motion are integrated with timesteps on the order of femtoseconds ($1 extrm{fs}=10^{-15} extrm{s}$). MD is often used to compute equilibrium properties, which requires sampling from an equilibrium distribution such as the Boltzmann distribution. However, many important processes, such as binding and folding, occur over timescales of milliseconds or beyond, and cannot be efficiently sampled with conventional MD. Furthermore, new MD simulations need to be performed for each molecular system studied. We present Timewarp, an enhanced sampling method which uses a normalising flow as a proposal distribution in a Markov chain Monte Carlo method targeting the Boltzmann distribution. The flow is trained offline on MD trajectories and learns to make large steps in time, simulating the molecular dynamics of $10^{5} - 10^{6}\: extrm{fs}$. Crucially, Timewarp is transferable between molecular systems: once trained, we show that it generalises to unseen small peptides (2-4 amino acids) at all-atom resolution, exploring their metastable states and providing wall-clock acceleration of sampling compared to standard MD. Our method constitutes an important step towards general, transferable algorithms for accelerating MD.

研究动机与目标

• Motivate the need for efficient sampling of Boltzmann distributions in MD, especially for processes spanning long timescales.
• Develop a transferable, machine-learned proposal mechanism that can generalize across unseen small peptides.
• Achieve wall-clock acceleration of asymptotically unbiased Boltzmann sampling via MCMC with learned proposals.

提出的方法

• Train a conditional normalising flow to approximate the time-tau MD transition distribution ϵ{ũ}{x(t+τ)}|x(t) from MD trajectories.
• Use the learned flow as an MCMC proposal with a Metropolis-Hastings correction to target the Boltzmann distribution augmented with auxiliary variables.
• Employ augmented normalising flows where velocities are treated as non-physical auxiliary variables to simplify learning and enforce permutation equivariance.
• Incorporate permutation, translation, and rotation symmetries through a permutation-equivariant transformer-based flow and data augmentation.
• Adopt a RealNVP-style flow with atom-transformer modules and kernel self-attention to capture local interactions.
• Optionally provide a fast exploratory mode without MH correction for rapid metastable-state discovery.

实验结果

研究问题

• RQ1 Can a learned conditional flow generalize to unseen small peptides in Cartesian coordinates while targeting Boltzmann sampling?
• RQ2 Does Timewarp provide transferable acceleration of sampling across different peptide systems compared to standard MD?
• RQ3 What are the trade-offs between asymptotically unbiased MCMC sampling and fast biased exploration for exploring metastable states?

主要发现

• Timewarp achieves wall-clock acceleration over standard MD in sampling metastable states for unseen peptides, with reported speed-ups in ESS per second (e.g., ~7x for alanine dipeptide, ~33x for 2AA dipeptides).
• Timewarp demonstrates transferability to unseen small peptides (2-4 amino acids) at all-atom resolution using training data from multiple peptides.
• The method can operate with an MH correction to ensure unbiased Boltzmann sampling, albeit with low acceptance rates (~1%), yet still enables faster exploration when able to make larger proposed moves.
• A biased exploration mode without MH correction yields substantial qualitative speedups in metastable-state discovery, with median speedups around ~600x when sampling parallel chains.
• Timewarp is designed as a general, transferable approach applicable to unseen systems without retraining, and is compatible with integration with other enhanced-sampling methods.