Research map / Image generation / Autoregressive and latent generation
Towards Unified and Lossless Latent Space for 3D Molecular Latent Diffusion Modeling
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Proposes UAE-3D, a multi-modal VAE that encodes 3D molecules (atoms, bonds, coordinates) into a unified near-lossless latent sequence, and UDM-3D, a plain Diffusion Transformer trained on it for 3D molecule generation.
Problem
3D molecule generation must jointly model discrete atom/bond features of different shapes and SE(3)-equivariant coordinates; prior latent diffusion models keep separate latent spaces for invariant and equivariant modalities, which reduces efficiency and consistency, and often have non-trivial recons
Approach
Train a Relational-Transformer-based VAE with a Transformer decoder and SE(3) random-transformation augmentations (instead of baked-in equivariant architectures) to compress all molecular modalities into one unified latent sequence with near-zero reconstruction error, then run a general-purpose DiT on it.
Key results
- On GEOM-Drugs, UDM-3D reduces FCD to 0.692 from prior best (2.523 JODO), a 72.6% reduction, and cuts bond length MMD by 88.4%.
- UAE-3D reconstructs QM9 and GEOM-Drugs molecules with 100% atom and bond accuracy and coordinate RMSD of 0.0002 Å and 0.0008 Å respectively.
- UDM-3D achieves lowest MAE on five of six QM9 conditional properties and training is 2.7x faster than JODO (52 h total on A100).
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| use a unified rather than separated latent space (ablation) | UDM-3D · QM9 de novo | FCD: 0.351 -> 0.130 | Helped |
| add both translation and rotation augmentations (ablation) | UAE-3D reconstruction · QM9 | coordinate RMSD: 1.4 × 10^{-4} -> 0.2 × 10^{-4} | Helped |
| add rotation augmentation (ablation) | UAE-3D reconstruction · QM9 | RMSD (× 10^-4 Å): 1.4 -> 1.1 | Helped |
| add translation augmentation (ablation) | UAE-3D reconstruction · QM9 | RMSD (× 10^-4 Å): 1.4 -> 0.9 | Helped |
| reduce latent dimension from 16 to 4 (ablation) | UDM-3D · QM9 de novo | V&U&N: 0.950 -> 0.566 | Hurt |
| increase latent dimension from 16 to 32 (ablation) | UDM-3D · QM9 de novo | V&U&N: 0.950 -> 0.918 | Hurt |
| increase bonded-distance loss weight from 0 to 10 (ablation) | UAE-3D reconstruction · QM9 | coordinate RMSD (Å): 0.0034 -> 0.0002 | Helped |
| reduce R-Trans depth from 6 to 3 layers (ablation) | UAE-3D reconstruction · QM9 | RMSD: 0.00203 -> 0.00326 | Hurt |
| increase R-Trans depth from 6 to 9 layers (ablation) | UAE-3D reconstruction · QM9 | RMSD: 0.00203 -> 0.00229 | No clear effect |
| increase R-Trans depth from 6 to 12 layers (ablation) | UAE-3D reconstruction · QM9 | RMSD: 0.00203 -> 0.00209 | No clear effect |
| replace DiT with a vanilla Transformer or PerceiverIO (ablation) | UDM-3D · QM9 de novo | generation metrics | Hurt |
| use translation and rotation augmentation rather than no augmentation (ablation) | UDM-3D · QM9 de novo | FCD: 0.581 -> 0.130 | Helped |
Surprising findings
- A general-purpose DiT with no molecular inductive bias, trained on a VAE latent learned via SE(3) augmentation alone, matches or beats equivariant diffusion models on 3D metrics.
Limitations
- Molecule size is sampled from the training-set size distribution with a separate module rather than generated autoregressively.
- Evaluated only on unconditional generation and quantum-property conditional generation; structure-based drug design is out of scope.
Experimental setup
Datasets: QM9 GEOM-Drugs
Benchmarks: QM9 GEOM-Drugs
Metrics: FCD atom stability molecule stability V&C V&U V&U&N SNN Frag Scaf FCD3D
Baselines: CDGS JODO MiDi EQGAT-diff EDM GeoLDM MDM ADiT SemlaFlow EEGSDE GeoBFN
Builds on: latent diffusion models Relational Transformer Diffusion Transformer GeoLDM EDM
Closest papers by meaning
- Scalable Non-Equivariant 3D Molecule Generation via Rotational Alignment (ICML 2025 · similarity 0.91)
- Navigating the Design Space of Equivariant Diffusion-Based Generative Models for De Novo 3D Molecule Generation (ICLR 2024 · similarity 0.91)
- Geometric-Facilitated Denoising Diffusion Model for 3D Molecule Generation (AAAI 2024 · similarity 0.90)
- All-atom Diffusion Transformers: Unified generative modelling of molecules and materials (ICML 2025 · similarity 0.89)
- Equivariant Blurring Diffusion for Hierarchical Molecular Conformer Generation (NeurIPS 2024 · similarity 0.89)
- NEAT: Neighborhood-Guided, Efficient, Autoregressive Set Transformer for 3D Molecular Generation (NeurIPS 2026 · similarity 0.88)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.