[論文レビュー] MUDiff: Unified Diffusion for Complete Molecule Generation

Chenqing Hua, Sitao Luan+5 more|arXiv (Cornell University)|Apr 28, 2023

Machine Learning in Materials Science被引用数 15

ひとこと要約

MUDiff は、2Dグラフ（エッジ）と3D座標（原子）双方に拡散を適用して完全な分子表現を共同生成し、回転・平行移動対称性を保つ形でノイズ除去を行う新しい MUformer を用いる。

ABSTRACT

Molecule generation is a very important practical problem, with uses in drug discovery and material design, and AI methods promise to provide useful solutions. However, existing methods for molecule generation focus either on 2D graph structure or on 3D geometric structure, which is not sufficient to represent a complete molecule as 2D graph captures mainly topology while 3D geometry captures mainly spatial atom arrangements. Combining these representations is essential to better represent a molecule. In this paper, we present a new model for generating a comprehensive representation of molecules, including atom features, 2D discrete molecule structures, and 3D continuous molecule coordinates, by combining discrete and continuous diffusion processes. The use of diffusion processes allows for capturing the probabilistic nature of molecular processes and exploring the effect of different factors on molecular structures. Additionally, we propose a novel graph transformer architecture to denoise the diffusion process. The transformer adheres to 3D roto-translation equivariance constraints, allowing it to learn invariant atom and edge representations while preserving the equivariance of atom coordinates. This transformer can be used to learn molecular representations robust to geometric transformations. We evaluate the performance of our model through experiments and comparisons with existing methods, showing its ability to generate more stable and valid molecules. Our model is a promising approach for designing stable and diverse molecules and can be applied to a wide range of tasks in molecular modeling.

研究の動機と目的

• 完全な分子表現のために、2D（グラフの連結性）と3D（幾何情報）という分子データをジョイントに生成するモチベーションを喚起する。
• 連続成分（原子特徴量、座標）と離散成分（エッジタイプ）の両方をノイズ除去する拡散ベースのフレームワークを開発する。
• 2Dと3D情報を統合しつつ、3Dの回転・平行移動対称性を保持する等価変換を持つMUformerを導入する。
• 生成と学習中に2Dまたは3Dデータが欠損しても頑健であることを可能にする。
• 既存手法と比較して生成分子の安定性と多様性が向上することを示す。

提案手法

• MUDiff を提案する：原子特徴量と座標には連続ノイズを、エッジタイプには離散ノイズを加え、全成分を共同でデノイズする拡散モデル。
• 回転・平行移動の制約の下、2Dおよび3D分子データを処理する不変・等価チャンネルを備えた統一型TransformerであるMUformerを導入する。
• 原子特徴量と座標のノイズを予測し、エッジタイプを分類する学習目的を定義する。
• 頑健な表現のために、2D近傍、3D近傍、およびグローバルグラフ特徴を統合するエンコーディング手法を採用する。
• 空間情報を捉え、等価性を保証するために3D特有のRBF（径状基底）とコサインカットオフ関数を適用する。
• 完全にノイズの多い潜在表現から徐々にデノイズして完全な分子へとサンプリングする手順を提供する。

Figure 1: The figure showcases our MUformer for processing 2D and 3D molecular data. Within the Transformer backbone, two channels exist: purple for 2D data and brown for 3D data. The blue part encodes 2D molecular structures, while the green part handles atom-level information and the red part proc

実験結果

リサーチクエスチョン

• RQ1 2Dグラフ構造と3D座標に対する拡散を共同で行うことで、単一表現モデルよりも安定で妥当な分子を生成できるか。
• RQ2 3D回転移動対称性に等価でありつつ、2Dと3D情報を効果的に統合するようなTransformerをどのように設計すればよいか。
• RQ3 生成または学習中に2Dまたは3Dデータが欠損するケースに対してモデルは頑健か。
• RQ4 2Dおよび3D構造を同時にモデル化することで得られる安定性・唯一性の性能向上はどの程度か。

主な発見

• MUDiff は既存手法より7.9%安定な分子を生成する（Sec 6.2）。
• MUDiff は既存手法と比較して分子の一意性を2%向上させる（Sec 6.2）。
• MUformer は原子特徴量、座標、エッジタイプを同時に予測し、回転移動同等性を持つ。
• 3D構造が限られた状態で訓練してもモデルは有効であり、完全な3Dデータで訓練された手法と対して競争力のある性能を示す（Sec 6.1）。
• 2D情報または3D情報のいずれかが欠損しても独立して機能することを支持し、頑健な完全表現を可能にする。