😽
論文解説：粗視化MDのための拡散モデルと力場 (1/3)
注意
半分メモです。勉強のために読みながら書いています。和訳を中心に適当に補足しています。
この記事では本論に入るまでのイントロとして、粗視化力場とは何か？粗視化力場の構築にはどんな手法が開発されてきたかということを解説していきます。
解説する論文
Two for One: Diffusion Models and Force Fields for Coarse-Grained
Molecular Dynamics
Microsoft Researchらのチームによって提案された手法。2023年2月にarxivに公開された。スコアベースの生成モデルにより訓練データ中に力の情報なしで粗視化力場を学習している。また、拡散生成モデルを用いることで、平衡分布からの粗視化構造のサンプリングができるだけでなく、スコア関数自体が粗視化力場として直接利用できることを示している。（おそらく、一つのモデルで構造サンプリングとダイナミクスの計算両方ができるというのがタイトルのtwo for oneの理由と思われる。）
拡散モデルを用いることで、訓練自体がシンプルになり、また、中規模程度までのタンパク質について平衡分布やダイナミクスの再現がよくなったとのこと。
Introduction
粗視化シミュレーションと粗視化力場
- 粗視化MDシミュレーションは全原子MDシミュレーションで計算可能な時空間スケールを超える方法である。MDの時空間スケールを数桁アップすることで、大規模なタンパク質の折りたたみイベントや大きな分子間の遅い相互作用など、より長い時間範囲での巨大分子動態に関する新しい研究が可能になる。
- 粗視化シミュレーションを行うためにはまずは全原子(fine-grained)のモデルを粗視化したモデルに射影する必要がある。その後、動力学シミュレーションを行うためには粗視化力場を設計する必要がある。
粗視化力場構築のアプローチ
- 粗視化力場の構築方法にはトップダウンのものとボトムアップのものがある。
- トップダウンの粗視化力場はしばしば実験的に観測可能なマクロな物理量や全原子モデルで得られる物理量を再現するように設計される。
- ボトムアップのアプローチは全原子モデルにおける熱力学量などのミクロなふるまいを再現するような粗視化モデルを構築することになる。この手法では全原子モデルを粗視化モデルに射影して得られた平衡分布と粗視化モデルの平衡分布が一致することが求められる。これまでの熱力学的な一貫性に立脚するボトムアップな粗視化はここ10年で特に機械学習の手法を用いたく結果が多数生まれてきている。
Variational force matching
- 粗視化力場を構築するための主要なアプローチにはvariational force matchingとrelative entropy minimizationの2つがある。
- Variational force matchingでは粗視化モデルにおける力と、全原子モデルを粗視化空間に射影した力の平均二乗誤差を最小化する。この方法では、全原子モデルの中に力を必ず含めておく必要がある。しかし、射影した力は確率的な性質があるため、このノイジーなforce-matchingの推定器は大きな分散を持ち、これは学習効率が良くない。
- 一方で、relative entropy minimizationは原子の力を必要とせず、粗視化空間における密度推定によって実現される。この種の手法はエネルギーベースモデルと等価である。このようなモデルの学習には、モデルからサンプルを反復的に抽出して対数尤度の勾配を推定する必要があるため、このような方法では計算コストが大幅に高くなる。
Flow-matching
- Flow-matcingは原子の力を訓練に必要とせず、同時にサンプリング効率を高めるという点でハイブリッドなアプローチである。この方法には二つの訓練の段階がある。最初は粗視化密度をaugumented normalizing flowによってモデル化する。次に、粗視化分子動力学シミュレーションで使用できる決定論的なCG力場を抽出するために、force-matchingのような目的を持つ第二の学習段階が必要となる。
- Kohlerはflow-mathcingがいくつかの折り畳みの早いタンパク質でパフォーマンスが改善されることを示した。しかし、学習された粗視化モデルは全原子モデルの熱力学的性質の再現という点で十分な精度があるとはいえず、より大きなタンパク質に対しては不安定を引き起こしてしまう。
今回の論文：拡散モデルによる粗視化構造の生成モデル
-
本研究では最近有名になったデノイジング拡散モデルを活用した。これはタンパク質の構造予測や配座生成、ドッキングで目覚ましい成果を上げている。
-
本研究では粗視化平衡分布から粗視化構造を生成するためのスコアベースの生成モデルを訓練している。スコアベースの拡散モデルと力場、分子動力学についての接続に重点を置き、標準的なノイズ除去損失とconservative scoreでこのような生成モデルを学習すると、独立同分布の粗視化構造の生成と粗視化力場として動力学シミュレーションに直接利用できる単一のモデルが得られることを示す。
-
また本論文の手法は単一の訓練プロセスで学習が実行できるというメリットに加え、この方法は小規模から中規模程度のタンパク質に対する粗視化平衡分布の再現性と、全原子シミュレーションで観察されるタンパク質のフォールディングのような動的なメカニズムの再現性が向上していることを示す。同時に我々の拡散粗視化モデルがflow-matchingの手法よりも大きなタンパク質へのスケーリングが可能であることも示す。
Background
粗視化と粗視化力場
- 粗視化は次元削減マップ
\Xi:\mathbb{R}^{3N} \rightarrow \mathbb{R}^{3n} \Xi \mathbf{x} \in \mathbb{R}^{3N} \mathbf{z} \in \mathbb{R}^{3n} n \ll N - 分子系において、粗視化マップは通常線形であり、
\Xi \in \mathbb{R}^{3n \times 3N} - ある温度
T q(\mathbf{x}) \propto \exp(-U(\mathbf{x}/k_BT)) U(\mathbf{x}) k_B \mathbf{z} \mathbf{x} \mathbf{z}
- ここで
\delta(\dot) - 追加の定数項を追加することでこの分布は熱力学的に一貫な平均力の粗視化ポテンシャル
V(\mathbf{z})
- しかし残念ながら、積分の計算は一般的に性質がよくない。ゆえに熱力学的に一貫した有効粗視化ポテンシャルを近似する方法が様々提案されている。次では簡単に二つの主要なアプローチについてまとめる。
Variational force matching
- Noidらは粗視化マッピング
\Xi -\nabla_{\mathbf{z}}V(\mathbf{z}) -\nabla_{\mathbf{x}}V(\mathbf{x}) \Xi
- ここで
\Xi_f \Xi - Noidはこの関係が熱力学的に一貫した粗視化ポテンシャル$V_{\theta}(\mathbf{z})の近似に使うことができ、次の変分損失を最小化することで得られる。
Relative entropy minimization
- 粗視化力場を学習するもう一つの方法はrelative entropy minimizationであり、これは密度分布の最適化が平均力ポテンシャル関数の最適化と明確につながっているからである。具体的にはrelative entropy minimizationの最小化やKLダイバージェンスにより粗視化密度を推定する。これらは
q(\mathbf{z}) - 粗視化力の推定は最適化されたモデル密度
p_\theta(\mathbf{z})
- variational force matchingとは違い、相対エントロピー最小化は粗視化マップや原子の力を訓練の際に必要としないという特徴がある。歴史的には、
p_\theta V_\theta
- 相対エントロピーを最小化するためには、モデルの自由エネルギーか勾配評価のためのモデルからの独立同時分布サンプリングのいずれかが必要となり、高次元の問題に対して適用するのは非現実的となる。
- 明示的な密度を用いたモデルとしてもう一つの方法がnormalizing-flowであり、最大尤度の密度推定や力場の学習のシンプルな学習が可能である。
- しかし、表現力豊かな可逆関数の学習は難しく、Kohlerらがaugmented normalizing flowに対して最適化した。相補的な乱数の導入により、flowの表現力は向上し、得られる最小化関数は相対エントロピーの変分上界となる。さらに、augumented normalizing flowからの粗視化力を確率的に推定することが可能となる。
- 粗視化動力学をシミュレートするための決定的な粗視化力を抽出するために、Kohlerは変分力マッチングに似た教師と生徒のセットアップを提案した。このtwo stageのアプローチはflow-mathchingの吹き替えともいえる。
Discussion