Research map / Probabilistic and scientific modeling / Variational and generative sampling
Accelerating 3D Molecule Generation via Jointly Geometric Optimal Transport
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 GOAT, a flow-matching framework for 3D molecule generation that jointly transports atom coordinates and features in an equivariant latent space, estimating and purifying optimal couplings to cut sampling steps to 90.
Problem
Fast and effective 3D molecule generation with flow matching: existing diffusion models need ~1000 steps, and prior flow methods transport atom coordinates and features along incompatible, non-optimal paths.
Approach
Map multi-modal atom coordinates and features into a joint continuous latent space via an equivariant autoencoder, compute optimal molecule transport (Hungarian permutation, Kabsch rotation, centering), then estimate and purify optimal distribution couplings using the trained flow to train a straight-line flow with non-increasing transport cost.
Key results
- On QM9, GOAT reaches 90 sampling steps versus 200 for EquiFM and 1000 for diffusion baselines, with novelty 78.6% and significance 72.3%.
- GOAT reduces geometric transport cost by nearly 89.65% relative to EquiFM, per the paper.
- On GEOM-DRUG, GOAT reaches 96.2% validity with 90 steps, versus 99.3% for GeoLDM at 1000 steps.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Use GOAT instead of EquiFM | 3D molecule generation · QM9 | novelty: 57.4 -> 78.6 | Helped |
| Use GOAT instead of EquiFM | 3D molecule generation · QM9 | significance: 53.7 -> 72.3 | Helped |
| Use GOAT instead of EquiFM | 3D molecule generation · QM9 | sampling steps: 200 -> 90 | Helped |
| Use GOAT instead of GeoLDM | 3D molecule generation · GEOM-DRUG | S-Time: 12.84 -> 0.94 | Helped |
| Use GOAT instead of GeoLDM | 3D molecule generation · GEOM-DRUG | validity: 99.3 -> 96.2 | Hurt |
| Add the equivariant autoencoder (ablation) | GOAT · QM9 | significance: 53.9 -> 72.3 | Helped |
| Add optimal distribution transport (ablation) | GOAT with λ=0.5 · QM9 | significance: 62.5 -> 72.3 | Helped |
| Add optimal distribution transport (ablation) | GOAT with λ=0.5 · QM9 | transport cost: 4.13 -> 3.14 | Helped |
| Add optimal molecule transport (ablation) | GOAT with λ=0.5 · QM9 | significance: 58.1 -> 72.3 | Helped |
| Set λ=0.5 rather than λ=1, λ=0.75, λ=0.25, or λ=0 (ablation) | GOAT without optimal distribution transport · QM9 | transport cost: 5.01; 4.63; 4.13; 4.87; 5.41 | Helped |
| Add optimal distribution transport (ablation) | GOAT · QM9 | sampling steps: 100 -> 90.0 | Helped |
| Remove optimal distribution transport (ablation) | GOAT · QM9 | training speed | Helped |
Surprising findings
- On QM9, GOAT's validity (90.0%) is lower than EquiFM (94.7%) and several baselines, so its gains are concentrated in novelty and significance.
- The paper reports GeoBFN at 1000 steps achieves higher validity than GOAT on QM9, and EDM-style diffusion models have near-zero molecule stability on GEOM-DRUG.
Limitations
- Optimal transport computation (Hungarian permutation and Kabsch rotation) is computationally intensive and adds training cost; training takes about 7 days on QM9 and 20 days on GEOM-DRUG.
- Coupling purification relies on RDKit validity filtering, and the theorem's guarantee holds in latent space rather than the original molecule space.
- Validity on QM9 is lower than EquiFM and several baselines.
Experimental setup
Datasets: QM9 GEOM-DRUG
Benchmarks: QM9 GEOM-DRUG
Metrics: atom stability validity uniqueness novelty significance sampling steps S-Time transport cost
Baselines: G-Schnet ENF EDM GDM-aug EDM-Bridge GeoLDM GeoBFN EquiFM
Builds on: flow matching EquiFM EGNN Kabsch algorithm Hungarian algorithm
Closest papers by meaning
- Unified Generative Modeling of 3D Molecules with Bayesian Flow Networks (ICLR 2024 · similarity 0.89)
- Energy-Based Flow Matching for Generating 3D Molecular Structure (ICML 2025 · similarity 0.88)
- 3D Molecule Generation from Rigid Motifs via $\mathrm{SE}(3)$ Flows (NeurIPS 2026 · similarity 0.87)
- NEAT: Neighborhood-Guided, Efficient, Autoregressive Set Transformer for 3D Molecular Generation (NeurIPS 2026 · similarity 0.87)
- Geometric-Facilitated Denoising Diffusion Model for 3D Molecule Generation (AAAI 2024 · similarity 0.87)
- Compositional Flows for 3D Molecule and Synthesis Pathway Co-design (ICML 2025 · similarity 0.86)
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.