[論文レビュー] MiDi: Mixed Graph and 3D Denoising Diffusion for Molecule Generation

Clément Vignac, Nagham Osman+2 more|arXiv (Cornell University)|Feb 17, 2023

Machine Learning in Materials Science被引用数 10

ひとこと要約

MiDiはエンドツーエンド拡散によって分子グラフとそれらの3D立体構造を同時に生成し、GEOM-DRUGSにおける従来手法と比較して安定で有効な分子の高い割合を達成します。

ABSTRACT

This work introduces MiDi, a novel diffusion model for jointly generating molecular graphs and their corresponding 3D arrangement of atoms. Unlike existing methods that rely on predefined rules to determine molecular bonds based on the 3D conformation, MiDi offers an end-to-end differentiable approach that streamlines the molecule generation process. Our experimental results demonstrate the effectiveness of this approach. On the challenging GEOM-DRUGS dataset, MiDi generates 92% of stable molecules, against 6% for the previous EDM model that uses interatomic distances for bond prediction, and 40% using EDM followed by an algorithm that directly optimize bond orders for validity. Our code is available at github.com/cvignac/MiDi.

研究の動機と目的

• グラフ構造と3D立体構造を同時に捉える分子のエンドツーエンド生成を促進する。
• グラフと座標の両方を同時にデノイズする拡散ベースのモデルを提案する。
• 新規なrEGNN強化のTransformerアーキテクチャを用いて3D等変処理を強化する。
• GEOM-DRUGSおよびQM9データセットにおいて、無条件生成性能の優越性を示す。

提案手法

• 分子を3Dに埋め込みされたグラフとして表現し、ノード特徴（原子種、電荷）とエッジ特徴（結合種）を持たせる。
• 座標の連続値ガウス分布と離散的特徴（原子種、電荷、結合）を混在させた拡散過程を、成分ごとに特化したノイズスケジュールとともに用いる。
• 座標と結合種を原子種と電荷より優先する適応型ノイズスケジュールを導入する。
• 翻訳不変性に依存しない特徴を含むRelaxed Equivariant Graph Neural Network (rEGNN) 層を開発する。
• SE(3)同変性を維持する Transformer ベースのデノイズネットワークへ rEGNNを組み込む。
• 座標回帰と離散特徴のクロスエントロピーを組み合わせた損失で学習する。

Figure 1 : Samples from our model. MiDi generates simultaneously a 2D graph structure and a 3D conformation that is consistent with this structure.

実験結果

リサーチクエスチョン

• RQ1 拡散モデルはエンドツーエンド微分可能な方法で分子グラフと3D座標を同時に学習できるか。
• RQ2 適応的な成分別ノイズスケジューリングは生成分子の品質と安定性を向上させるか。
• RQ3 緩和SE(3)等変性を持つアーキテクチャ（rEGNN）は標準のEGNNよりグラフ+コンフォーマ生成を改善するか。
• RQ4 QM9およびGEOM-DRUGSにおける無条件分子生成性能を、OpenBabelベースの後処理や従来の拡散法と比較してMiDiはどう結果を出すか？

主な発見

• MiDiはGEOM-DRUGSで安定分子の高い割合を達成し、EDMベースのアプローチとOpen Babel後処理を上回る。 (Adaptive MiDi: Mol stable 91.6%, At stable 99.8%, Validity 77.8%, Connected 100.0%).
• QM9ではMiDiはグラフベース指標でベースEDMモデルを上回るが、Open Babel最適化は単純なデータセットで依然強力。
• MiDiは2Dグラフ構造と3Dコンフォーマの両方をエンドツーエンド微分可能に生成し、別個の結合予測ステップを回避する。
• 適応ノイズスケジュールはサンプル品質を向上させ、座標と結合種を原子種と電荷より優先するようモデルを導く。
• MiDiは大規模な薬物様データセット（GEOM-DRUGS）へスケールし、後処理の結合予測に依存する手法よりもエンドツーエンド性能が高いことを示す。

Figure 2 : The noise schedule is tuned separately for each component. Atom coordinates and bond types are denoised earlier during sampling, while atom types and formal charges are updated later in the process. Experimentally, the adaptive schedule allows to obtain better 3D conformers and more stabl