[論文レビュー] Learning Joint 2D & 3D Diffusion Models for Complete Molecule Generation

Han Huang, Leilei Sun+2 more|arXiv (Cornell University)|May 21, 2023

Computational Drug Discovery Methods被引用数 5

ひとこと要約

本論文では、分子の原子種、形式的酸化数、結合種別、3次元座標を併せて生成するため、分子グラフと三次元幾何構造の相関をモデル化するための新規な拡散グラフトランスフォーマーを活用した、2次元と3次元を統合した拡散モデルJODOを提案する。このモデルはQM9およびGEOM-Drugsにおいて、有効性、FCD、スケルトン類似度の面でベースラインを上回り、状態打破性能を達成している。また、少ないステップでの高速サンプリングと逆分子設計を可能としている。

ABSTRACT

Designing new molecules is essential for drug discovery and material science. Recently, deep generative models that aim to model molecule distribution have made promising progress in narrowing down the chemical research space and generating high-fidelity molecules. However, current generative models only focus on modeling either 2D bonding graphs or 3D geometries, which are two complementary descriptors for molecules. The lack of ability to jointly model both limits the improvement of generation quality and further downstream applications. In this paper, we propose a new joint 2D and 3D diffusion model (JODO) that generates complete molecules with atom types, formal charges, bond information, and 3D coordinates. To capture the correlation between molecular graphs and geometries in the diffusion process, we develop a Diffusion Graph Transformer to parameterize the data prediction model that recovers the original data from noisy data. The Diffusion Graph Transformer interacts node and edge representations based on our relational attention mechanism, while simultaneously propagating and updating scalar features and geometric vectors. Our model can also be extended for inverse molecular design targeting single or multiple quantum properties. In our comprehensive evaluation pipeline for unconditional joint generation, the results of the experiment show that JODO remarkably outperforms the baselines on the QM9 and GEOM-Drugs datasets. Furthermore, our model excels in few-step fast sampling, as well as in inverse molecule design and molecular graph generation. Our code is provided in https://github.com/GRAPH-0/JODO.

研究の動機と目的

• 既存の生成モデルが2次元分子グラフや3次元幾何構造を個別にモデル化するという限界を是正すること。これにより、生成品質と後続の応用が制限される。
• 2次元結合グラフと3次元座標を統合的にモデル化する統一された生成フレームワークの開発。これにより、補完的となる分子記述子を捉える。
• 分子構造と幾何構造の同時分布を学習することで、サンプルの忠実性を向上させ、効率的な逆分子設計を可能とする。
• 幾何的対称性（SE(3)不変性）とグラフ順列不変性を保持する、微分可能で連続時間の拡散プロセスの設計。
• 後処理や追加シミュレーションに依存せずに、量子性質を条件とした高速かつ高品質なサンプリングと条件付き生成を実現すること。

提案手法

• 2次元グラフと3次元座標の両方の成分を同時に破損させる統一されたノイズスケジュールを有する連続時間の拡散モデルを提案。
• 関係性に依存するアテンションメカニズムを用いて、ノードおよびエッジ表現を統合的に更新する拡散グラフトランスフォーマーを導入。スカラー特徴量と幾何ベクトルを統合。
• ノイズ予測ではなく、データ予測モデル（ノイズ予測とは逆）を採用。最後のサンプリングステップの出力を補助的条件として用いる自己条件付き学習により、離散的構造の回復を改善。
• メッセージパッシングアーキテクチャを用いてデータ予測モデルをパラメータ化。分子グラフ全体にわたり、スカラー特徴量と3次元ベクトル特徴量を伝搬・更新。
• 3次元幾何構造生成におけるSE(3)不変性と2次元グラフ学習における順列不変性を、等変性および不変性を備えたニューラルネットワーク設計により強制。
• 生成を目的の量子性質に条件づけることで、逆設計への拡張を実現。これにより、性質を狙った分子生成を可能に。

実験結果

リサーチクエスチョン

• RQ1 2次元と3次元を統合した拡散モデルは、2次元グラフや3次元座標のみを生成するモデルよりも、より高い忠実度で完全な分子を生成できるか？
• RQ2 逆ノイズ除去プロセス中に、分子トポロジーと3次元幾何構造の相関を統一された拡散フレームワークがどの程度効果的にモデル化できるか？
• RQ3 関係性アテンションと幾何ベクトル更新を備えた提案された拡散グラフトランスフォーマーは、生成品質とサンプリング効率を向上させるか？
• RQ4 本モデルは、無条件および条件付き（性質を狙った）分子生成において、状態打破の性能を達成できるか？
• RQ5 既存のベースラインと比較して、本モデルは少ないステップでのサンプリングと量子性質のための逆設計において、どの程度の性能を示すか？

主な発見

• QM9データセットにおいて、JODO-2Dはバリエンシーのチェックなしで99.91％の有効性を達成し、CDGS（98.13％）やDiGress（92.10％）をすべて上回った。
• QM9においてJODO-2DはFCDスコア0.472を達成し、CDGS（2.069）やDiGress（3.597）よりも顕著に低い値を示しており、分布マッチングの面で優れた性能を示している。
• MOSESデータセットにおいて、JODO-2Dは88.9％の有効性、100％の一意性、91.0％の新規性を達成し、DiGress（85.7％の有効性、95.0％の新規性）を上回り、他の強力なベースラインと同等またはそれを上回った。
• 本モデルは逆分子設計において優れた性能を示しており、望ましい量子性質を持つ分子の効果的生成を可能としている。
• JODOは高品質な出力を得るための高速で少ないステップのサンプリングを可能としており、実世界の応用における実用的価値を示している。
• アブレーションスタディにより、自己条件付き学習を備えた提案されたデータ予測モデルと拡散グラフトランスフォーマーが、生成品質と構造的忠実度の向上に顕著な寄与をしていることが確認された。