拡散モデルはタンパク質相互作用の「エネルギー」を学習しているのか?― DFMDockから学習エネルギーを取り出す試み

論文タイトル

Can We Extract Physics-like Energies from Generative Protein Diffusion Models?

出典

Can We Extract Physics-like Energies from Generative Protein Diffusion Models?
Diffusion models have emerged as the state-of-the-art method in generative AI and have shown great success in image synthesis, video generation, molecular desig...

要旨

本研究は、タンパク質構造生成に用いられる拡散モデル(Diffusion Model)が、物理学におけるエネルギー関数に相当する情報を内部的に学習しているかを検証した研究です。著者らは、拡散モデルから負の対数尤度(−log p)を取り出し、それを学習されたエネルギーとして解釈する理論と実装を提案しました。

解説など

近年、RFDiffusionやAlphaFold3などの拡散モデルがタンパク質設計や構造予測で大きな成功を収めています。しかし、これらのモデルがなぜ高性能なのか、そして物理学でいう「自由エネルギー」のような概念を内部で学習しているのかは十分に理解されていません。

本研究では、

  • 拡散モデルが学習した確率分布
  • 統計熱力学におけるエネルギー
  • タンパク質ドッキングの評価関数

の関係を理論的に整理し、「学習されたエネルギー」を拡散モデルから直接抽出できるかを調べました。

基本的な考え方

統計熱力学では、ある構造が観測される確率とエネルギーはボルツマン分布によって結び付けられます。

論文では、

  • 出現確率が高い構造
  • エネルギーが低い構造

は本質的に同じ情報を表していると考えます。

そこで著者らは、

負の対数尤度(Negative Log-Likelihood, NLL)

を学習エネルギーとして解釈する枠組みを提案しました。

拡散モデルは通常、構造そのものの確率を直接出力しません。その代わりに「score」と呼ばれる確率勾配

を学習します。scoreは、”どちらに動かせば、より本物らしいタンパク質複合体になるか” を示す矢印です。著者らはこの score を統計力学の力(force)と対応付け、拡散経路上で積分することでエネルギーを復元できると考えました。score の矢印を、拡散過程の道すじに沿って足し合わせると、“どれくらい本物らしい領域に来たか”を数値化できます。その数値をエネルギーとして読む、というのが本研究の考え方です。

提案された2つの手法

1. DiffLikelihood

実際に生成時にたどる拡散経路(SDE trajectory)を利用して尤度を計算する方法です。

2. FlowLikelihood

確率流ODE(Probability Flow ODE)を利用する方法です。

検証の結果、どちらも真の確率分布をほぼ再現できています。特にFlowLikelihoodでは、ほぼノイズのない滑らかな確率推定が得られています。

タンパク質ドッキングモデルDFMDockへの適用

次に著者らは、自ら開発したタンパク質ドッキング用拡散モデルDFMDockに同じ考え方を適用しました。

DFMDockは、

  • 既知の複合体構造にノイズを加える
  • ノイズ除去を学習する
  • 推論時には未結合状態から複合体構造を生成する

というモデルです。

本研究では並進自由度のみを扱う簡略版を用いて解析を行っています。

得られた結果

DFMDockが正しくドッキングできたケースでは、

  • 実験構造付近でエネルギーが最小になる
  • いわゆる「energy funnel」が現れる

ことが確認されました。

これはRosettaなどの物理ベースのエネルギー関数で観察される現象とよく似ています。

さらに、生成したドッキング構造のランキング性能を比較したところ、

  • 25ケース中9ケースで
  • 学習エネルギーがRosetta interface energyと同等またはそれ以上の性能

を示したと報告されています。

これは、拡散モデル内部に実際に有用なエネルギー地形が形成されている可能性を示唆しています。

一方で本研究では、ドッキングの確からしさとの関係性を示したのみで、結合強度との相関は確認していません。著者らも学習エネルギーは物理的自由エネルギーと必ずしも一致しないという立場をとっています。