[論文レビュー] Navigating the Design Space of Equivariant Diffusion-Based Generative Models for De Novo 3D Molecule Generation

Tuan Le, Julian Cremer+3 more|arXiv (Cornell University)|Sep 29, 2023

Machine Learning in Materials Science被引用数 6

ひとこと要約

本論文は、連続的な原子座標、カテゴリカルな原子/結合特徴、時間依存損失重み付けを活用することで、QM9およびGEOM-Drugsで最先端の性能を達成し、訓練および推論が高速化された、3次元デノボ分子生成のための新規なE(3)-等長性拡散モデルEQGAT-diffを紹介する。さらに、暗黙的な水素を含むPubChem3Dで事前学習することで転移性を示し、ハイブリダイゼーション状態などの化学的に意味のある特徴を組み込むことで生成分子の妥当性が向上する。

ABSTRACT

Deep generative diffusion models are a promising avenue for 3D de novo molecular design in materials science and drug discovery. However, their utility is still limited by suboptimal performance on large molecular structures and limited training data. To address this gap, we explore the design space of E(3)-equivariant diffusion models, focusing on previously unexplored areas. Our extensive comparative analysis evaluates the interplay between continuous and discrete state spaces. From this investigation, we present the EQGAT-diff model, which consistently outperforms established models for the QM9 and GEOM-Drugs datasets. Significantly, EQGAT-diff takes continuous atom positions, while chemical elements and bond types are categorical and uses time-dependent loss weighting, substantially increasing training convergence, the quality of generated samples, and inference time. We also showcase that including chemically motivated additional features like hybridization states in the diffusion process enhances the validity of generated molecules. To further strengthen the applicability of diffusion models to limited training data, we investigate the transferability of EQGAT-diff trained on the large PubChem3D dataset with implicit hydrogen atoms to target different data distributions. Fine-tuning EQGAT-diff for just a few iterations shows an efficient distribution shift, further improving performance throughout data sets. Finally, we test our model on the Crossdocked data set for structure-based de novo ligand generation, underlining the importance of our findings showing state-of-the-art performance on Vina docking scores.

研究の動機と目的

• 既存の拡散モデルが3次元分子生成において大きな分子で性能が低く、データが限られているという問題を解決すること。
• E(3)-等長性拡散モデルの設計空間を体系的に探求し、連続的 vs. 離散的状態空間と損失重み付け戦略に焦点を当てる。
• 時間依存損失重み付けと強化された特徴工学を活用して、訓練収束、推論速度、サンプル品質の向上を図ること。
• 暗黙的な水素を含む大規模なPubChem3Dで事前学習することで、小規模なデータセットでの効率的なファインチューニングを可能にすること。
• ハイブリダイゼーション状態や芳香族性といった化学的に意味のある特徴を拡散プロセスに組み込むことで、分子の妥当性を向上させること。

提案手法

• 連続的な3次元原子座標とカテゴリカルな原子/結合タイプを、ノイズ除去拡散フレームワーク内でモデル化するE(3)-等長性グラフ自己注意ネットワークであるEQGAT-diffを提案する。
• 信号対ノイズ比（SNR）に基づく時間依存損失重み付けを採用し、訓練収束の高速化とサンプル品質の向上を図る。
• 連続座標、離散的原子タイプ、離散的結合タイプのハイブリッドデータモダリティを導入し、それぞれに別々のノイズ除去ヘッドを設ける。
• DDIMサンプリングを用い、離散的特徴の忠実度を保ちながら推論速度を向上させるために、修正された離散的更新を適用する。
• 暗黙的な水素を含む大規模なPubChem3Dデータセットで、無条件の事前学習を実施し、下流タスクでの効率的なファインチューニングを可能にする。
• ハイブリダイゼーション状態や芳香族性といった化学的に意味のある特徴を、入力ノードおよびエッジ特徴に組み込むことで、化学的妥当性を向上させる。

実験結果

リサーチクエスチョン

• RQ1 3次元分子拡散モデルにおいて、時間依存損失重み付け（例：SNRベース）は、一様または離散的損失重み付けと比較して、訓練収束およびサンプル品質にどのように影響するか？
• RQ2 連続的原子座標と離散的表現の両者を用いることで、モデルの性能と推論速度にどのような影響が生じるか？
• RQ3 大規模な3次元分子データセット（例：PubChem3D）で事前学習したモデルは、より小さな複雑なデータセットに対して、最小限の追加訓練で効果的にファインチューニング可能か？
• RQ4 ハイブリダイゼーションや芳香族性といった化学的に意味のある特徴は、拡散ベースの生成における分子の妥当性と多様性にどのように影響するか？
• RQ5 パrameterizationの選択（例：$x_0$ と $x_t$ の比較）は、3次元分子拡散モデルの性能と安定性にどの程度影響を与えるか？

主な発見

• EQGAT-diffはQM9およびGEOM-Drugsで最先端の性能を達成し、QM9では96.8%の妥当性率と96.6%の成功率を記録し、先行モデルを上回る。
• 切り捨てられたSNR(t)に基づく時間依存損失重み付けは、訓練収束とサンプル品質を顕著に向上させ、100ステップでのみ使用する場合、推論時間を最大5倍短縮する。
• 100ステップで学習しSNRベースの重み付けを適用したモデルは、500ステップで一様重み付けで学習したモデルを上回る性能を示し、適応的損失スケジューリングの有効性を裏付ける。
• PubChem3Dで事前学習したEQGAT-diffを、Crossdockedなどのターゲットデータセットでファインチューニングすることで、わずかなファインチューニングイテレーションで優れた性能が得られ、分布シフトへの効率的な適応が可能になる。
• ハイブリダイゼーション状態と芳香族性の特徴を組み込むことで、分子の妥当性と多様性が向上し、GEOM-Drugsでは42.2%の多様性スコアを達成し、ベースラインモデルよりも顕著に高い。
• DDIMサンプリングは500ステップ未満にステップ数を減らした場合、サンプル品質の向上に寄与しなかったが、100ステップで学習しSNRベースの重み付けを適用したモデルは、5倍の高速推論を実現しながらも、高い性能を維持した。