3D分子潜在拡散モデリングのための統一潜在空間に向けて

Towards Unified Latent Space for 3D Molecular Latent Diffusion Modeling

March 19, 2025

著者: Yanchen Luo, Zhiyuan Liu, Yi Zhao, Sihang Li, Kenji Kawaguchi, Tat-Seng Chua, Xiang Wang

cs.AI

要旨

3D分子生成は創薬や材料科学において極めて重要であり、原子種、化学結合、3D座標といった複雑なマルチモダリティを処理するモデルが求められています。主要な課題は、これらの異なる形状のモダリティを統合しつつ、3D座標に対するSE(3)等価性を維持することです。これを実現するため、既存のアプローチでは通常、不変性と等価性を持つモダリティに対して別々の潜在空間を維持しており、学習とサンプリングの両方において効率が低下しています。本研究では、3D分子の潜在拡散モデリングのための統一変分オートエンコーダ（UAE-3D）を提案します。これは、3D分子を統一された潜在空間から潜在シーケンスに圧縮しつつ、ほぼゼロの再構成誤差を維持するマルチモーダルVAEです。この統一潜在空間により、潜在拡散モデリングを行う際のマルチモダリティと等価性の処理の複雑さが解消されます。これを実証するため、分子帰納バイアスを持たない汎用拡散モデルであるDiffusion Transformerを潜在生成に採用しました。GEOM-DrugsとQM9データセットでの大規模な実験により、本手法がデノボおよび条件付き3D分子生成の両方において新たなベンチマークを確立し、優れた効率性と品質を達成することが示されました。

English

3D molecule generation is crucial for drug discovery and material science, requiring models to process complex multi-modalities, including atom types, chemical bonds, and 3D coordinates. A key challenge is integrating these modalities of different shapes while maintaining SE(3) equivariance for 3D coordinates. To achieve this, existing approaches typically maintain separate latent spaces for invariant and equivariant modalities, reducing efficiency in both training and sampling. In this work, we propose Unified Variational Auto-Encoder for 3D Molecular Latent Diffusion Modeling (UAE-3D), a multi-modal VAE that compresses 3D molecules into latent sequences from a unified latent space, while maintaining near-zero reconstruction error. This unified latent space eliminates the complexities of handling multi-modality and equivariance when performing latent diffusion modeling. We demonstrate this by employing the Diffusion Transformer--a general-purpose diffusion model without any molecular inductive bias--for latent generation. Extensive experiments on GEOM-Drugs and QM9 datasets demonstrate that our method significantly establishes new benchmarks in both de novo and conditional 3D molecule generation, achieving leading efficiency and quality.