論文解説：粗視化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}$の確率密度を次のようにあらわに表現することができる。

$$q(z) = \frac{\int \exp(-U(\mathbf{x})/k_B T) \delta(\Xi(\mathbf{x}) - \mathbf{z}) d\mathbf{x}}{\int \exp(-U(\mathbf{x}')/k_B T) d\mathbf{x}'}$$

• ここで$\delta(\dot)$はディラックのデルタ関数を表す。
• 追加の定数項を追加することでこの分布は熱力学的に一貫な平均力の粗視化ポテンシャル$V(\mathbf{z})$として一義に定義できる。

$$V(\mathbf{z}) = -k_B T \log q(\mathbf{z}) + cst. = -k_B T \log\int e^{-U(\mathbf{x}/k_B T} \delta(\Xi(\mathbf{x}) - \mathbf{z}) + cst.$$

• しかし残念ながら、積分の計算は一般的に性質がよくない。ゆえに熱力学的に一貫した有効粗視化ポテンシャルを近似する方法が様々提案されている。次では簡単に二つの主要なアプローチについてまとめる。

Variational force matching

• Noidらは粗視化マッピング$\Xi$のある制約の存在下で、粗視化力場$-\nabla_{\mathbf{z}}V(\mathbf{z})$と原子の力場$-\nabla_{\mathbf{x}}V(\mathbf{x})$とのより扱いやすい一貫性が得られることを示した。具体的には、$\Xi$が線形写像であり、各ビーズには一つ以上の原子が含まれ、各原子は特定のビーズに対してのみ非ゼロ係数を持つ場合、以下が成り立つ。

$$-\nabla_{\mathbf{z}}V(\mathbf{z})=\mathbb{E}_{q(\mathbf{x}|\mathbf{z})}\lbrack\Xi_{f}(-\nabla_\mathbf{x}U(\mathbf{x}))\rbrack$$

• ここで$\Xi_f$は線形写像であり、これらの係数は粗視化マップ$\Xi$の線形係数と関係している。
• Noidはこの関係が熱力学的に一貫した粗視化ポテンシャル$V_{\theta}(\mathbf{z})の近似に使うことができ、次の変分損失を最小化することで得られる。

$$\mathbb{E}_{q(\mathbf{x},\mathbf{z})}\lbrack ||\nabla_{\mathbf{z}}V_{\theta}(\mathbf{z})-\Xi_{f}(-\nabla_\mathbf{x}U(\mathbf{x}))||^2\rbrack$$

Relative entropy minimization

• 粗視化力場を学習するもう一つの方法はrelative entropy minimizationであり、これは密度分布の最適化が平均力ポテンシャル関数の最適化と明確につながっているからである。具体的にはrelative entropy minimizationの最小化やKLダイバージェンスにより粗視化密度を推定する。これらは$q(\mathbf{z})$から有限にサンプルされるときいんは最大尤度の最適化という点で等価である。
• 粗視化力の推定は最適化されたモデル密度$p_\theta(\mathbf{z})$から次のように抽出することができる。

$$-\nabla_{\mathbf{z}}V_{\theta}(\mathbf{z}) \propto \nabla_{\mathbf{z}}\log p_theta(\mathbf{z})$$

• variational force matchingとは違い、相対エントロピー最小化は粗視化マップや原子の力を訓練の際に必要としないという特徴がある。歴史的には、$p_\theta$の非正規化されたバージョンが粗視化ポテンシャル$V_\theta$の直接パラメータ化としてモデル化されており、この場合次のように表現できる

$$p_\theta(\mathbf{z}) \propto \exp (-V_\theta(\mathbf{z})/k_BT))$$

• 相対エントロピーを最小化するためには、モデルの自由エネルギーか勾配評価のためのモデルからの独立同時分布サンプリングのいずれかが必要となり、高次元の問題に対して適用するのは非現実的となる。
• 明示的な密度を用いたモデルとしてもう一つの方法がnormalizing-flowであり、最大尤度の密度推定や力場の学習のシンプルな学習が可能である。
• しかし、表現力豊かな可逆関数の学習は難しく、Kohlerらがaugmented normalizing flowに対して最適化した。相補的な乱数の導入により、flowの表現力は向上し、得られる最小化関数は相対エントロピーの変分上界となる。さらに、augumented normalizing flowからの粗視化力を確率的に推定することが可能となる。
• 粗視化動力学をシミュレートするための決定的な粗視化力を抽出するために、Kohlerは変分力マッチングに似た教師と生徒のセットアップを提案した。このtwo stageのアプローチはflow-mathchingの吹き替えともいえる。