[論文レビュー] Implicit Transfer Operator Learning: Multiple Time-Resolution Surrogates for Molecular Dynamics

Mathias Schreiner, Ole Winther+1 more|arXiv (Cornell University)|May 29, 2023

Machine Learning in Materials Science被引用数 8

ひとこと要約

ITO は SE(3)-等変_DIFFUSION モデルを用いて分子動力学の複数の時間解像度の代理モデルを学習し、自己一致的な長時間ダイナミクスとスケール間の効率的サンプリングを実現します。

ABSTRACT

Computing properties of molecular systems rely on estimating expectations of the (unnormalized) Boltzmann distribution. Molecular dynamics (MD) is a broadly adopted technique to approximate such quantities. However, stable simulations rely on very small integration time-steps ($10^{-15}\,\mathrm{s}$), whereas convergence of some moments, e.g. binding free energy or rates, might rely on sampling processes on time-scales as long as $10^{-1}\, \mathrm{s}$, and these simulations must be repeated for every molecular system independently. Here, we present Implict Transfer Operator (ITO) Learning, a framework to learn surrogates of the simulation process with multiple time-resolutions. We implement ITO with denoising diffusion probabilistic models with a new SE(3) equivariant architecture and show the resulting models can generate self-consistent stochastic dynamics across multiple time-scales, even when the system is only partially observed. Finally, we present a coarse-grained CG-SE3-ITO model which can quantitatively model all-atom molecular dynamics using only coarse molecular representations. As such, ITO provides an important step towards multiple time- and space-resolution acceleration of MD. Code is available at \href{https://github.com/olsson-group/ito}{https://github.com/olsson-group/ito}.

研究の動機と目的

• 分子系のボルツマン分布の期待値を推定する必要性を動機づけ、極めて小さな MD タイムステップを超える長時間スケールのダイナミクスの課題に対処する。
• MD 遷移確率のマルチタイム解像度代理モデルを学習するための Implicit Transfer Operator (ITO) フレームワークを導入する。
• 空間対称性を物理的に一貫性を保つために SE(3)-等変生成モデルを開発する。
• ITO が複数の時間スケールに渡る自己一致的な確率過程を生成できること、粗視化表現を用いた場合には全原子 MD を近似することを示す。

提案手法

• ITO を条件拡散確率モデル (cDDPM) によって条件付き遷移確率 p_{Nτ}(x_{Nτ}|x_0) を学習する形式として定式化する。
• 3D 回転と平行移動に対する不変性/等変性を保証する SE(3)-等変アーキテクチャ（ChiroPaiNN）を用いる。
• 転送演算子の固有関数射影によって遷移確率を時間変動成分と時間不変成分に分解し、訓練を複数のラグタイム (Nτ) に整合させる。
• モデルにラグ時間の分布を曝露するデータ拡張戦略を用いて、固有関数表現の学習をよりよく行う。
• 二つのアーキテクチャを提供する：分子系向けの SE3-ITO と粗視化タンパク質折りたたみデータ向けの CG-SE3-ITO。ともに拡散ベースのスコアモデリングと SE(3) 不変性を保証する不変事前分布を使用する。

実験結果

リサーチクエスチョン

• RQ1 MD データから学習された暗黙の転送演算子代理モデルは、複数の時間スケールにわたって自己一貫性のある確率過程を生成できるか。
• RQ2 確率的ラグ训练はメタ安定性の捉え方と長時間ダイナミクスを、固定ラグ訓練と比べて改善するか。
• RQ3 SE(3)-等変ITOモデルは、主要な動的・定常観測量を保ちつつ粗視化表現へスケールできるか。
• RQ4 ITO代理モデルは、リウェイティングを明示的に行わず、折りたたみ/展開時間などの動的観測量と自由エネルギーなどの定常観測量をどれだけ再現できるか。

主な発見

• ITO モデルは確率的ラグ訓練を用いると、 Müller–Brown ベンチマークにおけるメタ安定性予測で固定ラグモデルを上回る。
• SE3-ITO は長時間ダイナミクスを自己一貫的に生成し、アラニン二肽ペプチドの複数の時間スケールにわたり MD データと一致する。
• CG-SE3-ITO は高速折りたたみタンパク質（Chignolin, Trp-Cage, BBA, Villin）において、長時間原子 MD に比べ主要な動的・定常観測量を再現する。
• ITO を用いた祖先サンプリングは、長い軌跡（マイクロ秒まで）を生成でき、緩やかに緩和する座標での MD と一致する。
• CG-SE3-ITO から算出した折りたたみ自由エネルギーや平均初到達時間などの観測量は、参照 MD の統計（報告されたばらつき内）と一致する。
• この枠組みは GPU 上で高いサンプリング速度を維持しつつ、MD 軌道と定性的な一致を保つことで、実質的な効率向上を示す。