Towards Unified Latent Space for 3D Molecular Latent Diffusion Modeling
Yanchen Luo, Zhiyuan Liu, Yi Zhao, Sihang Li, Kenji Kawaguchi, Tat-Seng Chua, Xiang Wang
2025-03-21
Summary
This paper is about making AI models better at designing 3D molecules, which is useful for discovering new drugs and materials.
What's the problem?
AI models struggle to handle the different types of information needed to represent molecules, like atom types, bonds, and 3D coordinates, all at the same time.
What's the solution?
The researchers created a new AI model that combines all of this information into a single, unified space, making it easier for the AI to generate new and realistic molecules.
Why it matters?
This work matters because it can speed up the discovery of new drugs and materials by allowing AI to design better molecules.
Abstract
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.