# Accelerating 3D Molecule Generation via Jointly Geometric Optimal Transport

arXiv: https://arxiv.org/abs/2405.15252 · HTML: https://mlsanity.com/paper/2405.15252/

## Key facts

- Cited by 11 papers in the map (data through 14 Oct 2026).
- Placed in Flow matching, time series and science, one of 14 directions in Diffusion models.
- Builds on 26 of its 43 references that are in the map.
- First posted on arXiv 24 May 2024.

## Placed in

- Diffusion models / [Flow matching, time series and science](https://mlsanity.com/diffusion/flow-matching-time-series-and-science/)

## Builds on (most cited first, 26 of 26)

- [Denoising Diffusion Probabilistic Models](https://mlsanity.com/paper/2006.11239/) (2020)
- [Score-Based Generative Modeling through Stochastic Differential Equations](https://mlsanity.com/paper/2011.13456/) (2020)
- [Flow Matching for Generative Modeling](https://mlsanity.com/paper/2210.02747/) (2022)
- [Neural Ordinary Differential Equations](https://mlsanity.com/paper/1806.07366/) (2018)
- [Flow Straight and Fast: Learning to Generate and Transfer Data with Rectified Flow](https://mlsanity.com/paper/2209.03003/) (2022)
- [Variational Inference with Normalizing Flows](https://mlsanity.com/paper/1505.05770/) (2015)
- [E(n) Equivariant Graph Neural Networks](https://mlsanity.com/paper/2102.09844/) (2021)
- [Improving and generalizing flow-based generative models with minibatch optimal transport](https://mlsanity.com/paper/2302.00482/) (2023)
- [Equivariant Diffusion for Molecule Generation in 3D](https://mlsanity.com/paper/2203.17003/) (2022)
- [Junction Tree Variational Autoencoder for Molecular Graph Generation](https://mlsanity.com/paper/1802.04364/) (2018)
- [GraphAF: a Flow-based Autoregressive Model for Molecular Graph Generation](https://mlsanity.com/paper/2001.09382/) (2020)
- [Crystal Diffusion Variational Autoencoder for Periodic Material Generation](https://mlsanity.com/paper/2110.06197/) (2021)
- [Cormorant: Covariant Molecular Neural Networks](https://mlsanity.com/paper/1906.04015/) (2019)
- [GEOM: Energy-annotated molecular conformations for property prediction and molecular generation](https://mlsanity.com/paper/2006.05531/) (2020)
- [Geometric Latent Diffusion Models for 3D Molecule Generation](https://mlsanity.com/paper/2305.01140/) (2023)
- [Constrained Graph Variational Autoencoders for Molecule Design](https://mlsanity.com/paper/1805.09076/) (2018)
- [Symmetry-adapted generation of 3d point sets for the targeted discovery of molecules](https://mlsanity.com/paper/1906.00957/) (2019)
- [E(n) Equivariant Normalizing Flows](https://mlsanity.com/paper/2105.09016/) (2021)
- [Flow Matching in Latent Space](https://mlsanity.com/paper/2307.08698/) (2023)
- [Equivariant flow matching](https://mlsanity.com/paper/2306.15030/) (2023)
- [Diffusion-based Molecule Generation with Informative Prior Bridges](https://mlsanity.com/paper/2209.00865/) (2022)
- [Diffusion Bridge Mixture Transports, Schr"odinger Bridge Problems and Generative Modeling](https://mlsanity.com/paper/2304.00917/) (2023)
- [MolDiff: Addressing the Atom-Bond Inconsistency Problem in 3D Molecule Diffusion Generation](https://mlsanity.com/paper/2305.07508/) (2023)
- [Equivariant Flow Matching with Hybrid Probability Transport](https://mlsanity.com/paper/2312.07168/) (2023)
- [SE(3) Equivariant Augmented Coupling Flows](https://mlsanity.com/paper/2308.10364/) (2023)
- [An Equivariant Generative Framework for Molecular Graph-Structure Co-Design](https://mlsanity.com/paper/2304.12436/) (2023)

## Built on by (most cited first, 11 of 11)

- [Flow Matching Meets Biology and Life Science: A Survey](https://mlsanity.com/paper/2507.17731/) (2025)
- [Geometric Representation Condition Improves Equivariant Molecule Generation](https://mlsanity.com/paper/2410.03655/) (2024)
- [Feynman-Kac-Flow: Inference Steering of Conditional Flow Matching to an Energy-Tilted Posterior](https://mlsanity.com/paper/2509.01543/) (2025)
- [Rethinking Diffusion Models with Symmetries through Canonicalization with Applications to Molecular Graph Generation](https://mlsanity.com/paper/2602.15022/) (2026)
- [Predictive Feature Caching for Training-free Acceleration of Molecular Geometry Generation](https://mlsanity.com/paper/2510.04646/) (2025)
- [AbFlow : End-to-end Paratope-Centric Antibody Design by Interaction Enhanced Flow Matching](https://mlsanity.com/paper/2602.07084/) (2026)
- [GAGA: Gaussianity-Aware Gaussian Approximation for Efficient 3D Molecular Generation](https://mlsanity.com/paper/2507.09043/) (2025)
- [FlashMol: High-Quality Molecule Generation in as Few as Four Steps](https://mlsanity.com/paper/2605.07020/) (2026)
- [Toward Better Geometric Representations for Molecule Generative Models](https://mlsanity.com/paper/2605.07693/) (2026)
- [Generating Physically Consistent Molecules with Energy-Based Models](https://mlsanity.com/paper/2605.18381/) (2026)
- [Adopting Reinforcement Learning with Verifiable Rewards for Molecular Generation](https://mlsanity.com/paper/2607.19044/) (2026)

## Abstract (arXiv)

This paper proposes a new 3D molecule generation framework, called GOAT, for fast and effective 3D molecule generation based on the flow-matching optimal transport objective. Specifically, we formulate a geometric transport formula for measuring the cost of mapping multi-modal features (e.g., continuous atom coordinates and categorical atom types) between a base distribution and a target data distribution. Our formula is solved within a joint, equivariant, and smooth representation space. This is achieved by transforming the multi-modal features into a continuous latent space with equivariant networks. In addition, we find that identifying optimal distributional coupling is necessary for fast and effective transport between any two distributions. We further propose a mechanism for estimating and purifying optimal coupling to train the flow model with optimal transport. By doing so, GOAT can turn arbitrary distribution couplings into new deterministic couplings, leading to an estimated optimal transport plan for fast 3D molecule generation. The purification filters out the subpar molecules to ensure the ultimate generation quality. We theoretically and empirically prove that the proposed optimal coupling estimation and purification yield transport plan with non-increasing cost. Finally, extensive experiments show that GOAT enjoys the efficiency of solving geometric optimal transport, leading to a double speedup compared to the sub-optimal method while achieving the best generation quality regarding validity, uniqueness, and novelty. The code is available at https://github.com/WanyuGroup/ICLR2025-GOAT.
