Towards Unified Latent Space for 3D Molecular Latent Diffusion Modeling

2025-03-19

AI Summary

Existing 3D molecular generation methods employ separate latent spaces for atomic types, bond topology, and 3D coordinates, leading to inefficient training and difficulty in enforcing SE(3) equivariance. To address this, we propose the Unified Autoencoding Encoder-3D (UAE-3D), the first model that compresses multimodal molecular representations—atomic identity, connectivity, and Cartesian coordinates—into a single homogeneous latent sequence, thereby eliminating modality coupling while preserving SE(3) equivariance in 3D geometry. Building upon UAE-3D, we design a Diffusion Transformer framework that performs end-to-end latent-space diffusion without requiring explicit molecular priors. Evaluated on GEOM-Drugs and QM9, our method achieves state-of-the-art performance in both de novo and conditional generation: near-zero reconstruction error, significantly improved sampling efficiency, and superior generation quality across all metrics compared to prior approaches.

Technology Category

• Machine Learning: Deep Generative Models & Autoencoders
• Computer Vision: Diffusion Models for Vision
• Knowledge Representation and Reasoning: Geometric, Spatial, and Temporal Reasoning

Application Category

• Search and Retrieval-Augmented AI: Retrieval-Augmented Generation (RAG) and multi-modal RAG
• Graph Algorithms and Modeling for the Web: Graph embeddings and representation learning for Web-related graphs
• Semantics and Knowledge: Methods to enhance, augment, integrate or synergize semantic models such as knowledge graphs and LLMs

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 extbf{U}nified Variational extbf{A}uto- extbf{E}ncoder for extbf{3D} Molecular Latent Diffusion Modeling ( extbf{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 extit{de novo} and conditional 3D molecule generation, achieving leading efficiency and quality.

Problem

• Integrate multi-modal data for 3D molecule generation.
• Maintain SE(3) equivariance in unified latent space.
• Improve efficiency and quality in molecular generation.

Innovation

• Unified latent space for multi-modal 3D molecules
• SE(3) equivariance maintained in latent sequences
• Diffusion Transformer for efficient latent generation