[논문 리뷰] Timewarp: Transferable Acceleration of Molecular Dynamics by Learning Time-Coarsened Dynamics

Leon Klein, Andrew Y. K. Foong외 6명|arXiv (Cornell University)|2023. 02. 02.

Machine Learning in Materials Science인용 수 20

한 줄 요약

Timewarp는 Boltzmann distribution를 샘플링하기 위한 MCMC 제안으로 조건부 normalising flow를 사용하여, 시간 축소된 동역학을 학습함으로써 미지의 소형 펩타이드에 대해 전달 가능한 벽시계 가속을 달성한다.

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.

연구 동기 및 목표

• MD에서 Boltzmann 분포의 샘플링이 필요한 이유를 동작 시간스케일이 긴 프로세스에서 특히 강조한다.
• 보지 못한 소형 펩타이드에 대해 일반화될 수 있는 transferable, 머신러닝 기반 제안 메커니즘을 개발한다.
• 학습된 제안을 사용한 MCMC를 통해 asymptotically unbiased Boltzmann 샘플링의 벽시계 가속을 달성한다.

제안 방법

• Training a conditional normalising flow to approximate the time-tau MD transition distribution ði.e., 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.