迈向统一的三维分子潜在扩散建模潜在空间

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

摘要

三维分子生成在药物发现和材料科学中至关重要，要求模型能够处理复杂的多模态信息，包括原子类型、化学键和三维坐标。一个关键挑战在于整合这些不同形态的模态，同时保持三维坐标的SE(3)等变性。为实现这一目标，现有方法通常为不变性和等变性模态分别维护独立的潜在空间，这降低了训练和采样的效率。在本研究中，我们提出了统一变分自编码器用于三维分子潜在扩散建模（UAE-3D），这是一种多模态VAE，它将三维分子压缩至统一潜在空间中的潜在序列，同时保持近乎零的重构误差。这一统一潜在空间消除了在进行潜在扩散建模时处理多模态性和等变性的复杂性。我们通过采用扩散Transformer——一种无任何分子归纳偏见的通用扩散模型——进行潜在生成，来验证这一点。在GEOM-Drugs和QM9数据集上的大量实验表明，我们的方法在从头生成和条件生成三维分子方面显著设立了新的基准，实现了领先的效率与质量。

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.