[論文レビュー] 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 は Boltzmann 分布をサンプリングするための MCMC 提案として条件付き正規化フローを用い、時間の粗化したダイナミクスを学ぶことで、未知の小ペプチドに対して転送性を持つ壁時計時間での加速を達成します。

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.

研究の動機と目的

• 長時間スケールに跨るプロセスを特に対象と、分子動力学における Boltzmann 分布の効率的なサンプリングの必要性を動機づける。
• 未知の小さなペプチド全体に一般化できる転送可能な機械学習による提案機構を開発する。
• 学習提案を用いたMCMCによる漸近的に unbiased な Boltzmann サンプリングの壁時計時間での加速を達成する。

提案手法

• MD 軌跡から、条件付き正規化フローを訓練して time-tau MD 遷移分布 ルークの式を近似する。
• 学習済みフローを Metropolis-Hastings 補正を伴う MCMC 提案として用い、補助変数を付加した Boltzmann 分布を標的とする。
• 速度を非物理的な補助変数として扱う拡張正規化フローを用いて、学習を簡素化し置換対称性を保証する。
• 置換、平行移動、回転対称性を、置換同値性を持つトランスフォーマーベースのフローとデータ拡張を通じて組み込む。
• RealNVP 風のフローを採用し、原子トランスフォーマーモジュールとカーネル自己注意を用いて局所相互作用を捉える。
• 迅速な準安定状態の発見のために MH 補正なしの高速探索モードを任意で提供する。

実験結果

リサーチクエスチョン

• RQ1 Cartesian 座標系で未知の小さなペプチドに一般化できる条件付きフローを学習して Boltzmann サンプリングを標的にすることは可能か？
• RQ2 Timewarp は標準的な MD と比較して、異なるペプチド系に対するサンプリングの転送可能な加速を提供するか？
• RQ3 漸近的に unbiased な MCMC サンプリングと、準高速で偏りのある探索を用いた準安定状態の探索とのトレードオフは何か？

主な発見

• Timewarp は未知のペプチドの準安定状態のサンプリングにおいて標準MDより壁時計時間での加速を達成し、ESS/秒での速度向上を報告（例：アラニンジペプチドで約7倍、2AAジペプチドで約33倍）。
• Timewarp は複数のペプチドからの訓練データを用いて、全原子分解能で未知の小さなペプチド（2-4 アミノ酸）への転送性を示す。
• 本手法は unbiased Boltzmann サンプリングを保証するための MH 補正を用いて動作できるが、受け入れ率は低く（約1%程度）それでもより大きな提案移動が可能な場合は探索を高速化する。
• MH 補正なしの偏った探索モードは準安定状態の発見で顕著な定性的な速度向上をもたらし、並列連鎖をサンプリングする場合の中央値の速度向上は約600x程度。
• Timewarp は再訓練なしで未知のシステムに適用可能な一般的で転送可能なアプローチとして設計されており、他の強化サンプリング法との統合にも適している。