Zenn
🕌

論文解説：粗視化MDのための拡散モデルと力場 (3/3)

2023/08/21に公開
機械学習
diffusion
chemistry
Molecular dynamics
tech

注意

半分メモです。勉強のために読みながら書いています。和訳を中心に適当に補足しています

この記事ではこの論文の本質である粗視化MDのための拡散モデルと力場の抽出方法、実際のネットワーク構造に関して解説していきます。イントロに関するメモはこちら、手法に関するメモはこちらから。

なお、図はすべてarXiv:2302.00600 [cs.LG]より引用しています。

解説する論文

Two for One: Diffusion Models and Force Fields for Coarse-Grained
Molecular Dynamics

Microsoft Researchらのチームによって提案された手法。2023年2月にarxivに公開された。スコアベースの生成モデルにより訓練データ中に力の情報なしで粗視化力場を学習している。また、拡散生成モデルを用いることで、平衡分布からの粗視化構造のサンプリングができるだけでなく、スコア関数自体が粗視化力場として直接利用できることを示している。（おそらく、1つのモデルで構造サンプリングとダイナミクスの計算両方ができるというのがタイトルのtwo for oneの理由と思われる。）

拡散モデルを用いることで、訓練自体がシンプルになり、また、中規模程度までのタンパク質について平衡分布やダイナミクスの再現が先行研究であるCG-NetやFlowよりも改善されたとのこと。

計算対象

  1. アラニンジペプチド
  2. 5種類のfast-floding protein（Chignolin, Trp-cage, Bba, Villin, Protein G）

比較対象の機械学習モデル

  1. DFF
    • この論文のモデル
    • i.i.d (独立同分布)とMDシミュレーション両方を使用
  2. Flow i.i.d, Flow-CGNet sim
    • Kohler (2023)らのモデル
    • Flow i.i.dはflow-matchingの最初の段階で粗視化密度推定として訓練されたモデル
    • Flow-CGNetはflow-matchingの第二段階でaugmented normalizing flowの購買から抽出された力場を用いてMDシミュレーションをしたもの
  3. CGNet sim
    • Wang (2019)らのモデル
    • 粗視化ビーズの力を純粋なforce-matchingによって訓練したモデル
    • MDのみに利用

結果1. アラニンジペプチドの粗視化シミュレーション

訓練条件

  • 500 nsのMDシミュレーションのトラジェクトリから2 psごとに250万サンプルを取得。このシミュレーションをn=4で実施し、3つを訓練、1つをテストデータにした4-fold クロスバリデーションを実施。Parallel temperingなどは行わない。
  • トラジェクトリは温度10 K ～ 500 Kで取得。
  • モデルは96個の隠れ層を持つ2つのGraph transformer layerで構成

評価指標

  • ラマンチャンドランプロット（アラニンジペプチドの2つの二面角の2次元ヒストグラム）と二面角の分布に対するJensen-Shannon (JS)ダイバージェンスの2つ。

結果

  • DFFはFlowやCGNetに比べて少ないデータ数でJSダイバージェンスが小さくなり、ラマチャンドランプロットの再現もより良い。

結果2. Fast-floding proteinの粗視化シミュレーション

訓練条件

  • タンパク質モデルによって異なるが、数百μsのトラジェクトリから200psごとに構造をサンプルし、50万～200万弱のスナップショットを取得した。
  • タンパク質ごとに異なる機械学習モデルを作成した。層数や特徴量の数、学習率などはモデルによって異なる。

評価指標（平衡構造）

  1. TIC JS。タンパク質の配座に対してTime-lagges independent component analysis (TICA)を行った。TICAの第1, 2座標の2次元分布を機械学習モデルによって生成した分布と元のMDから得られた分布のJSダイバージェンスを評価した。
  2. PWD JS。タンパク質のグローバル構造を評価するために、対角オフセットが3より大きい対角距離行列の上三角形内のすべての距離について、テストMD分布に対するJSダイバージェンスを計算した
  3. Cα炭素のRMSDに対する自由エネルギープロファイルとTICの第1, 2座標に対する自由エネルギープロファイル
  4. Contact mapの正規化カウント

結果（平衡構造）

  • いずれのタンパク質においてもFlow, Flow-CGNetよりもDFFの方がTIC JS, PWD JSのどちらも良好な結果となった。
  • DFFはFlow, Flow-CGNetよりも自由エネルギーの再現性もよい
    • （個人的メモ）といいつつ、DFF sim.はそれなりにreferenceから外れているように見えるが・・・？

評価指標（ダイナミクス）

  • 2つのTIC座標の時系列データと2次元TIC空間 (2D-TIC)上での対応関係を解析
  • 遷移確率を計算するために2D-TIC空間をエルボー法でK個のクラスターに分割し、すべての遷移をカウント・正規化し、推定されたマルコフモデルに対応する遷移確率行列を得、ここからJSダイバージェンスを評価した。

結果（ダイナミクス）

  • 評価したすべてのタンパク質において、DFFの方がFlow-CGNetよりも良い結果を示した。
  • また、DFFではFlow-CGNetで評価できなかったProtein Gを評価することができた。

Discussion