Accelerating 3D Molecule Generation via Jointly Geometric Optimal Transport

TL;DR

This paper introduces GOAT, a novel 3D molecule generation framework leveraging geometric optimal transport, which significantly accelerates the process while maintaining high quality, by estimating and purifying optimal couplings in a joint equivariant space.

Contribution

The paper proposes a new geometric transport formula and a coupling purification mechanism for fast, high-quality 3D molecule generation using optimal transport theory.

Findings

GOAT doubles the speed of previous methods.

It achieves higher validity, uniqueness, and novelty in generated molecules.

Theoretical and empirical results confirm the effectiveness of optimal coupling estimation.

Abstract

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.