[論文レビュー] Geometric Latent Diffusion Models for 3D Molecule Generation

Minkai Xu, Alexander K. Powers+3 more|arXiv (Cornell University)|May 2, 2023

Model Reduction and Neural Networks被引用数 32

ひとこと要約

GeoLDM はSE(3)-等変な潜在空間（不変スカラーと等変テンソル）を用いた潜在拡散フレームワークを3D分子生成に導入し、従来手法より有効性と制御性を向上させる。

ABSTRACT

Generative models, especially diffusion models (DMs), have achieved promising results for generating feature-rich geometries and advancing foundational science problems such as molecule design. Inspired by the recent huge success of Stable (latent) Diffusion models, we propose a novel and principled method for 3D molecule generation named Geometric Latent Diffusion Models (GeoLDM). GeoLDM is the first latent DM model for the molecular geometry domain, composed of autoencoders encoding structures into continuous latent codes and DMs operating in the latent space. Our key innovation is that for modeling the 3D molecular geometries, we capture its critical roto-translational equivariance constraints by building a point-structured latent space with both invariant scalars and equivariant tensors. Extensive experiments demonstrate that GeoLDM can consistently achieve better performance on multiple molecule generation benchmarks, with up to 7\% improvement for the valid percentage of large biomolecules. Results also demonstrate GeoLDM's higher capacity for controllable generation thanks to the latent modeling. Code is provided at \url{https://github.com/MinkaiXu/GeoLDM}.

研究の動機と目的

• 回転・並進の等変性を尊重するジオメトリを意識した潜在空間で3D分子生成を前進させる。
• 不変成分と等変成分の両方を生成する潜在コードを出力する幾何自己エンコーダを導入する。
• 潜在空間で拡散モデルを訓練し、3D分子の有効性・多様性・制御可能性を向上させる。
• latent GeoLDM が QM9 および DRUG ベンチマークで生成指標を改善し、制御可能な生成を実現することを示す。
• 再現とさらなる研究のための公開コードを提供する。

提案手法

• 3D分子幾何を潜在空間へマッピングし、ノードごとに不変成分（z_h）と等変成分（z_x）を持つ潜在空間を得るための幾何自己エンコーダをEGNNを用いて構築する。
• 潜在遷移を等変性に設計し、潜在事前分布を不変とすることで、結果として得られる周辺分布がSE(3)-不変となるようにしてSE(3) 不変性を保証する。
• 時系列条件付きのEGNNsとして実装されたdenoisingネットワークε_θを用い、構造化された潜在空間上で潜在拡散モデル（LDM）を訓練する。
• 二段階訓練を採用：まず再構成と正則化を用いて幾何AEを訓練し、次にエンコーダを固定して潜在DMを訓練する；生成にはp_θ(z_x,z_h)を用い、デコードにはp_xi(x,h|z_x,z_h)を適用する。
• デノイザーとエンコーダ／デコーダを性質ベクトルsで条件付けすることにより、クロスアテンションではなく連結による制御可能な生成をサポートする。
• 任意に、トレーニング分布p(N)から分子サイズNをサンプルし、それに応じて生成することも可能。

Figure 1: Illustration of GeoLDM . The encoder ${\mathcal{E}}_{\phi}$ encodes molecular features ${\mathbf{x}},{\mathbf{h}}$ into equivariant latent variables ${\mathbf{z}}_{\textnormal{x}},{\mathbf{z}}_{\textnormal{h}}$ , and the latent diffusion transitions $q({\mathbf{z}}_{{\textnormal{x}},t},{\m

実験結果

リサーチクエスチョン

• RQ1 幾何を意識した潜在空間（不変成分と等変成分の両方を持つ）上で動作する潜在拡散モデルは、3D分子生成の品質を向上させるか。
• RQ2 潜在空間でSE(3) 等変性を課すと、RAWな原子特徴空間で動く潜在モデルと比べて化学的有効性と構造的多様性が向上するか。
• RQ3 GeoLDM は標的化学特性の制御可能な生成をどの程度実現できるか、標準的な3D分子ベンチマークでの性能はどうか。
• RQ4 QM9 と DRUG データセットに対して、条件なしと条件付き生成の両方で GeoLDM はどのように性能するか。

主な発見

データ	原子 Sta (%)	分子 Sta (%)	有効性 (%)	有効性と一意性 (%)	原子 Sta (%)	有効性 (%)
QM9	ENF	85.0	4.9	40.2	39.4	G-Schnet	95.7	68.1	85.5	80.3	GDM	97.0	63.2	-	-
QM9	GDM-aug	97.6	71.6	90.4	89.5	EDM	98.7	82.0	91.9	90.7
QM9	EDM-Bridge	98.8	84.6	92.0*	90.7	GraphLDM	97.2	70.5	83.6	82.7
QM9	GraphLDM-aug	97.9	78.7	90.5	89.5	GeoLDM	98.9	89.4	93.8	92.7
DRUG	GeoLDM	84.4	99.3

• GeoLDM はベンチマーク全体で生成品質を向上させ、QM9/DRUG の大規模生体分子で有効率を最大7%高く達成する。
• 不変変数と等変変数の両方を潜在変数として用いた潜在モデリングは、スカラーのみのモデルと比べて有効性と多様性が優れる。
• 条件付き生成実験では、潜在空間設計と単純な性質連結による条件付けにより、制御性が改善される。
• GeoLDM は学習された周辺分布において強いSE(3)-不変性を示し、回転・平行移動への一般化を支援する。
• 実証的な結果は、GeoLDM が複数の指標（有効性、安定性、有効性×一意性）で従来の3D分子生成法を上回ることを示している。
• 著者らは再現とさらなる探究のためのコードをプロジェクトサイトに提供している。

Figure 2: Molecules generated by GeoLDM trained on QM9 (left three) and DRUG (right four).