Towards Unified and Lossless Latent Space for 3D Molecular Latent Diffusion Modeling
Yanchen Luo, Zhiyuan Liu, Yi Zhao, Sihang Li, Hengxing Cai, Kenji Kawaguchi, Tat-Seng Chua, Yang Zhang, Xiang Wang
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. On GEOM-Drugs, it reduces FCD by 72.6% over the previous best result, while achieving over 70% relative average improvements in geometric fidelity. Our code is released at https://github.com/lyc0930/UAE-3D/.
Ask about this paper
Your agent reads all of it.
Lune indexed this paper to the last equation, along with the top-tier papers that cite it. Ask a question and the answer quotes them.
Your agent calls
Luneget_paper_fulltext
Free to start. No credit card required.
Terminal
Install the CLIlune papers fulltext f5536fcb-223d-43ed-9b6b-41b1c4e5b4ebCited by top-tier papers2
- Learning 3D Anisotropic Noise Distributions Improves Molecular Force FieldsXixian Liu, Rui Jiao, Zhiyuan Liu, Yurou Liu et al.NeurIPS 2025 · 3 citations
- 3D-GSRD: 3D Molecular Graph Auto-Encoder with Selective Re-mask DecodingChang Wu, Zhiyuan Liu, Wen Shu, Liang Wang et al.NeurIPS 2025 · 1 citation
Builds on29
- Denoising Diffusion Probabilistic ModelsJonathan Ho, Ajay Jain, Pieter AbbeelNeurIPS 2020 · 35,902 citations
- Diffusion Models Beat GANs on Image SynthesisPrafulla Dhariwal, Alexander Quinn NicholNeurIPS 2021 · 13,211 citations
- High-Resolution Image Synthesis with Latent Diffusion ModelsRobin Rombach, Andreas Blattmann, Dominik Lorenz, Patrick Esser et al.CVPR 2022 · 13,123 citations
- Scalable Diffusion Models with TransformersWilliam Peebles, Saining XieICCV 2023 · 5,568 citations
- E(n) Equivariant Graph Neural NetworksVictor Garcia Satorras, Emiel Hoogeboom, Max WellingICML 2021 · 1,432 citations
Related papers
- Geometric Latent Diffusion Models for 3D Molecule GenerationMinkai Xu, Alexander S. Powers, Ron O. Dror, Stefano Ermon et al.ICML 2023 · 252 citations
- All-atom Diffusion Transformers: Unified generative modelling of molecules and materialsChaitanya K. Joshi, Xiang Fu, Yi-Lun Liao, Vahe Gharakhanyan et al.ICML 2025
- UniMoMo: Unified Generative Modeling of 3D Molecules for De Novo Binder DesignXiangzhe Kong, Zishen Zhang, Ziting Zhang, Rui Jiao et al.ICML 2025
- Equivariant Diffusion for Molecule Generation in 3DEmiel Hoogeboom, Victor Garcia Satorras, Clément Vignac, Max WellingICML 2022 · 865 citations
- Unified Generative Modeling of 3D Molecules with Bayesian Flow NetworksYuxuan Song, Jingjing Gong, Hao Zhou, Mingyue Zheng et al.ICLR 2024 · 36 citations
