[論文レビュー] GeoDiff: a Geometric Diffusion Model for Molecular Conformation Generation

Minkai Xu, Lantao Yu+4 more|arXiv (Cornell University)|Mar 6, 2022

Machine Learning in Materials Science被引用数 173

ひとこと要約

GeoDiffは原子座標を直接モデル化し、SE(3)-等元拡散フレームワークを用いて分子の立体構造を生成する。これにより、特に大規模分子で立体配置生成における最新性能を達成する。

ABSTRACT

Predicting molecular conformations from molecular graphs is a fundamental problem in cheminformatics and drug discovery. Recently, significant progress has been achieved with machine learning approaches, especially with deep generative models. Inspired by the diffusion process in classical non-equilibrium thermodynamics where heated particles will diffuse from original states to a noise distribution, in this paper, we propose a novel generative model named GeoDiff for molecular conformation prediction. GeoDiff treats each atom as a particle and learns to directly reverse the diffusion process (i.e., transforming from a noise distribution to stable conformations) as a Markov chain. Modeling such a generation process is however very challenging as the likelihood of conformations should be roto-translational invariant. We theoretically show that Markov chains evolving with equivariant Markov kernels can induce an invariant distribution by design, and further propose building blocks for the Markov kernels to preserve the desirable equivariance property. The whole framework can be efficiently trained in an end-to-end fashion by optimizing a weighted variational lower bound to the (conditional) likelihood. Experiments on multiple benchmarks show that GeoDiff is superior or comparable to existing state-of-the-art approaches, especially on large molecules.

研究の動機と目的

• 分子の3D原子座標を直接出力する拡散型生成モデルの動機づけと開発。
• SE(3)等価設計とCoM-free不変初期化によって回転・並進不変性を保証する。
• 分子グラフを条件としたエンドツーエンド学習可能なフレームワークを、条件付き生成のための適切なELBO目的で提供する。

提案手法

• 3D座標 C^t の前方拡散を固定分散スケジュールで定式化し、C^T ~ N(0, I) に至る。
• 分子グラフ G を条件としてノイズを徐々に除去する条件付き逆マルコフ連鎖 p_θ(C^{t-1}|G, C^t) を定義する。
• 回帰平均 μ_θ をSE(3)等価性を持つニューラルネットワーク（Graph Field Network）を用いてパラメータ化し、回転・並進等価性を確保する。
• 翻訳：CoM-free不変初期密度とCoM-freeマルコフ核を用いて平行移動不変性を保証する。回転等価性は μ_θ とネットワークアーキテクチャを通じて課す。
• εとε_θの間の加重L2損失に還元される加重ELBOで訓練し、等価性を課す二つの戦略（アライメントベースと連鎖規則アプローチ）を採用する。
• C^T ~ N(0, I) から開始して C^{t-1} を順次サンプリングする p_θ(C^{t-1}|G, C^t) を用いたサンプリングアルゴリズムを提供する。

実験結果

リサーチクエスチョン

• RQ1 3D座標上で直接動作する拡散モデルは、不変ジオメトリ中間体に依存するモデルと比較して、競争力のあるまたは優れたコンフォマー生成性能を達成できるか。
• RQ2 幾何生成において訓練時と推論時の回転・並進不変性をどのように保証できるか。
• RQ3 分子の立体配置の逆拡散ステップをモデル化する有効なパラメータ化とアーキテクチャ（例：SE(3)-等価ネットワーク）とは何か。
• RQ4 エンドツーエンドの座標ベース拡散生成は、より大きな薬物様分子へスケールし、多様で高品質なコンフォマーを提供するか。

主な発見

• GeoDiffはGEOM-Drugs上で一貫して最先端のMLベースラインを上回り、COV-RおよびMAT-R指標で顕著な向上を示す。
• 大規模分子に対して基準と比べて高いCOV-Rと低いMAT-Rを達成し、網羅性と精度の両方を示す。
• GeoDiff-Cは一般にGeoDiff-Aよりわずかに最適化性能が良く、連鎖則ベースの訓練が学習をより適切に導くことを示唆する。
• 本手法は生成されたコンフォマーの多様性と質を大幅に向上させ、特に大規模分子が評価されるDrugsデータセットで顕著な効果を示す。
• 定性的サンプルはGeoDiffが局所的およびグローバルな構造特徴の両方を効果的に捉えることを示している。