Accelerating 3D Molecule Generation via Jointly Geometric Optimal Transport
Authors: Haokai Hong
Wanyu Lin
KC Tan
Venue: ICLR 2025 Poster (2025)
Keywords: Molecule Generation
Flow Matching
Fast Generation
Primary area: applications to physical sciences (physics, chemistry, biology, etc.)
TL;DR: 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.

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.

Published as a conference paper at ICLR 2025
ACCELERATING 3D MOLECULE GENERATION VIA JOINTLY GEOMETRIC OPTIMAL TRANSPORT
Haokai Hong, Wanyu Lin; Kay Chen Tan Department of Data Science and Artificial Intelligence
Department of Computing The Hong Kong Polytechnic University, Hong Kong SAR, China.
haokai. hong@connect polyu. hk, {wan- lin, ketan]@polyu edu .hk
ABSTRACT
This paper proposes 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 coor- dinates and categorical atom types) between a base distribution and target data distribution. Our formula is solved within a joint, equivariant, and smooth repre- sentation space. This is achieved by transforming the multi-modal mooaal features into a continuous latent space with equivariant networks. In addition, we find that iden- tifying 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 estimated optimal transport plan for fast 3D molecule generation. The purification filters out the subpar molecules ensure the ultimate generation quality. We theoretically and empirically prove that the proposed opti- mal 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 github.
INTRODUCTION
The problem of 3D molecule generation is essential in various scientific fields, such as material science, biology, and chemistry (Hoogeboom et al., 2022; Watson et al., 2023; Xie et al., 2022). Typically, 3D molecules can be represented atomic geometrie graphs (Hoogeboom et al., 2022; Xu et al., 2023; Song et al., 2023b), where each atom/node is embedded in the Cartesian coordinates and encompasses rich features, such as atom types and charges. There has been fruitful research progress on geometric generative modeling for facilitating the process of 3D molecule generation. Specifically, geometric generative models are proposed to estimate the distribution of complex ge- ometries and are used for generating feature-rich geometries. The success of these generative mod- eling mainly comes from the advancements in using the notion of probability measurement transport for generating samples. Generative modeling aims to generate samples via mapping a simple prior distribution, e.g., Gaussian, to a desired target probability distribution. This mapping process can be termed distribution transport/geeraative problem (Peluchetti, 2023).
Recent representative models for sampling 3D molecules in silicon include diffusion-based mod- els (Hoogeboom et al., 2022: Wu et al., 2022; Xu al., 2023) and flow matching- based mod- els (Song et al., 2023a; Klein et al., 2023). Diffusion-based models have shown superior results on molecule generation tasks (Xu et al., 2023; Hoogeboom
al., 2022; Jung al., 2024). They simulate a stochastic differential equation (SDE) to transport a base distribution (e.g., Gaussian) to the data distribution. However, a major drawback of diffusion-basse models is their slow inference
"ooreeponniigg author.

---

Published as a conference paper at ICLR 2025
Transport
Diffusion-Baan Basee
Hybrid
Jointly (Ours)
Atom Coordinates
Atom Features f
Molecule in
Distribution p
Figure I: The Illustration of Probability Paths based on Different Molecule Generative Models. 1. The diffusion path (Hoogeboom et al., 2022; Xu et al., 2023), which typically deviates from a straight line map, necessitates a large number of sampling steps. 2. The hybrid transport (Song et al., 2023a) ensures straight transport for atomic coordinates, but it does not guarantee the same for atom features. Furthermore, this cost does not consider the optimal distribution couplings, leading to suboptimal transport between distributions. 3. GOAT simultaneously considers the optimal transport for atom coordinates and features, providing a joint and straight path for fast sampling.
speed with the learned stochastic transport trajectory (Hoogeboom et al., 2022; Wu et al., 2022; Xu et al., 2023); they typically need approximately 1, 000 sampling steps to produce valid samples. This could make large-scale inference prohibitively expensive. Accordingly, flow matching built upon continuous normalizing flows - has emerged as a new paradigm that could potentially provide effective density estimation and fast inference (Lipman et al., 2022; Tong et al., 2023; Dao et al., 2023; Liu et al., 2022; Klein et al., 2023; Song et al., 2023a).
This paper aims to deal with the problem of and effective 3D molecule generation based on flow matching optimal transport principle. In particular, our objective is to obtain a distribution transport trajectory with optimal transport cost and generation quality regarding molecule validity, unique- ness, and novelty. A few recent works have been proposed to improve the sampling efficiency of the geometrical domains via flow matching principle. (Klein al., 2023) was proposed for efficient, simulation- training of equivariant continuous normalizing flows, which can produce samples from the equilibrium Boltzmann distribution of molecule in Cartesian coordinates. However, can only work for many-body molecular systems and does not consider atomic features.
In the context of molecule generation, properly characterizing the transport cost to optimize over is indispensable and challenging. There are two main challenges. Firstly, the multi-modal property of molecular space, typically consisting of continuous atom coordinates and categorical atom types, makes the transport cost measurement hard to optimize. Secondly, the optimal transport problem essentially is to search optimal distribution couplings for mapping. This process typically requires similarity computation of the two distributions. However, the various sizes of geometric graphs to transport introduce difficulties in evaluating the distribution similarity. The closest to ours is EquiFM (Song et al., 2023a), which attempts to address the multi-modality issue by using different probability paths to transport each modality separately. The proposed equivariant optimal transport (OT) for transporting atom coordinates indeed forms straight-line trajectory for training, while the sariance-preserving principle could not ensure a straight-line trajectory for atom features. Therefore, the fused flow paths might deviate strongly from the OT paths and could not ensure optimal coupling between two probability measurements, leading to high computational costs and numerical errors.
In this work, we propose a new framework for fast 3D molecule generation based on a novel and principled optimal transport flow-mater not-miaching objective, dubbed as Geometric OptimAl Transport (GOAT). In particular, we formulate a geometric transport cost measurement for optimally transport- ing continuous atom coordinates and categorical features, which is inherently a Bilevel optimization problem. To deal with the first challenge induced by transporting multiple modalities, GOAT lever- ages a latent variable model equipped with equivariant networks to map the multi-modal features into a joint, equivalent, and smooth representation space. This equivariant latent variable model has been proven to be flexible and expressive for modeling complex 3D molecules (Satorras al., 2021; Xu et al., 2023). A latent flow matching then operates over the latent space, which can provide distributional coupling estimation.
2

---

Published as a conference paper at ICLR 2025
To tackle the second challenge obtaining the optimal distribution couplings, we propose frain from directly working with distribution similarity computation. Specifically, we propose to rectify the flow with the estimated ones based on the latent flow matching. Because the estimated distributional couplings are identified based on the synthesized samples, they might deviate from the real-world samples in terms of quality. We provide purification process for high-quality samples regarding validity, uniqueness, and novelty. With this process, we can turn arbitrary couplings into deterministic and causal ones, leading to the optimal transport path for fast and effective generation.
We theoretically demonstrate that optimal coupling estimation and purification result in non- increasing geometric transport costs, Moreover, we empirically highlight the superiority of GOTA by conducting experiments on widely used benchmarks. The proposed method reduced the trans- port cost by nearly 89.65%, halving the sampling times compared to EquiFM. terms of generation quality, our method achieves up to a 17.1% % improvement over existing algorithms.
2 PROBLEM SETUP
(3D) molecule with N atoms can be represented geometric graph denoted as = (x, h), where x = (XI,..., Xx) E Rg represents the atom coordinates and h (h),..., hy) Rx is the atom features containing atomic types, charges, etc. d is the dimensionality of the atom features. A zero center-of-mase -fass (Zero CoM) space is defined as X = {x E RN3 1/2 Σ x' = 0}. indicating that the mean of the N atoms' coordinates should be O. In what follows, we will introduce some necessary concepts, including flow matching and optimal transport, to facilitate the definition of our problem.
Notations. A
General Flow Matching'. Given noise X E RN ~ Po and data x1 E R ~ pr. the general flow- based model considers mapping f R + R RN + RVV which defines an ordinary differential equation (ODE): dx = u, (x)ddt. Continuous normalizing flows (CNFs) were introduced with black-box ODE solvers to train approximate ut- However, CNFs are hard to train as they need numerous evaluations of the vector field.
smooth varying vector field u [0, 1] x
Flow matching (Lipman et al., 2022), simulation-fra frree approach for training CNFs, is pro- posed to regress the neural network (x(x, t) to some target vector field (x): (FM( := Ex-2((0,1) soll/loxx, () uz (x)) where px(x) is the corresponding probability path which conditioned on X1 ~ P1 and then defined as Pr (x(x1) tion, common probability paths include variance exploding (VE) diffusion path (Song et al., 2020), variance preserving (VP) diffusion path (Ho et al., 2020), and straight transport path (Lipman al.,
I Pc (x((x)))) JPi (x1) dxx1 In implementa-
2022; Liu et al., 2022).
Optimal Transport (OT). Transport plan between Po and pi is also called coupling and we denoted it as P1) (Ambrosio et al., 2021). OT addresses the problem of finding the optimal coupling that minimizes the transport cost, typically in the form of R[c(x1 (0x-xx)) xo)]. where e : R - is a cost function, such as c(-) || 1|a. Formally, a coupling I'(po.)) P1) is called optimal only if it achieves the minimum value of E[c(x; - (x) xo) E I'po,, among all couplings that share the same marginals. An ideal optimal transport trajectory is map of optimal couplings.
Geometric Optimal Transport. In our task, we consider a pair of geometric probability distri- butions over T x(+4)) with densities p(gn) and p(gi) (or denoted as po and pı). Geometric generative modeling considers the task of fitting a mapping f from (3++d) to Nx that transforms go to g1. More specifically, if go is distributed with density po then f(g) is distributed with density Pt- Typically, Po is an easily sampled density, such as a Gaussian.
In our specific problem, beyond geometric distribution transport, we concentrate on fast. 3D molecule generation based on the flow model with optimal transport (OT), which has been proven effective in accelerating non- non-geometric flow models (Liu et al., 2022; Tong et al., 2023). Therefore, our problem is defined as geometric optimal transport flow matching.
We are aware of the drawbacks of reusing the notation x, which represents general data point here. "Geometric probability distribution denotes the molecular data distribution and it is to reflect the geometric property of molecules as opposed non-geometric data, such as text.
3

---

Published as a conference paper at ICLR 2025
3 OUR METHOD: GOAT
3.1 FORMULATING GEOMETRIC OPTIMAL TRANSPORT PROBLEM
Our objective is to obtain an optimal transport trajectory for fast 3D molecule geometry generation based on optimal transport flow models. In this regard, we can reformulate our problem into search- ing for optimal coupling for geometric optimal transport. Specifically. involves the transportation of molecules via optimal coupling, where each atom follows straight and shortest path. In other words, each molecule coupled with a noisy sample that incurs the minimum cost, and each atom within the target molecule is paired with its counterpart in the noise, leading to the minimum cost. To encapsulate the optimal scenario, we consider the problem of geometric optimal transport with two components: 1) optimal molecule transport (OMT) with equivariant OT for atom coordinates and invariant OT for atom features; 2) optimal distribution transport (ODT) with optimal molecule coupling.
Geometric Transport Cost. Transporting a molecule includes transforming atom coordinates and features. We can depict the molecule transport cost below:
(g g1) = | xolla l + ||b; - hol12,
(1)
where go Pa and gi P1- In addition, given coupling [(po. P1). we measure the distribution similarity between two distributions denoted as Po and p1 based on the probability transport cost as follows:
C =B|cg (go g1)], (go+ g1) ) ['po.
(2)
Geometric Optimal Transport Problem. Building upon the established transport cost, we can formulate the geometric optimal transport problem for fast and effective 3D molecule generation. In particular, molecule remains invariant for any rotation, translation, and permutation, while the transport cost is not invariant or equivariant to the above operations. Therefore, there exists minimum molecule transport cost with 1) optimal rotation and translation transformations such that the molecules from the data and noise are nearest each other at the atomic coordinate level and 2) optimal permutation transformation such that the atomic features of the two are also nearest.
We supplement the detailed analysis of equivariance and invariance in Appendix A and present geometric optimal transport as follows:
min Bieg (go g1)],
s. (go. g1) (pooPi):
(g0, 81) min [Rx] + t, Rx² + t,...,,,A RR, + -xxx x...,.,. hN) (h) bg....,, 1l2, VT, R, and t,
x² Xl)
R,t
+ (1 - X) min || h?.....
where ?g denotes optimal molecule transport cost, π represents a permutation of N elements, trade-off coefficient balancing the transport costs of atom coordinates and atom features for searching optimal permutation (#). Additionally, R and denote a rotation matrix and translation, respectively. The defined geometric optimal transport problem forms a bi-level optimization prob- lem that involves two levels of optimization tasks. Specifically, minimizing molecule transport cost nested inside the optimizing distribution transport cost.
The Challenges of Solving the Geometric Optimal Transport Problem. First, optimal molecule transport involves searching for a unified optimal permutation for atom coordinates and features with minimum transport cost. The paths for transporting continuous coordinates and categorical features are incompatible and require sophisticated, hybrid modeling multi-modal variables (Song al., 2023a), leading to sub-optimal solution. Second, molecular distribution comprises molecules with diverse numbers atoms, introducing difficulties quantifying the transport cost for search- ing optimal coupling. As result, the minimization of geometric transport C within molecular distributions poses a more significant challenge compared to other domains such as computer vi- sion (Tong et al., 2023) many-body systems (Klein et al., 2023). Moreover, the proposed geomet- ric optimal transport problem, which involves a nested optimization structure, presents significant computational challenge for optimization.
4

---

Published as a conference paper at ICLR 2025
3.2 SOLVING GEOMETRIC OPTIMAL TRANSPORT PROBLEM
In this section, we will address the above challenges under the depicted problem for fast and effec- tive 3D molecule generation from two aspects, including optimal molecule transport and optimal distribution transport.
3.2.1 SOLVING OPTIMAL MOLECULE TRANSPORT
As illustrated in Eq. (3), our objective for optimal molecule transport is to find R. i, and á that minimize the transport cost denoted as cy for two given molecule geometries represented as go and gi:
R, t, ी arg min Cg( g1), Vn, R, and t. R.t.
(4)
As per Eq. (3), the optimal molecule transport problem entails the consideration of both atom co- ordinates and features for comprehensive representations of 3D molecules. Though coordinates and features represent different modalities, they need to be considered in tandem. Previous search (Song et al., 2023a; Klein et al., 2023) either solely focused on equivariant optimal transport for coordinates transported atom features via distinct yet non-optimal path. In contrast, propose to unify the transport of atom coordinates and features. If we can unify different modalities within an equivariant and smooth representation space, optimal transport from base distribution to the data distribution, leading to fast and effective molecule generation, is possible.
Specifically, we map the atom coordinates and features from the data space into latent space with an equivariant autoencoder (Satorras et al., 2021), which enables us to compute a unified optimal permutation (+).
After the mapping, we ascertain the optimal rotation (R) for two atomic coordinate sets for trans- portation, utilizing the Kabsch algorithm (Kabsch, 1976) as did in (Song et al., 2023a; Klein et al., 2023). The computed rotation matrix ensures that the coordinates of the target molecules in the la- tent space are in closest proximity to the noise molecules. In addition, to achieve optimal translation (t), we establish the data distribution p1 and base distribution Po in the zero CoM space (Hoogeboom al., 2022; Xu et al., 2023). Equipped with the equivariant autoencoder, the latent representation also resides in the zero CoM space, thus ensuring optimal translation in the latent space.
Implementation, Initially, the distributions Po and pr are aligned to the Zero CoM by subtracting the center of gravity using t.
Subsequently, an equivariant autoencoder is designed to project gi ~ pi into the latent space. Here, the encoder E. transforms gi into latent domain z1 (g1), where z1 = (Zx,1 E RN Zh,1 NN and k represents the latent dimensionality for the atomic features. The decoder D. then learns to decode 2 back to molecular domain formulated as gi (z)). The equivariant autoen- coder can be trained by minimizing the reconstruction objective, which is d(D(E(g))) g1). With the encoded z and the sampled noise Zg = (zx.u E IR x
Zh.0 E R NNk) from Po. we then formu-
late the atom-level cost matrix M., [ the 2-norm distance between i-th atom of 21 and j-th atom of 2, including the latent coordinates Zx and the latent features Zh- With Mc,- the optimal permutation á is induced with the Hungarian algorithm (Kuhn, 1955). The coordinates of the noise molecule Zx,0 and the latent coordinates of the target molecule 1 are then aligned through rotation solved by the Kabsch algorithm (Kabsch, 1976). In summary, we perform the above-calculated translation, encoding, rotation, and permuta- tion on the target molecule gı to obtain 21, which forms the optimal molecule transport cost with the sampled noise Zg. The complete process is denoted as 2 = π ( i)).
|z| 2/]12 = | 2 12 ||2,, 2h,01 which
3.2.2 SEARCHING OPTIMAL COUPLING FOR OPTIMAL DISTRIBUTION TRANSPORT
By solving Eq, (4), can obtain R, é, and स yielding optimal molecule transport trajectory straight one given two data points from the base distribution and target distribution, respectively (see the gray trajectory in Figure 2). Nevertheless, ensuring straight trajectory does not necessarily yield optimal transport for generative modeling because a straight map for two data points does not indicate straight map between two distributions. Figure 2 depicts two possible trajectories for
5

---

Published as a conference paper at ICLR 2025
generative modeling. The gray one shows straight but not the shortest map, while the red one represents the shortest map, indicating an optimal trajectory. As discussed in Sec. 3.1, an optimal trajectory can only be achieved with optimal coupling, leading the shortest path for mapping the base distribution to the target distribution. To approximate optimal coupling and further boost the sampling speed, introduce the second part of framework, optimal flow estimation and purification, which is dedicated to solving optimal distribution transport.
The pathway to optimal distribution transport is identify the optimal coupling [(poo P1) that satisfies the condition formulated as:
Gaussian Noise
Distribution Po
f arg min Eiey (Zo. z1)]
r
(5)
S. t. Zo. z1) E (ppo.
where I' is an arbitrary coupling plan between Po and Pi and 6g is optimal molecule transport cost defined Eq. (3).
However, measuring the distribution transport cost for searching for optimal coupling challenging due to the various sizes of molecules. Inspired by (Liu al., 2022), we circumvent the need to quantify transport costs by estimating the optimal coupling, (po,p1 based on a trained flow model with the ini- tial coupling denoted as I'(po.P)..
Target Data Distribution P1
Straight But Not Optimal
Straight and Optimal
The estimated optimal coupling minimize the transport cost but may introduce generation error for the following reasons. The first type of error arises from estimating the flow path between Po to P1 via a neural network vp. implying that p/. characterized by ve, does not perfectly match the data distribu- tion Pl. The second type of error stems from the discreteness of molecular data and the continuity of the distribution. In essence, intermediate value between two similar and valid molecules, which are closely distributed, may not be biochemically valid. To compensate for such discrepancies, we implement purification process on the generated coupling ensure effective generation. We present detailed implementation below.
An Illustration of the Difference Between Straight Coupling and Optimal Cou- pling. GOAT approximates optimal coupling for fast generation.
Figure 2:
Implementation. First, based on Sec. 3.2.1, we can obtain a set noise and target molecule pairs with optimal molecule transport cost. We leverage the corresponding transport path as the conditional probability path for training the flow with the loss formulated as:
(f)
pt. (1,, (2,, t) (2) -z0)l13
(6)
where 2, = 12) + (1 -1--10 l)zo, E 10, 1]. Compared with using conditional optimal transport path and pariance-preserving path in hybrid fashion (Song et al., 2023a), our method employs a unified linear interpolation of 2, Zo as the flow probability path. Such straight trajectory adheres to the
naive ODE formula denoted as dz, (2) zo )tt, thereby providing a more straight flow path for fast sampling. The optimum of LF1 is achieved when uz(Zz, t) = E Zo Zr].
The proposed framework then samples data pairs (Zo, 2)) via trained flow model 0n the estimated optimal coupling. Specifically, z can be sampled following dz, = V/ (z,, t)dt starting from Zo ~ Po and the process of sampling is denoted as ODE.,, Pair of and z is set as fixed and batch of pairs will generated as the estimated optimal coupling represented as ['po, p1) = {(zo, }. where zi ODE,g (zo). Finally, we estimate and purify the flow. Specifically, z is decoded by D. for gí and evaluated in terms of stability and validity by RdKit (Landrum et al., 2016). This provides a criterion for filtering out invalid molecules to purify the coupling. Subsequently, the optimal flow trained using the loss in Eq. (6) with estimated and purified coupling.
Provably Reduced Geometric Transport Cost. The estimated optimal coupling I can boost gen- eration only when geometric transport cost is reduced. We theoretically show that our approach can indeed reduce geometric transport costs as follows:

---

Published as a conference paper at ICLR 2025
Theorem 3.1. The coupling I't incurs no larger geometric transport cost than the random coupling I I(Po.Pt) in that Elc, (Zo, z1]l s Elćy(Zo) (zo,Zz)), Zı]], where (Zo, zz) I'ppo, (2o, z1) = I'po. P1), and 0g is optimal molecule transport cost.
With this theorem, the proposed GOAT is guaranteed Pareto descent the geometric transport cost, leading to faster generation. A comprehensive proof is given in Appendix B, and the pseu- docode for training and sampling is presented in Appendix C.
4 EXPERIMENTAL STUDIES
Datasets. We evaluate over benchmark datasets for 3D molecule generation, including QM9 (Ra- makrishnan al., 2014) and the GEOM-DRUG (Axelrod & Gómez-Bombarell 2022). QM9 is a standard dataset that contains 130k 3D molecules with up to 29 atoms. GEOM-DRUS encompasses around 450K molecules, each with an average of 44 atoms and a maximum of 181 atoms. More dataset details are presented in Appendix E.
Baselines. We compare GOAT with several competitive baseline models. G-Schnet (Gebauer et al., 2019) and Equivariant Normalizing Flows (ENF) (Chen et al., 2018) are equivariant gen- erative models utilizing the autoregressive models and continuous normalizing flow, respectively. Equivariant Graph Diffusion Model (EDM) and its variant GDM-Aug (Hoogeboom et al., 2022), EDM- -riider (Wu et al.. 2022). GeoLDM (Xu et al., 2023) are diffusion-baan -based approaches. GeoBFN (Song et al., 2023b) leverages Bayesian flow networks for distributional parameter approx- imation. EquiFM (Song et al., 2023a) is the first Now-matching method for 3D molecule generation.
4.1 EVALUATION METRICS
Evaluating Generation Quality. Without loss of generality, we use validity, uniqueness, and nov- elty to evaluate the quality of generated molecules (Zhang et al., 2023). Existing experiments caleu- late validity, uniqueness, and novelty, which are nested; novelty measures novel molecules among unique and valid molecules. However, such calculation cannot reflect the ultimate quality among all samples. We further propose a new metric toward the significance of generative models (Walters & Murcko, 2020).
Below, we provide the detailed definitions of these metrics. 1) Validity. An essential criterion for molecule generation is that the generated molecules must be chemically valid, which implies that the molecules should obey chemical bonds and valency constraints. We use RdKit (Landrum et al., 2016) to check if molecule obeys the chemical valency rules. Validity calculates the percentage of valid molecules among all the generated molecules; 2) Uniqueness. An important indicator of a molecule generative model is whether it can continuously generate different samples, which is quantified by the uniqueness. We evaluate uniqueness by measuring the fraction of unique molecules among all the generated valid ones; 3) Novelty. An ideal generative model for de novo molecule design should be able to generate novel molecular samples that do not exist in the training set. Therefore, we report novelty that quantifies the percentage of novel samples among all the valid and unique molecules; 4) Significance. To comprehensively evaluate the molecule generative models, we represent a new metric, significance, to quantify the percentage of valid, unique, and novel molecules among the generated samples.
Evaluating Generation Efficiency. 1) We report sampling steps to measure the generation speed. The time cost of each sampling step in most baselines, including EDM, EDM-Bridge. GeoBFN, GeoLFM, and EquiFM, is identical because they all applied EGNN (Satorras et al., 2021) with the same layers and parameters. Fewer steps indicate higher generation efficiency. For EquiFM and the proposed GOAT, we applied the same adaptive stepsize on ODE solver Dopri5 (Dormand & Prince, 1980) for fair comparison. 2) Considering the generation quality and efficiency simultaneously. propose report the il (valid, unique, and novel) molecule, denoted as S-Time. The metric is calculated by the number of significant molecules over the total time consumed by generating all molecules, and it comprehensively reflects the performance of generation quality and efficiency. 3) We measure the generation efficiency by comparing geometric transport cost, which is calculated by Eq. (2).
7

---

Published as a conference paper at ICLR 2025
4.2 RESULTS AND ANALYSIS
Table I: Comparisons of generation quality regarding Atom Stability, Validity, Uniqueness, Novelty, and Significance. And comparisons of generation efficiency regarding Steps and Time. The best results are highlighted in bold.
QM9
Quality (1)
Efficiency (L)
Metrics
Atom Sta
Valid
Uniqueness
Novelty
Significance
Steps
S-Time
Data
99.0
97.7
100.0
ENF
85.0
40.2
98.0
G-Schnet
95.7
85.5
93.9
GDM-aug
97.6
90.4
99.0
74.6
66.8
1000
1.50
EDM
98.7
91.9
98.7
65.7
59.6
1000
1.68
EDM-Bridge
98.8
92.0
98.6
1000
GeoLDM
98.9
93.8
98.8
58.1
53.9
1000
1.86
GeoBFN
98.6
93.0
98.4
70.3
64.4
100
0.16
EquiFM 98.9
94.7
57.4
53.7
200
0.37
GOAT (Ours)
99.2
92.9
99.0
78.6
72.3
90
0.12
Table 2: Comparisons of generation quality regarding Atom Stability, Validity, Steps, and Time on GEOM- DRUG. The best results are highlighted in bold.
Total Distribution Transpor Cast
36.73
Arbitmrry Coupling
-4
Average Transprrt Ont for An Alom
60 M
30.25
Optinall Coupling
Oya.nal Coupling
2
40
GEOM-DRUG
Quality (t)
Efficiency (L)
105
-20
Metries
Atom Sta
Valid
Steps
S-Time
20
Data
86.5
99.9
4.13
3.07
GDM-aug
77.7
91.8
1000
EDM
81.3
92.6
1000
14.88
OM9
EmiFM
Uniifed
L-MM
LODT
2OMT
200T
EDM-Bridge
82.4
92.8
1000
GeoLDM
84.4
99.3
1000
12.84
GeoBFN
78.9
93.1
100
1.27
Figure 3: The blue histogram plots the comparisons of distribution transport cost. The red line chart depicts the average trans- port cost atom (best view in color).
EquiFM
84.1
98.9
200
2.02
GOAT (Ours)
84.8
96.2
90
0.94
In this study, we generate 10K molecular samples for each method and compute the aforementioned metries for comparisons. The evaluation results are presented Tables and with Figure 3.
Performance Comparisons with Diffusion-Based Methods. We observe that all diffusion based generation methods indeed need 1000 sampling steps to achieve comparable generation quality. Sur- prisingly, with the least sampling steps, GOAT achieves the best atom stability, uniqueness, novelty, and significance over QM9. Specifically, it improves novelty by up to 35.2% and significance by up to 34.1%, respectively. Among these diffusion models, GeoLDM achieves the best validity per- formance. However, it owns relatively poor novelty and significance, 58. 1% and 53.9% on QM9, respectively. These results indicate that the latent diffusion models model the complex geomet- ric 3D molecules well but introduce a serious overfitting problem generating more molecules that the same as the training samples. Though GDM-Aug can achieve the second-best novelty among all methods, it needs 1000 sampling steps for 3D molecule generation. As for GEOM-DRUG, we directly compare the validity as ultimate significance since all compared methods achieved almost 100% uniqueness (Xu et al., 2023). Table shows that the proposed algorithm also achieves com- petitive performance while maintaining a leading edge in generation speed on such large-scale dataset. Specifically, GOAT only spends 0.94 seconds for each valid, unique, and novel molecule on average and reaches 96.2% validity, while GeoL.M takes more than 10x seconds to reach 99.3%. We believe this performance is competitive and more efficient.
Performance Comparisons with Nlow-Matching-Batt Based Methods. EquiFM and GOAT are all based on flow matching, using an ODE solver for generation. We can see that fow-matching-bent methods can obtain faster generation speeds than diffusion models. In particular, GOAT only needs 90 steps, while EquiFM 200 steps for sampling. EquiFM solely considers optimal transport for atom coordinates. Therefore, the generation speed is still inferior to ours. Because the proposed GOAT solves optimal molecule transport and optimal distribution transport together, the number
8

---

Published as a conference paper at ICLR 2025
Table 3: MAE for molecular property prediction. A lower number indicates better controllable generation result. The best results are highlighted in bold.
GOAT
EquiPM GeoBFN
MAE (a) [quality]
3.5
EDM
Geol.DM
Property
Steps
a
Ar
THOMO
FLUMO
Co
(HUL, 2.74)
GOAT
EDM (1000,2.66)
Units
Bohr"
meV
meV
meV
sdl
2.5
QM9
0.100
64
39
36
0.043
0.040
(200, 2:410)
Random
9.010
1470
645
1457
16666
6.857
(1000 2.770
N atoms
3,860
806
426
813
1.053
1.971
GUAT (288, 1.72)
EDM
1000
2.760
655
356
583
1111
1.101
Ge
1000
2.370
567
340
522
1.108
1,025
200
400
600
800
1000
EquiFM
220
2.410
591
337
530
1.106
1,033
Sampling Steps [speed]
GOAT (Ours)
200
1.725
585
330
521
0.906
0.881
Figure 4: Quality vs. Speed (ex).
GooBFN
100
3,875
768
426
855
1.381
1,401
EquiFM
100
3,006
830
392
735
1,064
1.177
GOAT shows the optimal trade-off be-
GOAT (Ours)
2.740
605
350
534
1,010
0.883
tween generation quality and speed.
of sampling steps is further reduced by 2x compared to EquiFM with the same ODE solver. This verifies our hypothesis that joint optimal transport path can further boost the generation efficiency.
Though EquiFM can perform well in terms of molecule validity, it achieves unsatisfactory perfor- mance in novelty and significance on QM9 among all methods. More specifically, nearly half of the generated samples are the same as the training samples, which unacceptable in the context of de novo molecule design. In contrast, GOAT can obtain 78.6% novelty with 37% improvement and 72.3% significance with 34.6% improvement compared to EquiFM. On GEOM-DRUG, the pro- posed method achieves approximate performance compared to EquiFM while taking only half the sampling steps. GeoBEN (Song et al., 2023b) can have comparable sampling efficiency to ours, which is neither diffusion-based nor flow-matching based methods. We find that its generation qual- ity over GEOM-DRUG is around 3% lower than GOAT regarding the validity, and it owns around 8% decrease in novelty with similar sampling speed.
Geometric Transport Cost Comparisons. As EquiFM and GOAT are both flow-matching-bant transport methods, we compare their transport costs and present the visualized results in Figure 3. We present distribution transport cost (Po + P1) in blue bars and molecule transport cost av- eraged over the number of atoms (go > g1) in red lines. Compared to EquiFM transports with hybrid method, the proposed method reduced the geometric transport cost with 1) unified trans- port (Unified), 2) optimal molecule transport (1-OMT), and 3) optimal distribution transport cost (1-ODT), thereby achieving significant reduction in geometric transport cost by nearly 89.65% leading to faster generation. We further minimize molecule and distribution transport costs (2-OMT and 2-ODT) and observe that the transport cost is reduced marginally, indicating nearly optimum of the proposed method. The above analysis reveals that the proposed method indeed reduced the geometric transport cost by unifying transport, minimizing molecule transport cost, and estimating optimal couplings. The most intuitive manifestation of the reduction in transport cost is the boost in generation speed, which has been demonstrated in the previous section.
Controllable Molecule Generation. Without loss of generality, GOAT can be readily adapted to perform controllable molecule generation with a desired property s by modeling the neural network as vy (z, t|s). We evaluated the performance of GOAT on generating molecules with properties including a, Ae. EHомO, ELUMO. μ. and Cy. The quality of the generated molecules concerning their desired property was assessed using the Mean Absolute Error (MAE) between the conditioned property and the predicted property. This measure helps to determine how closely the generated molecules align with the desired property.
We use the property classifier network & from (Garcia Satorras et al., 2021) and split the QM9 training partition into two halves with 50K samples each. The classifier şp is trained in the first half, while the Conditional GOAT trained in the second half. Then, &p is applied to evaluate conditionally generated samples by the GOAT. We report the numerical results in Table 3. Random means we simply do random shuffling of the property labels in the dataset and then evaluate şp on it. Natons predicts the molecular properties by only using the number of atoms in the molecule.
Compared to existing methods, our proposed approach demonstrates superior performance in con- trollable generation tasks. Specifically, it achieves the best results across all six tasks when eval- uated with 100 sampling steps, outperforming other variable sampling step methods (EquiFM and GeoBFN). Furthermore, even with increased sampling steps, our method maintains outstanding per-

---

Published as a conference paper at ICLR 2025
Table 4: Ablation Studies. OMT represents optimal molecule transport, and ODT stands for opti- mal distribution transport. The best results are highlighted in bold, and the second-bess results are highlighted with underlines.
Metrics (QM9)
Quality (1)
Efficiency (U)
Components
A
Atom Sta
Valid
Uniqueness
Novelty
Significance
Steps
Time
Cost
w/o FAE
98.8
92.8
92.2
58.4
53.9
280
0.63
30.25
w/o ODT
97.7
89.5
98.7
70.1
61.9
120
0.17
5.01
w/o ODT 0.75
97.9
89.7
98.9
70.2
62.3
110
0.16
4,63
W/o ODT 0.5
98.1
89.9
98.8
70.4
62.5
100
0.14
4.13
w/o ODT 0.25
97.8
89.6
98.7
70.1
62.0
120
0.17
4.87
w/o ODT
0
97.5
89.3
98.8
70.0
61.8
130
0.19
5.41
w/o OMT 0.5
96.5
85,0
98.9
69.1
58.1
170
0.31
5.32
GOAT
0.5
99.2
92.9
99.0
78.6
72.3
90.0
0.12
3.14
formance in generating molecules with properties a, μ. and Co. whereas other methods require longer sampling steps to achieve comparable results. Notably, only GeoBFN, which requires more than doubled sampling steps, shows marginal advantage in other properties. To better illustrate the advantages of the proposed method, we present a performance comparison in accuracy and effi ciency, measured by property a, in Figure 4. The figure demonstrates that the proposed method achieves a new trade-off between accuracy and efficiency in conditional molecule generation.
Parameter Analysis and Ablation Studies. In this section, we first analyze the effectiveness of the parameter A, which determines the trade-off between transporting atom coordinates and atom features. Since the parameter only affects the optimization of molecule transport, we compare its influence without considering optimal distribution transport. Our results in Table 4 indicate that when 0.5, meaning the weights for optimizing transport costs in two modalities are equal, the proposed algorithm achieves the best transport plan. We hypothesize that this is because the transport cost calculated by equation with generative model in transporting noise to the data distribution.
0.5 accurately reflects the actual cost for the
Additionally, we conduct ablation studies on equivariant autoencoder (EAE), optimal molecule transport, and optimal distribution transport. Without considering optimal transport, the model trained solely with flow matching in the latent space (w/o OMT) shows a significant increase in training speed. This can be attributed to the reduced transport cost resulting from the unified space, although the performance remains suboptimal. When solving OMT without ODT (w/o ODT, X = 0.5), both performance and speed improve, but they still do not reach the final results, which account for both molecule and distribution in geometric optimal transport.
L Limitations. Addressing the optimal transport costs, particularly those involving rotation and per- mutation aspects, can be computationally intensive (Song et al., 2023a; Klein et al., 2023). However, these operations can be efficiently parallelized on CPUs enhance the training speed. Besides, re- fining the flow may require additional time-consuming training, but such an operation boosts the generation speed and improves novelty without compromising quality. In summary, the above- mentioned operations will accelerate the generation of molecules once and for all after training, which prioritized in this research. We leave improvements concerning training efficiency and other methods for boosting generation speed, such as distillation (Liu et al., 2022), for future work.
5
CONCLUSION
This paper introduces GOAT, a 3D molecular generation framework that tackles optimal transport for enhanced generation quality and efficiency molecule design. Recognizing that in silico molecule generation is a problem of probability distribution transport and the key to accelerating this lies in minimizing the transport cost. To this end, we formulated the geometric optimal transport problem tailored for molecular distribution. This proposed problem led us to consider the transport cost of atom coordinates, atom features, and the complete molecules. This motivates the design of joint transport to solve optimal molecule transport with different modalities and the framework to mini- mize the distributional transport cost. Both theoretical and empirical validations confirm that GOAT reduces the geometric transport cost, resulting in faster and more effective molecule generation. Our method achieves sfate-of-the-ar performance in generating valid, unique, and novel molecules, thereby enhancing the ultimate significance of in-silico molecule generation.
10

---

Published as a conference paper at ICLR 2025
ACKNOWLEDGMENTS
This work was supported in part by the Research Grants Council the Hong Kong (HK) SAR under Grant No. C5052-23G, Grant PolyU 15229824, Grant PolyU 15218622, Grant PolyU 15215623 and Grant PolyU 15208222; the National Natural Science Foundation of China (NSFC) under Grants U21A20512; NSFC Young Scientist Fund under Grant PolyU A0040473.
REFERENCES
Luigi Ambrosio, Elia Brué, Daniele Semola, et al. Lectures on Optimal Transport, volume 130. Springer, 2021.
Brandon Anderson, Truong Son Hy, and Risi Kondor. Cormorant: Covariant Molecular Neural
Networks. Advances in Neural Information Processing Systems, 32, 2019.
Simon Axelrod and Rafael Bombaz--ombareli. GEOM, Energy-Annotated Molecular Conforma- tions for Property Prediction and Molecular Generation. Scientific Data, 9()):185, 2022. ISSN 2052-4463.
Enrico Celeghini, Riccardo Giachetti, Emanuele Sorace, and Marco Tarlini. The Three-Dimensional Euclidean Quantum Group E (3)q and Its R-Matrix. Journal of Mathematical Physics, 32(5): 1159-1165, 1991.
Ricky TQ Chen, Yulia Rubanova, Jesse Bettencourt, and David K Duvenaud. Neural Ordinary Differential Equations. Advances in Neural Information Processing Systems, 31, 2018
Quan Dao. Hao Phung, Binh Nguyen, and Anh Tran. Flow Matching in Latent Space. arXiv preprint arXiv:2307.08698, 2023.
John R Dormand and Peter J Prince. A Family of Embedded Runge-Kutta Formulae. Journal of computational and applied mathematics, 6(1):19-26, 1980.
Victor Gareia Satorras, Emiel Hoogeboom, Fabian Fuchs, Ingmar Posner, and Max Welling. E(n) Equivariant Normalizing Flows. In M. Ranzato, A. Beygelzimer, Y. Dauphin, P.S. Liang, and J. Wortman Vaughan (eds.), Advances in Neural Information Processing Systems, volume 34, pp. 1181-4192. Curran Associates, Inc., 2021.
Niklas Gebauer, Michael Gastegger, and Kristof Schütt. Symmetry-Adapter Generation of 3D Point Sets for the Targeted Discovery of Molecules. Advances in Neural Information Processing Sys- tems, 32, 2019.
Xu Han, Caihua Shan, Yifei Shen, Can Xu, Han Yang, Xiang Li, and Dongsheng Li. Training- free Multi-objective Diffusion Model for 3D Molecule Generation. In The Twelfth International Conference on Learning Representations, 2023.
Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising Diffusion Probabilistic Models. Advances in Neural Information Processing Systems, 33:6840 -8511, 2020.
Emiel Hoogeboom, Victor Garcia Satorras, Clément Vignac, and Max Welling. Equivariant Dif- fusion for Molecule Generation in 3D. In International Conference Machine Learning, pp. 8867-8887, PMLR, 2022.
Wengong Jin, Regina Barzilay, and Tommi Jaakkola. Junction Tree Variational Autoencoder for Molecular Graph Generation. In International Conference Machine Learning, pp. 2323-2332. PMLR, 2018
Hojung Jung, Youngrok Park, Laura Schmid, Jachyeong Jo, Dongkyu Lee, Bongsang Kim, Se- Young Yun, and Jinwoo Shin. Conditional Synthesis of 3D Molecules with Time Correction Sampler. Advances in Neural Information Processing Systems, 37:75914-7599 2024.
Wolfgang Kabsch. A Solution for the Best Rotation to Relate Two Sets of Vectors. Acta Crystal- lographica Section A: Crystal Physics. Diffraction, Theoretical and General Crystallography, 32 (5):922-923, 1976.
11

---

Published as a conference paper at ICLR 2025
Diederik Kingma and Jimmy Ba. Adam: A Method for Stochastic Optimization. In International Conference on Learning Representations (ICLR), San Diega, CA, USA, 2015.
Leon Klein, Andreas Krämer, and Frank Noé. Equivariant Flow Matching, Advances in Neural Information Processing Systems, 36, 2023.
Harold W Kuhn. The Hungarian Method for The Assignment Problem. Naval Research Logisties Quarterly, 2(1-2):83-97, 1955.
Greg Landrum et al. Rdkit: Open-Soure Cheminformaties Software. Open-source Cheminformat- ics, 2016.
Yaron Lipman, Ricky TQ Chen, Heli Ben-Hamu, Maximilian Nickel, and Matthew Le. Flow Match- ing for Generative Modeling. The Eleventh International Conference on Learning Representa- tions, 2022.
Qi Liu, Miltiadis Allamanis, Marc Brockschmidt, and Alexander Gaunt. Constrained Graph Varia- tional Autoencoders for Molecule Design. Advances in Neural Information Processing Systems, 31, 2018
Xingchao Liu, Chengyue Gong, al. Flow Straight and Fast: Learning to Generate and Transfer Data with Rectified Flow. In The Eleventh International Conference on Learning Representations, 2022.
Youzhi Luo and Shuiwang Ji. An Autoregressive Flow Model for 3D Molecular Geometry Genera- tion From Scratch. In International Conference on Learning Representutions (ICLR), 2022.
Laurence Midgley, Vincent Stimper, Javier Antorán, Emile Mathieu, Bernhard Schölkopf, and José Miguel Hernández-Lobato SE (3) Equivariant Augmented Coupling Flows. Advances in Neural Information Processing Systems, 36, 2023.
Stefano Peluchetti. Diffusion Bridge Mixture Transports, Schrödinger Bridge Problems and Gener- ative Modeling. Journal of Machine Learning Research, 24(374):1-51 2023.
Xingang Peng, Jiaqi Guan, Qiang Liu, and Jianzhu Ma. MolDiff: Addressing the Atom-Bond Incon- sistency Problem in 3D Molecule Diffusion Generation. In Proceedings of the 40th International Conference Machine Learning, pp. 27611-27629, 2023.
Raghunathan Ramakrishnan, Pavlo O Dral, Matthias Rupp, and O Anatole Von Lilienfeld. Quantum Chemistry Structures and Properties of 134 Kilo Molecules. Scientific data, 1(1):1-7,201
Danilo Rezende and Shakir Mohamed. Variational Inference with Normalizing Flows. In Interna- lional Conference on Machine Learning, pp. 1530-1538. PMLR, 2015.
Lars Ruddigkeit, Ruud Van Deursen, Lorenz C Blum, and Jean-Louis Reymond. Enumeration of 166 Billion Organic Small Molecules in the Chemical Universe Database GDB-17. Journal of chemical information and modeling, 52(1 1)26644-28755 2012.
Victor Garcia Satorras, Emiel Hoogeboom, and Max Welling. E (n) Equivariant Graph Neural Networks. International Conference on Machine Learning, pp. 9323-9332. PMLR, 2021.
Jean-Pierre Serre et al. Linear Representations of Finite Groups, volume 42. Springer, 1977.
Chence Shi, Minkai Xu, Zhaocheng Zhu, Weinan Zhang, Ming Zhang, and Jian Tang. GraphAF: a Flow-based Autoregressive Model for Molecular Graph Generation. In International Conference on Learning Representations, 2019.
Yang Song, Jascha Sohl-Dickstein, Diederik P Kingma, Abhishek Kumar, Stefano Ermon, and Ben Poole. Score-Based Generative Modeling through Stochastic Differential Equations. Interna- tional Conference on Learning Representations, 2020.
Yuxuan Song, Jingjing Gong, Minkai Xu, Ziyao Cao. Yanyan Lan, Stefano Ermon, Hao Zhou, and Wei- Ying Ma. Equivariant Flow Matching with Hybrid Probability Transport for 3D Molecule Generation. Advances in Neural Information Processing Systems, 36, 2023a.
12

---

Published as a conference paper at ICLR 2025
Yuxuan Song, Jingjing Gong, Hao Zhou, Mingyue Zheng, Jingjing Liu, and Wei- Ying Ma. Unified Generative Modeling of 3D Molecules with Bayesian Flow Networks. In The Twelfth Interna- tional Conference on Learning Representations, 2023b.
Alexander Tong, Nikolay Malkin, Guillaume Huguet, Yanlei Zhang, Jarrid Rector-Brooks, Kilian FATRAS, Guy Wolf, and Yoshua Bengio. Improving and Generalizing Flow-Based Generative Models with Minibatch Optimal Transport. ICML Workshop on New Frontiers Learning, Control, and Dynamical Systems, 2023.
W Patrick Walters and Mark Murcko. Assessing the Impact of Generative Al on Medicinal Chem- istry. Nature Biotechnology, 38(2):143-145 2020.
Joseph L. Watson, David Juergens, Nathaniel R. Bennett, Brian L. Trippe. Jason Yim, Helen E. Eisenach, Woody Ahern, Andrew J. Borst, Robert J. Ragotte, Lukas F. Milles, Basile I. M. Wicky, Nikita Hanikel, Samuel J. Pellock, Alexis Courbet, William Sheffler, Jue Wang, Preetham Venkatesh, Isaac Sappington, Susana Vázquez Torres, Anna Lauko, Valentin De Bortoli, Emile Mathieu, Sergey Ovchinnikov, Regina Barzilay. Tommi S. Jaakkola, Frank DiMaio, Minkyung Back, and David Baker. De Novo Design of Protein Structure and Function with RFdiffusion. Nature, 620(7976): 1089-11000000 2023. ISSN 1476-4687.
Lemeng Wu, Chengyue Gong, Xingehao Liu, Mao Ye, and Qiang Liu. Diffusion-Basee Molecule Generation with Informative Prior Bridges. Advances in Neural Information Processing Systems, 35:36533-36548 2022.
Tian Xie, Xiang Fu, Octavian- Ganea, Regina Barzilay, and Tommi S. Jaakkola. Crystal Diffusion Variational Autoencoder for Periodic Material Generation. In International Conference on Learning Representations, 2022.
Minkai Xu, Alexander Powers, Ron O Dror, Stefano Ermon, and Jure Leskovee. Geometric Latent Diffusion Models for 3D Molecule Generation. In International Conference Machine Learning, pp. 38592-38610. PMLR, 2023.
Zaixi Zhang, Qi Liu, Chee-Kong Lee, Chang- Yu Hsieh, and Enhong Chen. An Equivariant Gen- erative Framework for Molecular 8392, 2023.
Co-Design. Chemical Science, 14(81):8380-
13

---

Published as a conference paper at ICLR 2025
APPENDIX
A EQUIVARIANCE AND INVARIANCE IN GEOMETRIC OPTIMAL TRANSPORT
Equivariance. Molecules, typically existing within threc-dimensional physical space, are sub- ject to geometric symmetries, including translations, rotations, and potential reflections. These are collectively referred to as the Euclidean group in 3 dimensions, denoted as E(3) (Celeghini al., 1991).
function F is said to equivariant to the action of group G Gf o F(x) = Fo o (x) for all G, where Sg. Tg are linear representations related to the group element (Serre et al., 1977). Invariance. A function F is said to be invariant to the action group if π (x) = F(x) for all g E G and every permutation π E Sp.
Equivariance and Invariance in Molecules. For geometric graph generation, we consider the special Euclidean group SE(3), involving translations and rotations. Moreover, the transformations Sg or Tg can be represented by a translation and an orthogonal matrix rotation R.
For molecule g (x. h), the node features h -Euuiveqqu which can be expressed as Rx + (Rx; + +t...., + +t).
while the coordinates
Equivariance and Invariance in Geometric Optimal Transport. For non-topological data, such as images, the transport cost between two given data points is fixed. However, this does not apply to topological graphs. For instance, when topological graph (molecule) undergoes rotation or translation, the inherent properties of the molecule remain unchanged, but the cost of transporting coordinates may vary. Similarly, if the atom order in of the molecules changes in silico, the molecule remains constant, but the transport cost of coordinates and features may alter. Therefore, the proposed optimal molecule transport problem aims to find an optimal rotation, translation, and permutation transformation for one molecule to minimize the distance, considering both coordinates and features, from another molecule.
B PROOF FOR THEOREM 3.1
The theorem 3.1 is reproduced here for convenience:
Theorem 3.1 The coupling f incurs larger geometric transport costs than the arbitrary coupling
I(Po.P. P1) in that E(c, (Zo: z/}] < E(C,(
( Zo. 21)l where (Zo, z1) F(Po, P(), (Zn, Z1) E T(Po. PI).
and ē,p( (Z0 Z1) min 7| [Rz + t, RzZ + t...., R t) (z 2 x,00 ........ z 2 +
min [z) I......... 1, (z),0= zo. ................................ N 122 VT, R, and
z is geometry g in the latent space, which is composed of zx E Nx3 and Zn E IR Nxk, where ls is the latent dimension characterized by Eg.
With graeulga optimal transport R. and á we have:
ElCy(Zo, 21) -E [min || (Rz% - t, Rx + t,..,, R + +t) - (zk 22.0..., ,....,,,))) No
+ min ||=(Σ|1111 2/2 1......... ..... 1/111111 h zh.0.... 2 VT, R, and t]
= =[[ + f, R2
+ + f...., Ruz 'N + +f) (z 22.0 22
[ 2/ 1........ . _
NN
(2 2h.0.... 2 |2]
14

---

Published as a conference paper at ICLR 2025
Let 2x A( + i, R2 t,.... Ray + t), 2n
ft ......,
................................ 2hh) and
[2 zinl
R Nxx+3))) then have:
Biey (zo, 21)
_E]][ 1 +x, .....,, 22, (zx,0- 2x 2x.0..... x0
+ +1 h........
(Z 2 ..... N. 11,2)
=R[ +, 29 +,
(Z); z0-
40 }||2]
= Zo [2]-
Likewise, we have:
E| (0, 21 = B|||2: 2o|12]
At this point, what we aim to prove is simplified to:
E| ス zu [|]]] |||| zo [12]
(8)
Proof. Given that zi ( (zo). dz: =1g (z,, t)dt, we have:
Elcy (0, 21)] = E 1/// rg (z,. Ddt
(9)
2
|| ll2 (3+(3) + R. is the Euclidean norm of and it convex, therefore, with Il Sa odtll s Sa [|ll| dt induced by Jensen's inequality have:
Elcg(20 21 s E V. (Iva(ee, dt
(10)
With defined (z((2.1) 1) = 18|z1 20/20 po] we then have:
[ || - 20 2el
(20 x = E
(11)
Again, with the finite form of Jensen's inequality, we have:
E E (zo. 21) < E
| zoll2 |z
// Jensen's 's inequality
TE [[E[[lz] 21 zull2 |2:]]ct
(12)
[B[[| zallz|dt
// [R[12 [2 20112|2e] 1|2:
zoll2
E[[[2; 20ll2]
Eiey (zo. z1)l
// By Eq.
Combining equations 12, Eq. proved.
It is important note that solving the geometric optimal transport problem in the latent space does not necessarily ensure that the molecule itself or its distribution also satisfies the optimal transport in the original space. However, given that the proposed flow model is trained in the latent space, is sufficient to ensure that latent molecules and distributions are transported with optimal cost, thereby accelerating the flow model in the generation of molecules.
C ALGORITHMS
This section contains the main algorithms of the proposed GOAT. First, we present the algorithm for solving optimal molecule transport and unified flow in Algorithm and Algorithm 2, respectively. Algorithm presents the pseudo-code for training the GOAT. Algorithm 4 presents the process of fast molecule generation with GOAT.
15

---

Published as a conference paper at ICLR 2025
Algorithm Optimal Molecule Transport
Input: 21 and Zo-
2: Output: 21 and %o.
3: Optimal Molecule Transport:
4: Mcyli:). < || 2/ll 2 | - xx,01 2 |
- 22,ll // Construct Atom-level Trunsport
Cost Matrix
5: f + Hungarian algorithm (Kuhn, 1955)
I/ Optimal Permutation
6: R < Kabsch algorithm (Kabsch, 1976)
// Optimal Rotation
7: 21 = x(Rzz)
I/ Optimal Molecule Transport
8: return 21, Zo
Algorithm 2 Equivariant Autoencoder
1: Input: geometric data point g (x, h), equivariant encoder &
2: Output: encoded data point z
3: Unified Flow:
4: x+ x
G(x)
I/ Translate to CoM Space // Encode
5: Me ph t Ex h)
6: (cx, Gh) ~ I)
I/ Sample noise for Equivariant Autoencoder Translate CoM Space // Obtain Latent Representation
7: 6x f 6x G((x)
8: 2x 2h t + (xx, Ch/ O
9: 2 + (xx bnl
10: return
Algorithm 3 Geometric Optimal Transport
1: Input: data distribution pt. equivariant encoder Ex, decoder De, flow network
2: Output: GOAT: (fg)
3: for gi (x, h) P1 do
4: Z1 Equivariant Autoencoder(g:)
I/ Algorithm2 1/ Sample noise from base distribution po Algorithm Loss for the flow I/ Optimize
5:
Zo (zx,0. Zh.0) ~ I)
6:
21. Zg Optimal Molecule Transport (z), 20)
7:
LFi (0) pe. / [(u,,,,,)) 1)
(2)-2))11
8:
i < optimizer (F..0)
9: end for
10: for g1 = (x, h) ~ P1 do
11:
Zo. Zo.zí.gi +Sampling(D...)
11 Algorithm 4
if gi meets quality (measure by RdKit (Landrum al., 2016)) then
13:
21. = Optimal Molecule Transport (z/, 20)
11 Algorithm I Loss for the flow // Optimize
(0) Et,po.pi (110)2), t) - (2) zo)l|
15:
i + optimizer(C2 (Cp. 0)
end if
17: end for return
Algorithm Sampling
Input: equivariant decoder D., flow network 0.
2: Output: noise: Zp- generated latent sample: z generated molecule: gí.
3: Zo f (Zx,0, Zh.o) ~ (O
Sample noise from base distribution Po
4: 2 ODE,g (20)
5: gi e (z')
// Solve ODE
6: return Zo, 2
16

---

Published as a conference paper at ICLR 2025
D RELATED WORKS
Molecule Generation Models. Initial research in molecule generation primarily concentrated on the creation of molecules as 2D graphs (Jin et al., 2018; Liu et al., 2018; Shi et al., 2019). How- ever, the field has seen shift in interest towards 3D molecule generation. Techniques such as G-SchNet (Gebauer et al., 2019) and G-SphereNet (Luo & Ji, 2022) employ autoregressive methods to incrementally construct molecules by progressively linking atoms molecular fragments. These approaches necessitate either a detailed formulation of complex action space or an ordering of actions.
Motivated by the success of Diffusion Models (DMs) in image generation, the focus has now turned to their application in 3D molecule generation from noise (Hoogeboom et al., 2022: Xu al., 2023; Wu et al., 2022; Han et al., 2023). To address the inconsistency of unified Gaussian diffusion across various modalities, latent space was introduced by (Xu et al., 2023). To resolve the atom-bond in- consistency issue, (Peng et al., 2023) proposed different noise schedulers for different modalities to accommodate noise sensitivity. However, diffusion-basee models consistently face the challenge of slow sampling speed, resulting in a significant computational burden for generation. To enhance the speed, recent proposals have introduced flow matching-base -ased (Song et al., 2023a) and Bayesian flow network-bass based (Song et al., 2023b) models. Despite these advancements, there remains substantial potential for improvement in these frameworks regarding speed, novelty, and ultimate significance.
Flow Models. Introduced in (Chen al., 2018), Continuous Normalizing Flows (CNFs) represent a continuous-time variant Normalizing Flows (Rezende Mohamed, 2015). Subsequently, flow matching (Lipman et al., 2022) and rectified flow (Liu et al., 2022) were proposed to circumvent the need for ODE simulations during forward and backward propagation in CNF, and they introduced optimal transport for faster generation. Leveraging these advanced flow models, (Garcia Satorras et al., 2021) pioneered the use of flow models for molecule generation, which was later followed by the proposal of (Song al., 2023a), based on hybrid transport. Beyond the realm of 3D molecule generation, the concept of flow matching and optimal transport has also found applications in many- body systems (Garcia Satorras et al., 2021) and molecule simulations (Midgley et al., 2023). Despite these ad addancements, existing models primarily focus on atomic coordinates. leaving the challenge of geometric optimal transport unresolved.
E DATASET
E. QM9 DATASET
QM9 (Ramakrishnan et al., 2014) is a comprehensive dataset that provides geometric, energetic, electronic, and thermodynamic properties for subset of the GDB-17 database (Ruddigkeit et al., 2012) comprises a total of 130,831 molecules We utilize the train/validation/ess partitions de- lineated in (Anderson et al., 2019), comprising 100K, 18K, and 13K samples for each respective partition.
E.2 GEOM-DRUG DATASET
HEOM-DRUCE (Geometric Ensemble Of Molecules) dataset (Axelrod & Boanez-Bommarrlli, 2022) encompasses around 450.000 molecules, each with an average of 44.2 atoms and a maximum of 181 atoms We build the GEOM-DRUG dataset following (Hoogeboom et al., 2022) with the provided code.
F IMPLEMENTATION DETAILS
In this study, all the neural networks utilized for the encoder, flow network, and decoder are imple- mented using EGNNs (Satorras et al., 2021). The dimension of latent invariant features, denoted as k, is set to 2 for QM9 and for GEOM-DRUG, to map the molecule for unified flow matching.
31933998
17

---

Published as a conference paper at ICLR 2025
For the training of the flow neural network, we employ EGNNs with 9 layers and 256 hidden features on QM9, and layers and 256 hidden features on GEOM-DRUG, with a batch size of 64 and 16, respectively.
In the case equivariant autoencoders, the decoder is parameterized the same manner as the encoder, but the encoder is implemented with 1-layer EGNN. This shallow encoder effectively constrains the encoding capacity and aids in regularizing the latent space (Xu et al., 2023).
All models utilize SiLU activations and are trained until convergence. Across all experiments, the Adam optimizer (Kingma & Ba, 2015) with a constant learning rate of 10 is chosen as our default training configuration. The training process for QM9 takes approximately 3000 epochs, while for GEOM-DRUG, it takes about 20 epochs.
With the flow model trained QM9 or GEOM-DRUG, we then generate and purify the coupling to obtain a total of 100K molecular pairs, which form the estimated couplings.
Hardware Configuration
1. GPU: NVIDIA GeForce RTX 3090
2. CPU: Intel(R) Xeon(R) Platinum 8338C CPU
3. Memory: 512 GB
4. Time: Around 7 days for QM9 and 20 days for GEOM-DRUG.
G MORE EXPERIMENTAL RESULTS
We present the full results in Tables and our detailed experimental results on QM9, we reproduced EDM, GeoLDM, and EquiFM the QM9 dataset to obtain the actual generation time consumption with the same compute configuration. As a result, the proposed method achieves the fastest sampling speed, which is consistent with the measurement of sampling steps. We also witness a huge generation speed improvement by the proposed GOAT for GEOM-DRUG.
In addition supplementing the actual time used for generation, molecule stability, and it obvious that all methods achieve nearly 0% molecule stability in GEOM- DRUG. This is because metries, atom and molecule stability, create errors during bond type predic- tion based on pair-wise atom types and distances. Therefore, we concentrate on metrics measured by RdKit.
also added the metrics of
Lastly, we produced the full results of GeoBFN using sampling steps from 50 to 1,000. It is worth noting that the novelty and significance continue decrease on QM9 datasets sampling steps increase, which aligns with our conjecture in the experiments. Besides, we also observed that its performance on CEOM-DRUG also decreased in terms of validity. Combined with its efficiency and quality, we believe that our method, GOAT, has competitive performance compared with GeoBFN.
We present the visualization of generated molecules on QM9 and GEOM-DRUG in Figures 5 and 6.
18

---

Published as a conference paper at ICLR 2025
Table 5: Comparisons of generation quality (larger is better) in terms of Atom Stability, Molecule Stability, Validity, Uniqueness, Novelty, and Significance. And comparisons of generation efficiency regarding generation time and sampling steps for one molecule (less better). The best results are highlighted in bold.
Q.99
Efficiency
Quality (%) Uniqueness
# Metrics
S-mem
Steps
Atom Sta
Mol
Valid
Novelty
Significance
Data ENF
99.0
95.2
97.7
100.0
85.0
4.9
40.2
98.0
G-Schnet
95.7
68.1
85.5
93.9
GDM-aug
1.50
1000
97.6
71.6
9U.4
99.0
66.8
73.9
EDM
1.68
1000
98.7
82.0
91.9
65.7
64.8
EDM-Bridge
1000
98.8
84.6
92.0
98.6
Geol.DM
1.86
1000
98.9
89.4
93.8
98.8
58.1
53.9
50
98.3
85.1
92.3
98.3
72.9
66.1
0.16
100
98.6
87.2
93.0
98.4
70.3
64.4
GeoBEN
500
98.8
88.4
93.4
98.3
67.7
62.1
1000
99.1
90.9
95.3
97.6
66.4
61.8
EquiFM
0.37
200
98.9
88.3
94.7
98.7
57.4
53.7
GOAT
0.12
90
98.4
84.1
90.0
99.0
78.6
72.3
Table 6: Comparisons of generation quality (larger is better) in terms of Atom Stability, Molecule Stability, Validity, Uniqueness, Novelty, and Significance. And comparisons generation efficiency regarding generation time and sampling steps per molecule (less is better). The best results are highlighted in bold.
GEOM-DRUG
Efficiency S-Time
Quality (%
Metrics
Steps
Atom Sta
Mol Sta
Valid
Uniqueness
Data
86.5
0.0
99.9
100.0
ENF
G-Schnet
GDM-aug
1000
77.7
91.8
EDM
14.88
1000
81.3
0.0
92.6
99.9
EDM-Bridge
1000
82.4
92.8
50
78.9
93.1
100
81.4
93.5
GeoBEN
500
85.6
92.1
1000
86.2
91.7
GeoL.DM
12.84
1000
84.4
0.0
99.3
99.9
EquiFM
200
84.1
98.9
GOAT
0.94
90
84.8
0.0
96.2
99.9
19

---

Published as a conference paper at ICLR 2025
Figure 5: Molecules Generated by GOAT trained on QM9.
Figure Molecules Generated by GOAT trained on GEOM-DRUG.
20

---

Published as a conference paper at ICLR 2025
H DISTANCE BETWEEN NOISES AND GENERATED MOLECULES
We presented a comparison of the distance between generated molecules and the initial noise, in- cluding compared methods, the proposed GOAT, and its variants (w/o ODT and w/o OMT). We presented the experimental results in the Table:
Table 7: Distance Between Noises and Generated Molecules
QM9
GEOM-DRUG
Metries
Average distance
Average distance per atom
Average distance
Average distance
per atom
EDM
651.52
22.47
1834.22
10.13
GeoLDM
185.01
6.38
1046.67
5.78
EquiFM
530.86
18.31
1543.23
8.57
GOAT- w/o ODT
72.48
2.50
924.50
5.14
GOAT-w/o OMT
93.36
3.22
1190.88
6.62
GOAT
55.10
1.90
702.89
3.88
The experimental results on the distance between molecules and noises validate that the proposed method achieves the minimum transport distance from the noise and thereby also verifies the supe- riority of our method in generation speed.
I IMPACT STATEMENTS
This paper contributes to the advancement of generative Artificial Intelligence (AT) in scientific do- mains, including material science, chemistry, and biology. The insights gained will significantly enhance generative AI technologies, thereby streamlining the process of scientific knowledge dis- covery.
The application of machine learning to molecule generation expands the possibilities for molecule design beyond therapeutic purposes, potentially leading the creation of illicit drugs or hazardous substances. This potential for misuse and unforeseen consequences underscores the need for strin- gent ethical guidelines, robust regulation, and responsible use of these technologies safeguard individuals and society.
21


Reviews:
[{"review_id": "J1jQeTjgTb", "reviewer": "Reviewer_DEWZ", "summary": "Because of the known multi-modal (Cartesian coordinates and other features) distribution issue in molecule generation and distribution coupling issue, the paper proposed an optimal molecule transport (OMT) algorithm based on optimal transport (OT) in Flow Matching (FM) on latent codes encoded by equivariance autoencoder.", "questions": "Based on the listed weakness points above, I want to ask following questions:\n\n1. Why using equivariance autoencoder to produce latent codes can yield in unified permutation? \n2. Can I assume your paper is based on the already on-the-shelf work which proposed EOT and EAE? If so, what is your major novel improvement scientifically based on that? If not, please clarify why.\n3. Please clarify the terminology issues on point 2 and point 4 mentioned above in the weakness part.\n\nI will adjust my rating to the paper based on the answers to the above questions during rebuttal and discussion phase.", "limitations": "", "rating": 6, "confidence": 3, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "1. **Clear problem statement**: optimal transport of multi-modal probability and distribution coupling in the specific case of molecule generation are well-known problems and to be addressed. The writing provides enough information for general readers to comprehend. \nThe experiments seem abundant to make the effectiveness of the method stand. \n2. Table 4 is crucial for the core statement of this paper, which is that Optimal Molecule Transport (OMT) is important for faster and better molecule generation based on optimal transport (OT) of Flow Matching (FM). From the numbers there, it seems justified.", "weaknesses": "1. I do not see why after equivariance autoencoder mapping to the latent space, we thus have a unified optimal permutation. In equation (5), we still have the same \\pi for both coordinates and features right? The reasoning for this part is missing.\n2. I am not sure the distribution coupling, $\\Gamma$ is a correct terminology used in this paper. The distribution coupling should be referring to pairing points between two distributions, rather than pairing within one pair of data sample by permutation, rotation and translation. \n3. Maybe limited novelty: In EquiFM[1], we know the idea of equivariant optimal transport (EOT) is proposed, which has a large overlapping with the core contributions of this paper. Can this paper’s main contributions be concluded as  EOT + Equivariant Autoencoder? If so, the novelty of paper may be limited.\n4. Is Geometric Probability Distribution a formal terminology in molecule generation community? I do not see it is widely adopted in other related work and it can be misunderstood with geometric distribution. \n5. The abbreviation EAE in line 251 pops out without specification (I am assuming it is referring to Equivaraint Autoencoder?)\n\n[1] Song, Yuxuan, et al. \"Equivariant flow matching with hybrid probability transport for 3d molecule generation.\" Advances in Neural Information", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "T8iTroWv2Y", "reviewer": "Reviewer_AHDV", "summary": "This paper proposes a 3D molecule generation framework which formulates a geometric transport formula for\nmeasuring 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. They further propose a flow refinement and purification mechanism for optimal coupling identification, which filters out the subpar molecules to ensure the ultimate generation quality.", "questions": "Please refer to weaknesses.\n\nI honestly do not know much about the literature on molecule generation so my confidence would not be high. I am happy to check other reviewers' opinions.", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "1. The problem setup is well-motivated. The OT path and equivariance properties are indeed needed in diffusion models/flow matching methods for fast sampling process. \n\n2. The empirical performance is strong in terms of various metrics. The proposed methods is good in both generation quality and inference steps.", "weaknesses": "1. In Eq. (5), the authors seem to sample permutation, rotation, and translation matrix. Does this mean that the model is not strictly equivariant but approximately equivariant? In other words, the equivariance is learned from data augmentation but not learned by construction.\n\n2. I am wondering maybe it is possible to distill the trained model, like rectified flow, to have even faster inference (1 to 5 steps). I am especially curious to see the generation quality of distillation + purification.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "x84Egx6ln6", "reviewer": "Reviewer_cmjL", "summary": "This paper introduces a novel framework for accelerating 3D molecule generation, named GOAT, which leverages the principles of flow-matching optimal transport to efficiently generate molecules with improved quality and speed.", "questions": "1. It would be interesting if the paper could provide a comparison of the distance between generated molecules and the initial state, with/without the use of optimal transport.\n2. The use of optimal transport may raise concerns about the diversity of the generated molecules. This could be addressed by generating N molecules and comparing the coverage of the reference conformation set within a specified distance threshold (you can refer to the design of this criterion in papers like GeoDiff[1]).\n[1] Xu, Minkai, et al. \"Geodiff: A geometric diffusion model for molecular conformation generation\"", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. I think that using optimal transport to optimize the training process of diffusion models is a very clever and reasonable design, especially when applied to scenarios like molecular generation, where the generation efficiency of diffusion models can be a concern.\n2. I'm glad that the paper points out the optimal transport optimization might bring extra computational cost, especially those involving rotation and permutation, can be computationally intensive. However , it's inevitable, but it's a great trial.\n3. The framework achieves the satisfiable downstream generation quality regarding validity, uniqueness, and novelty, which are crucial metrics in molecule generation.", "weaknesses": "1. While the paper includes some ablation studies, a more comprehensive set of experiments could provide further insights into the contribution of each component of the framework.\n2. The training process itself may be time-consuming, which could be a drawback for some applications, though as mentioned in strengths, it's inevitable.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "cpcy7VmjR6", "reviewer": "Reviewer_N7NU", "summary": "The paper presents a novel approach to molecule generation by introducing a variation of optimal transport for multimodal features within a general flow-matching objective. An equivariant neural network is utilized to transform these multimodal features into a latent space, where multimodal data optimal transport is applied. The results show that this method outperforms existing models like EquiFM and EDM, and it also demonstrates superior computational speed compared to other models.", "questions": "1. Could you provide a more comprehensive evaluation of the generated 3D molecular structures?\nI acknowledge that determining the most effective 3D metric is an open and important question. There are several possible solutions you might consider:\n- The approach uses the Wasserstein distance between the distributions of bond lengths and bond angles. Some papers, like MiDi, have implemented this method. However, be cautious, as a single poorly predicted molecule can significantly skew the metric due to the way MiDi implemented it.\n- SemlaFlow, for example, proposes using the optimization energy drop with the MMFF94 force field as a measure. Given that your ground truth data is based on GFN-xTB calculations, it might be more appropriate to assess the energy drop using GFN-xTB, as it aligns with the potential energy landscape of your dataset.\n- Models like JODO and others have utilized MMD for bond angles, bond lengths, and torsion distributions, focusing only on the most frequent bonds and angles. While none of these metrics are perfect, they are generally more informative and insightful than relying solely on 3D atom stability.\n\nThe important part here is to ensure that when comparing models, you use exactly the same implementation of these metrics to obtain a reliable comparison. By incorporating these additional evaluations, you can provide a more thorough and insightful assessment of the geometric accuracy of your generated molecular structures. If I am adding one of these, I would rather do it on GEOM-Drugs dataset because it is much more realistic, for QM9 it is quite fast to compute GFN-xTB geometry optimization, and a lot of these molecules do not have considerable variability of 3D structure. \n\n2. Could you perform comparisons with recent models like EQGAT-Diff, JODO, or SemlaFlow? I'm especially interested in a comparison with SemlaFlow because both papers place significant emphasis on improving the performance of molecule generation and utilize the flow matching objective. The main difference is that your approach performs flow matching in latent space, while SemlaFlow conducts diffusion explicitly on bonds, atoms, and coordinates.\n\nIt would be greatly appreciated if you could compare your method with SemlaFlow in terms of:\n- Topological Metrics: Such as 2D molecule stability, validity, or introduced significance.\n- 3D Metrics: Possibly using some of the metrics described in previous comments (e.g., bond angles, torsion distributions).\n- Performance Metrics: Considering that both models generate molecules of similar size, metrics like the average time to generate a molecule would be insightful.\n\nAdditionally, could you discuss the advantages of your method over the version of flow matching used in SemlaFlow? Providing a brief discussion of the key algorithmic differences between your proposed method and SemlaFlow's flow matching approach—focusing on how these differences might impact performance or computational efficiency—would enhance the understanding of your contributions.", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. **Significant Theoretical Development in Joint Optimal Transport:**\n\nThe theoretical advancement of joint optimal transport presented in the paper has the potential to greatly impact future flow-matching (FM)-based model development. This innovation could propel the entire field of molecule and conformer generation toward faster and more efficient methods.\n\n2. **Emphasis on Performance Improvement and Comprehensive Comparisons:**\n\nThe paper places a strong focus on performance enhancement and provides thorough performance comparisons with models like EquiFM, EDM, and others. Since dataset conformations are obtained through computational methods, there is an upper bound on computational cost beyond which applying AI is not justified. The proposed model addresses this by improving performance while staying within acceptable computational limits.", "weaknesses": "1. **Comparison with Recent Edge-Modeling Methods and Potential Integration of Edge Features:**\n\nSome recent models, such as JODO and EQGAT-Diff, have achieved better performance by explicitly modeling edges in molecular graphs. Including comparisons with these models would strengthen your paper. Moreover, adding edge features to your joint optimal transport framework seems feasible and could unlock the full potential of your method. Additionally, a more recent model called Semla Flow (a preprint published in June 2024) demonstrates superior accuracy and speed. Comparing your approach with these methods would greatly enhance the value of the proposed latent optimal transport (OT) technique.\n\n2. **Concerns About Atom Stability Metric and 3D Evaluation:**\n\nWhile the task is 3D molecule generation, there is no real evaluation of 3D coordinates for unconditional molecule generation in your paper. The term \"atom stability\" is mentioned without a clear definition. Given that you report an 86.5% figure for the GEOM Drugs dataset, I assume this refers to 3D atom stability. This metric is based on comparing bond lengths with tabulated values, allowing a tolerance of about 0.5 Å. The issue is that only 2.8% of GEOM Drugs molecules fully comply with these criteria, making the metric potentially misleading. I strongly encourage avoiding the propagation of this metric in new papers. The optimal distances between atoms are primarily defined by the potential energy landscape underlying the data—for GEOM Drugs, it's GFN2-xTB—and depending on atom configurations, deviations in bond lengths can exceed 10%.\n\n3. **Difficulty in Claiming State-of-the-Art Performance:**\n\nRelated to the first point, it is difficult to claim that this method achieves state-of-the-art (SOTA) performance because it has not been compared with some other relevant methods mentioned above, including MiDi. Including such comparisons would strengthen the claim of achieving SOTA performance.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)