Equivariant Flow Matching with Hybrid Probability Transport for 3D Molecule Generation
Authors: Yuxuan Song
Jingjing Gong
Minkai Xu
Ziyao Cao
Yanyan Lan
Stefano Ermon
Hao Zhou
Wei-Ying Ma
Venue: NeurIPS 2023 poster (2023)
Keywords: Drug Design
Molecule Generation
Deep Learning
Computational Biology

Abstract: The generation of 3D molecules requires simultaneously deciding the categorical features (atom types) and continuous features (atom coordinates). Deep generative models, especially Diffusion Models (DMs), have demonstrated effectiveness in generating feature-rich geometries. However, existing DMs typically suffer from unstable probability dynamics with inefficient sampling speed. In this paper, we introduce geometric flow matching, which enjoys the advantages of both equivariant modeling and stabilized probability dynamics. 
More specifically, we propose a hybrid probability path where the coordinates probability path is regularized by an equivariant optimal transport, and the information between different modalities is aligned. Experimentally, the proposed method could consistently achieve better performance on multiple molecule generation benchmarks with 4.75$\times$ speed up of sampling on average.

Equivariant Flow Matching with Hybrid Probability Transport for 3D Molecule Generation
Yuxuan Song JJingjing Gong Minkai Xu², Ziyao Cao Yanyan Lan 1,3, Stefano Ermon", Hao Zhou', WWei-Wiinn Ying Ma
Institute of Al Industry Research(AIR) Tsinghua University Computer Science Department, Stanford University
s "Bejjing Academy of Artificial Intelligence (songyuxuan, gongj inggingg www.gnua.nnnnnnneae. edu.
Abstract
The generation of 3D molecules requires simultaneously deciding the categorical features (atom types) and continuous features (atom coordinates). Deep generative models, especially Diffusion Models (DMs), have demonstrated effectiveness in generating feature-rich geometries. However, existing DMs typically suffer from unstable probability dynamics with inefficient sampling speed. this paper, we introduce geometric flow matching, which enjoys the advantages of both equivariant modeling and stabilized probability dynamics. More specifically, we propose a hybrid probability path where the coordinates probability path is regularized by equivariant optimal transport, and the information between different modalities is aligned. Experimentally. the proposed method could consistently achieve better performance on multiple molecule generation benchmarks with 4.75x x speed up sampling average
Introduction
Geometric generative models aim approximate the distribution of complex geometries and emerge as an important research direction in various scientific domains. general formulation of the geometries in scientific fields could be the point clouds where each point is embedded in the Cartesian coordinates and labeled with rich features. For example, the molecules are the atomic graphs in 3D 4 and the proteins could be seen the proximity spatial graphs (15] Therefore, with the ability of density estimation and generating novel geometries, geometric generative models have appealing potentials in many important scientific discovery problems, e.g., material science [36]. de
novo drug design [I] and protein engincering [45]
With the advancements of deep generative modeling, there has been series of fruitful research progresses achieved in geometric generative modeling, especially molecular structures. For example, 0 and 42 proposed data-driven methods generate 3D molecules (in silico) with autoregressive and flow-based models respectively. However, despite great potential, the performance is indeed limited considering several important empirical evaluation metrics such validity, stability, and molecule size, due to the insufficient capacity of the underlying generative models [40]. Most recently, diffusion models (DMs) have shown surprising results on many generative modeling tasks which generate new samples by simulating a stochastic differential equation (SDE) to transform the prior density to the data distribution. With the simple regression training objective, several attempts [ on applying DMs in this field have also demonstrated superior performance. However, existing
* Eqqual Contribution, Hao Zhou @whouhhaa@@iiitsiigguuu.co 2The code is available at htpp://gthub.commlggoooleeMala
the corresponding author.
37th Conference on Neural Information Processing Systems (NcurIPS 2023).

---

N N
N
N
N
C
C
Sample
Noise
X
Figure Illustration of EquiFM. We define a hybrid path for generating molecules g (x, h), where is trained on equívariant optimal transport path and h is trained a path whose information quantity aligned with x's path. The sampling is conducted by solving an ODE, i.e. go ODESolveeg.. 0).
DM-based methods typically suffer from unstable probability dynamics which could lead to an inefficient sampling speed also limit the validity rate of generated molecules.
In this work, we propose a novel and principled flow-matching objective, termed Equivariant Flow- Matching (EquiFM), for geometric generative modeling. Our method is inspired by the recent advancement of flow matching [26]. a simulation strated appealing generation performance with stable training and efficient sampling. Nevertheless, designing suitable geometric flow-mate wow-matching objectives for molecular generation is non-tne trivilli
objective for training CNFs that has demon-
(1) the 3D skeleton modeling is sensitive, i.., a slight difference in the atom coordinates could affect the formulation some certain types of bonds; (2) the atomic feature space consists of various physical quantities which lies in the different data modality, e.g., charge, atom types, and coordinates are correspondingly discrete, integer, and continuous variables. To this end, we highlight our key innovations follows:
For stabling the 3D skeleton modeling, we introduce an Equivariant Optimal-Tranm Transport guide the generative probability path of atom coordinates. The improved objective implies an intuitive and well-motivated prior, i.e. minimizing the coordinates changes during generation, and helps both stabilize training and boost the generation performance.
Towards the modality inconsistency issues, proposed to differ the generative probability path of different components based on the information quantity and thus introduce hybrid generative path. The hybrid-path techniques distinguish different modalities without adding extra modeling complexity or computational load.
The proposed model lies the scope of continuous normalizing flow, which is parameterized by ODE. We can use an efficient ODE solver during the molecule generation process to improve the inference efficiency upon the SDE simulation required DMs.
A unique advantage of EquiFM lies in the framework enriching the flexibility to choose different probability paths for different modalities. Besides, the framework is very general and could be easily extended to various downstream tasks. We conduct detailed evaluations of the EquiFM multiple benchmarks, including both unconditional and property- property-conditionee conditioned molecule generation. Results demonstrate that EquiFM consistently achieve superior generation performance all the metrics, and 4.75x speed up on average. Empirical studies also show a significant improvement in controllable generation. All the empirical results demonstrate that the EquiFM enjoys significantly higher modeling capacity and inference efficiency.
2 Related Work
Flow Matching and Diffusion Models Diffusion models have been studied various research works
such as 47 12 49]. and have recently shown success in fields like high-dimenn digenssinni statistics [1]
language modeling (23). and equivariant representations [33]. Loss-rescaling techniques for diffusion models have been introduced in [48), while enhancements to the architecture incorporating classifier guidance are discussed in [8]. Noise schedule learning techniques have also been proposed in [3 18]
Diffusion models suffer from unstable probability dynamics and inefficient sampling, which limits

---

their effectiveness in some scenarios. Flow matching is a relatively new approach that has gained attention recently. Research works such as [26] 28 have proposed this simulation-fre free objective for training continuous normalizing flow. It involves other probability paths besides the diffusion path and could potentially offer better sampling efficiency through ODE solving. Furthermore, the follow-ups [37 (50] proposed to use the OT couplings to straighten the marginal probability paths. However, the application of flow matching to geometric domains requires designing appropriate probability paths, which is an area that remains unexplored.
3D Molecule Generation Previous studies have primarily focused on generating molecules as 2D graphs [1] 27
46). but there has been increasing interest in 3D molecule generation. G-Schnet and G-SphereNet 9 30 have utilized autoregressive methods to construct molecules by sequentially attaching atoms or molecular fragments. These frameworks have also been extended to structure- based drug design [4
However, this approach requires careful formulation of a complex action space and action ordering. Other approaches use atomic density grids that generate the entire molecule in a single step by producing a density over the voxelized 3D space [31]. Nevertheless, these density grids lack the desirable equivariance property and require separate fitting algorithm.
35 38
In the past year, the attention has shifted towards using DMs for 3D molecule generation I 52 53]. with successful applications in target drug generation [25], antibody design [29]. and protein design 2 5T. However, our method based on the flow matching objective and hence lies in a
different model family, i.e. continuous normalizing flow, which fundamentally differs from this line of research in both training and generation.
3 Backgrounds
3.1 Flow Matching for Non-geometric Domains
In this section, we provide an overview of the general flow matching method to introduce the necessary notations and concepts based on (26). The data distribution is defined as g, zo represents a data point from q and Ii represents a sample from the prior distribution p1- The time-dependent probability path is defined as Pre[o,1] : F - Rx and the time-dependent vector field is defined as
Vtejo,1] : Rd + Rd, The vector field uniquely defines time-dependent flow theeoo,1 Rd - R by the following ordinary differential equation (ODE):
d
dt th (2) = (1hx v (x)
(1)
4 proposed to train the parameterized flow model v/, called continuous normalizing flow (CNF)
with black-box ODE solvers. Such a model could reshape a simple prior distribution p1 to the complex real-world distribution g. CNFs are difficult to train due to the need for numerical ODE simulations. (26) introduced flow matching, simulation-free objective, by regressing the neural network (2 (r, f) to some target vector field (z)::
(F) (x())())) [10)(z, t) (z))
The objective LFM requires access to the vector field un (z) and the corresponding probability path p.(x). However, these entities are difficult to define in practice. Conversely, the conditional vector field u,(xx zn) and the corresponding conditional probability path (x 2o) are readily definable.
The probability path can be marginalized from mixture of conditional probability path (x = J [x =o)q(=o)drog and the vector field u, (z) can be marginalized from conditional vector field as (x) = Bzonn (x(x)) This illustrates how (x( and (x are related to their conditional form, and (26 further proved that with the conditional vector field (z (2) generating
p=(x)
the conditional probability path p=(x: zo). the marginal vector field ux(x) will generate the marginal
probability path Pe (xx). The observation inspires the new conditional flow matching (CFM) objective:
(0) Br.g((e)) [[w t) [x(x 20)]12
The CFM objective enjoys the tractability for optimization, and optimizing the CFM objective is equivalent optimizing Eq. 2 i.e., voL (6) VoLCrm(O). For the inference phase, ODE solvers could be applied to solve the Eq. [] e.g., 2o = ODESolve 0g 1, In this paper, we
consider using the Gaussian conditional probability path, which lies in the form of pz 20)
N (æ pe(2o), (x)o)) We introduce two probability paths utilized in the following of paper:
3

---

Conditional Optimal Transport Path With the prior distribution defined as a standard Gaussian
distribution, the empirical data distribution Po (x #n) is approximated with a peaked Gaussian
centered in 20 with small variance min as N (ar Fo. min-). The probability path is pt(x [o))
N (2) (1 1)aro (amin + (1 miin)()²1) and the corresponding flow is (/)(1) = (amin + (1
miin )t)z + (1 t) Then the vector field could be obtained by Eq. as: (e((((r)) xo)
di (x)) -r (1 min )r.
d
Put the above terms into 3 the reparameterized objective is as:
(0) (z) ([=)((((2)),)) (ete - (-zo + -mimmm)))
(4)
Intuitively, the conditional optimal transport objective tends to learn the transformation direction from noise to data sample in straight line which could hold appealing geometric properties.
Variance Preserving Path The pariance-preservee (VP) path is defined as Pe | mo)
N (x Cg (1 az 2)) where a (+()) and T(t) S" 3(s)ds. Here 3 is some noise schedule function. Following the Theorem. 3 in 266] the target conditional vector field of VP path could be derived ux(x (o) I-a7
of (are: mp). of denotes the derivative with respect time.
And the objective for VP conditional flow matching is as:
af (a.r zu) 1-at aj
M (0) (.,,(
vo (2,
The VP path flexible to control the information dynamics, e.g. correlation changes towards zo on the conditional probability path, by selecting different noise schedule functions.
4 Methodology
In this section, we formally describe the Equivariant Flow Matching (EquiFM) framework. The proposed method is inspired by the appealing properties of recent advancements in flow matching (26). but designing suitable probability paths and objectives for the molecular generation is however challenging [13.. We address the challenges by specifying a hybrid probability path with equivariant flow matching. The overall framework is introduced in Section. 4.1 And then elaborate on the design details of the hybrid probability path of the equivariant variable and invariant variable in
Section. 4.2 and Section. respectively. A high-level schematic is provided in Figure. I
4.1 Equivariant Flow Matching
X is the atom coordinates matrix and h = (h'...., h² E R N xd is the node feature matrix, such as atomic type and charges. Here X E RR : > x² = =0} is the Zero Center-of-Mass (Zero CoM) space, which means the average of the N elements should be O. We introduce the general form of equivariant flow matching in the following.
Recall molecule could be presented as the tuple g (x, h), where x = (x'....,x') xN)
SE(3) Invariant Probability Path For modeling the density function in the geometric domains, is important to make the likelihood function invariant to the rotation and translation transformations. We could always make the probability path of equivariant variable x invariant to the translation by setting N EN u(z t)' = O. Formally,
the prior distribution and vector field in the Zero CoM space, i.e. 1 the rotational invariance could be satisfied by making the parameterized vector field equivariant and the prior invariant to the rotational transformations as shown in the following statement:
Theorem 4.1. Let (v% (g. t), up (g. t)) = ug t), where vz (g. t) and vh (g. t) are the parameterized vector field for x and h. If the vector field is equivariant any rotational transformation R, i.e., vo (R., h).() = (Rxx (g. ()), (gg()) With an rotational invariant prior function pı (x, h), i.e., pi (Rx, h) pı (x, h), then the probability path pa,e generated by the vector field (o(()) is also rotational invariant.
To make the vector field satisfy the equivariance constraint, we parameterize with an Equivariant Graph Neural Network (EGNN) [43]. And more details can be found in Appendix C

---

(a) EquiFM
(b) FM
(c) EDM
Figure 2: Generation route visualization of different models. Note that a lighter color indicates an earlier step of an atom and denser color corresponds to later step. A change of base color indicates change of atom type. EquiFM generates molecules in straightforward route shown in 2 Vanilla flow matching method 2(b) on the other hand, takes detour while generating molecules, resulting route inward then outward before
converging to a molecule. The generation process in EDM2(c) is rather chaotic until the last few steps before
converging into a molecule.
Hybrid Probability Modeling We refer to the target conditional vector field on each part as u [g go) and (g go) correspondingly, then could get the objective in the following formulation:
(0) (r,g lo) u go) 2
1| (g. t) uu ( 10)
(6)
Proposition 4.2. There could be joint probability path pr(glgo) which satisfies that pt (glgg)))
(x( p (h) (oo))). and the conditional vector field on x and h is independent: u7 go) = ( xo), (g go) = u.hh] ho).
The above Proposition 4.2 states special property of conditional flow matching, i.e., in the multi- variable setting the probability path of each variable could be designed independently. Such property appealing in setting, x and h hold different data types and come from different manifolds, thus is intuitive to use different probability paths for modeling and generating the two variables.
Coordinates Matching with Equivariant Optimal Transport
We focus the generation of coordinates variable The conditional OT path (Eq. 4 could be a promising candidate as it tends to move the atom coordinates directly towards the ground truth atom coordinates along straight line. However, directly applying the objective could be problematic in 3D molecule generation. With Xy as the point cloud from molecule distribution and x₃ from the prior distribution, the objective in Eq. tends to move the atom based between the atoms. Optimizing the vector field toward such direction could involve extra variance for training and lead twisted and unstable generation procedure shown in molecule generation visualization Fig. 2(b)and and Fig.
random alignment
[(c)
To address the above-mentionel issue, we first introduce the concept of equivariant optimal transport (EOT) between two geometries as follows:
Definition 4.3. Given point clouds, (z) ...., 2 x3 and y (y) y Nx3
We define the equivariant optimal transport plan as
R argmin||I(RRz: RR?
RR (y²
>
R.R
Here a permutation N elements and E 3x3 stands for a rotation matrix in the 3D space. With and lie in the zero of mass space, the mappings in Eq. 7 are optimal towards any E(3) equivariant operations on either side of point clouds. Therefore, the mappings are referred as equivariant optimal trans)
The equivariant optimal transport finds the minimum straight-line distance between the paired atom coordinates upon all the possible rotations and alignment. We could then build a probability path based on the EOT map which could minimize the movement distance of atom coordinates for the transformation between the molecule data from and sampled point cloud from P1 as:
Pt [o J<P,, where (ETT(()
(x) (amin + (1 mt (R'x x) (1 t)x0
(8)
5

---

Proposition 4,4. The probability path implied by the EOT map, i.e. Eq. 8 is also an SE(3) invariant probability path.
The proposition could be proved following the Definition 4.3 and the Theorem 4.1 Combining the
above terms, the final equivariant optimal transport based training objective is:
BOT CPM
|oo (vo(v (w: (o(BoT(x)) (x1). (-xu + (1 mmin) v [x))))]
(xx)) (x1)
A good property of the objective with EOT is that the training characteristics are invariant to translation and rotation of initial X1, and equivariant with respect to both sampled noise x, and data point Xo.
which empirically contributes to more effective training.
Solving Equivariant Optimal Transport We propose an iterative algorithm to obtain the equiv- ariant optimal transport map. The algorithm first conducts the Hungarian algorithm (0) to align the atoms between the initial geometry from P1 and the ground truth geometry from Po: and then conducts the Kabsch algorithm [17] to solve the optimal rotation matrix based on the atom alignment. The proposed algorithm asymptotically converges to the optimal solution. Besides, the method holds a close relationship with the Iterative Closest Point (ICP) [5] algorithm, while our settings require the node alignment could only be one-one mapping. We leave the detailed description in Appendix C
4.3 Information Aligned Hybrid Probability Path
In this section, we address the challenges posed by the multi-modality nature of 3D molecular data. Specifically, we focus on the distinct generation procedures required for various modalities, such coordinates and atom types, within the flow-matching framework. It is crucial to recognize that altering atom types carries a different amount of chemical information compared to perturbing coordinates. To better understand this intuition, we provide the following corner case:
Example 1: pr = (x Vt < 6x and (x xx)= pı (x(xx)) Vt 2 Ex-
We define the pe [h (bo) similarly with a different parameter Ep Here we consider the corner case that 6x 0 and <h 1, i.e. no noise for atom types from timestep 0 timestep c and max noise level from <h to timestep 1. (Reversely for (x ) Under such a probability path, the model will be encouraged to determine and fix the node type at around en step (very early step in the whole generation procedure), even if the coordinates are far from reasonable 3D structures. However, this particular case may not be optimal. The subsequent steps of updating the structure could alter the bonded connections between atoms, leading to a potential mismatch in the valency of the atoms with the early fixed atom types. Therefore, selecting a suitable inductive bias for determining the probability paths of different modalities is crucial for generating valid 3D molecules. In this paper, we utilize an information- -thooreticc inspired quantity as the measurement to identify probability paths for learning the flow matching model on 3D molecules.
Definition 4.5. For distribution po on the joint space g. and two corresponding conditional probability path pe (xl go) and pe (b)g)), we denote the [xx,, he) as the mutual information for x; with distribution Pt (x Po (g ddg and h, with distribution f Pa (h (B))Po(gs (go) dgoo
Proposition 4.6. For the independent conditional probability path prlgigo) = (x )P: ho). when the conditional probability path of x and h lies in OT path or VP path, if I [(xo. ho) > O. then
Vt E (0. 1), [xx h.) > O and I [(xt, ht,) > I(Xe,, he,), ). Vt: < tz.
We use the quantity I, [x,,, he) as the key property to distinguish different probability paths. Given the conditional probability path pe (x xn), it implies an information quantity change trajectory from T(x, h,) = 0 to 7 7(x0, ho) following T(x,, ho). Thus, one well-motivated probability path on is to align the information quantity changes by setting / (h, ho) = [xx,, ho). Based on such intuition, we design our probability path on h. We follow the data representation in [13). This, for atom types, we represent it by one-hot encoding and charges are represented as integer variables. Empirically, the VP path involves a noise schedule function 3 which could naturally adjust the information change by choosing different noise schedules, so we explore the probability path mainly on the VP path. For [(, ho), we decompose it as [(,, ho) = H (ha) - H (ho (h)) where H (ho)) is constant and (ho h,) the conditional entropy towards ho with h, as the logits. Similarly, the difficulty of estimation 7(x,, ha) lies in 7(ha) xt). Following the difference of entropy estimator in (32,, we build prediction model (b)oo)))) to estimate [(hp) x,) for selected time I. More details can be found in Appendix C We demonstrate 20 time steps for 1 1(x,, ho), and I (h, ho) for vanilla VP path

---

Table 1: Results of atom stability, molecule stability, validity, and validity x xuniquensss. A higher number indicates a better generation quality. The results marked with an asterisk were obtained from own tests.
QM
DRUG
Metries
Atom Sta (%)
Mol Sta (%)
Valid (%)
Valid & Unique (%)
Atom Sta (%)
Valid (%)
Data
99.0
95.2
97.7
97.7
86.5
99.9
ENF
85.0
4,9
40.2
39.4
G-Schnet
95.7
68.1
85.5
80.3
GDM
97.0
63.2
75.0
90.8
GDM-AUG
97.6
71.6
90.4
89.5
77.7
91.8
EDM
98.7
82.0
91.9
90.7
81.3
92.6
EDM- EDM-Bridge
98.8
84.6
92.0*
90.7
82.4
92.8*
EQUIFM
98.9 + 01.1
88.3 + 0.3
94.7 - 0.4
93.5 + 0.3
84.1
98.9
with the linear schedule (VPiinear) on 33 VP path with cosine cchedules(VP.oo. schedules(VPprr [13, and the OT path in Fig.
33 and polynomial
We observe that the information quantity of [xx, does not change uniformly, this is, it stays stable at the start and drops dramatically after some threshold. It is in line with the fact that when the coordinates x are away from the original positions to certain extent, the paired distance between the bonded atoms could be out of the bond length range [6] In this case, the point cloud x then loses the intrinsic chemical information. Reversely, the dynamics T(x,, ho) also implies generation procedure where the coordinates x transform first and the atom types h is then determined when are relatively stable.
Experiments
In this section, we justify the advantages of the proposed EquiFM with comprehensive experiments. The experimental setup introduced in Section [5.]] Then we report and analyze the evaluation
results for the unconditional and conditional settings in Section 5.2 and [.3 We provide detailed
ablation studies in Section [5.4tt further gain insight into the effect of different probability paths. We
leave more implementation details in Appendix B.4
5.1
Experiment Setup
Evaluation Task. With the evaluation setting following prior works on 3D molecules generation 9
30
13 52]. we conduct extensive experiments of EquiFM on three comprehensive tasks against several state-of-the-art approaches. Molecular Modeling and Generation assesses the capacity to learn the underlying molecular data distribution and generate chemically valid and structurally diverse molecules. Conditional Molecule Generation focuses testing the ability to generate molecules with desired chemical properties. Following [13), we retrain a conditional version EquiFM on the molecular data with corresponding property labels.
Datasets We choose OM9 dataset [39], which has been widely adopted in previous 3D molecule generation studies 19. 10), for the setting of unconditional and conditional molecule generation. We also test the EquiFM on the GEOM-DRUG (Geometric Ensemble Of Molecules) dataset for generating large molecular geometries. The data configurations directly follow previous work 3
5.2 Molecular Modeling and Generation
Evaluation Metries. The model performance evaluated by measuring the chemical feasibility of generated molecules, indicating whether the model can learn underlying chemical rules from data. Given molecular geometries, the bond types are first predicted (single, double, triple, or none) based on pair- wise atomic distances and atom types [B].
Next, we evaluate the quality our predicted molecular graph by calculating both atom stability and molecule stability metrics. The atom stability metric measures the proportion of atoms that have correct valency, while the molecule stability metric quantifies the percentage of generated molecules in which all atoms stable. Additionally, we report validity and uniqueness metries that indicate the
7

---

Table 2: Mean Absolute Error for molecular property prediction. A lower number indicates better control- lable generation result.
Table 3: Ablation study, EquiFM models trained with different probability path, the effect of EOT is also evaluated.
Property Units
As
EHOMO meV
ELUMO
Medhod
Atom Stable
Mol Stable (%)
Bohr
meV
meV
EquiFMoor+wh.
98,9-0.1
0,3
QM9
0.10
64
39
36
0.043
0.040
Equi SquiFMMr++vn.
98.7=0.1
84.920.4
Random*
9.01
1470
645
1457
1616
6,857
Equi EquilM MMrvvn
98.4+0.1
81.6+0.3
Nasus
3.86
866
426
813
1.053
1.91
Equi PM
98.7100.1
84.7=0.2
EDM
2.76
655
356
584
1.111
1.101
Equil Wihwa
98.7+0.1
83.4+00,5
EQUIFM
2.41
591
337
530
1.106
1,033
EquiFMaor Mootror oz
97.3-0.11
77.120.4
percentage of valid (determined by RDKIT) and unique molecules among all generated compounds. Furthermore, we also explore the sampling efficiency of different methods.
Baselines. The proposed method is compared with several competitive baselines. G-Schret [ is the previous equivariant generative model for molecules, based on autoregressive factorization. Equivariant Normalizing Flows (ENF) 42 is another continuous normalizing flow model while the Equivariant Graph Diffusion Models (EDM) with its non-equivariant variant (GDM) [13] recent progress on diffusion models for molecule generation. Most recently, [22 proposed an improved version of EDM FEDM-Bridge' which improves upon the performance of EDM by incorporating well-designed informative prior bridges. To yield fair comparison, all the model-agnostic configurations are set as the same as described in Sec. 5.1
objective is simulation-base
Results and Analysis. We generate 10, 000 samples from each method to calculate the above metrics, and the results are reported in Table T As shown in the table, EquiFM outperforms competitive baseline methods on all metrics with an obvious margin. In the benchmarked 3D molecule generation task, the objective is to generate atom types and coordinates only. To evaluate stability, the bonds subsequently added using predefined module such as Open Babel following previous works. It is worth noting that this bond-adding process may introduce biases and errors, even when provided with accurate ground truth atom types and coordinates. As result, the atom stability evaluated on ground truth may be less than 100%. Note the molecule stability is approximately the N-th power of the atom stability, N is the atom number in the molecule. Consequently, for large molecules in the DEOM-DRUG dataset, the molecule stability is estimated to be approximately 0%. Furthermore, as DRUG contains many more molecules with diverse compositions, we also observe that unique metric is almost 100% for all methods. Therefore, we omit the molecule stability and unique metrics for the DRUG dataset. Overall, the superior performance demonstrates EquiFM's higher capacity to model the molecular distribution and generate chemically realistic molecular geometries. We provide visualization of randomly generated molecules Appendix [P and the efficiency study in Appendix
5.3 Controllable Molecule Generation
Evaluation Metrics. In this task, we aim to conduct controllable molecule generation with the given desired properties. This can be useful in realistic settings of material and drug design where we are interested in discovering molecules with specific property preferences. We test our conditional version of EquiFM on QM9 with 6 properties: polarizability a, orbital energies HoMM., ELUMO and their gap Ae, Dipole moment μr, and heat capacity C.. For evaluating the model's capacity to conduct property-conditione generation, we follow the 42 to first split the QM9 training set into two halves with 50K samples in each. Then we train property prediction network w on the first half and train conditional models the second half. Afterward, given a range of property values 8, we conditionally draw samples from the generative models and then use to calculate their property values as 8. The Mean Absolute Error (MAE) between s and á is reported to measure whether generated molecules close to their conditioned property. We also test the MAE of directly running on the second half QM9, named QM9 in Table 2 which measures the bias of w. A smaller gap with QM9 numbers performance.
indicates better property-conditionn
Baselines. We incorporate existing EDM as our baseline model. In addition, we follow 13 to also list two baselines agnostic to ground-truth property named Random and Natour- Random means we simply do random shuffling of the property labels in the dataset and then evaluate w on it. This operation removes any relation between molecule and property, which can be viewed as upper bound of MAE metric. Natoms predicts the molecular properties by only using the number of atoms in the molecule. The improvement over Random can verify the method is able to incorporate conditional
8

---

0.9
relative I An) or I( hw)
Molecule Stability
0.8
0.8
0.6
0.7
0.6
ByarM Digei Feyars Tale By altd mMM
0.2
0.5
Rex
0.4
600
1,000
0.6
Number of Function Evaluations (NFE)
Figure 3: Resulting performance in molecule sta- bility w.t.f change of NFEs for each integration al- gorithm applied in the sampling process. EquiFM better than compared EDM model with every integration algorithm.
Figure 4: The x-axis is time, and the y-axis is the normalized mutual information estimation. It could observed that VP path holds the clos- est tendency to that / (Xe, ho). and OT path has the largest discrepancy.
property information into the generated molecules. And overcoming Natoms further indicates the model can incorporate conditioning into molecular structures beyond the number of atoms.
Results and Analysis. The visualizations of conditioned generation be found in Appendix As shown in Table 2 for all the conditional generation tasks, our proposed EquiFM outperforms other comparable models with margin. This further demonstrates the generalization ability of the proposed framework upon different tasks.
5.4 Ablations On the Impacts of Different Probability Paths
In this section, we aim to answer the following questions: 1) how is the impact of the different probability paths the coordinate variable and the categorical variable h? 2) how does the equivariant optimal transport path boost the generation?
To answer these questions, apply several different probability paths and compare them on the QM9 dataset, including the variance-preserviie (VPLiia path(Eq, (5). vanilla optimal transport (OT) pEt and the equivariant optimal 9 transport PEOT(EOT) the coordinate variable x; And variance-preserving (VPLinear) path(Eq. Eq. vanilla optimal transport (OT) (Eq, 4 the variance- preserving path with polynomial decay (VPruly), variance preserving path with cosine schedule (VPc.s). The result is illustrated in Tab. 3 We notice that based paths on coordinates in general show superior performance than the others due to the stability and simplicity of the training objective. Furthermore, regularizing the path with the equivariant- boost the performance by large margin. To gain more intuitive understanding, the generation dynamic comparison in Fig. 2(b) As shown, the generation procedure trained with twisted
prior, the EOT path could further
further provide
vanilla OT path, though more stable than the EDM generation procedure, also exits phenomenon, i.e., all atoms tend to first contract together and then expand; Such phenomenon disappears in the generation procedure of EOT path due that the generation direction well constrained. For the probability path on the categorical variable, we find the VP path, holds the superior performance due the closest alignment with the information quantity changes. If there is significant discrepancy in the information quantity dynamics, e.g., OT path, it may result in substantial decline in performance.
5.5 Sampling Efficiency
We also evaluate the sampling efficiency our model, as shown Fig. the results EquiFM with 4 different integrating algorithms converge to state-of-the-as results in much less NFE compared to baseline model EDM. Remarkably, the red triangle is result of EquiFM with Dopri5 integrating algorithm, it converges in approximately 210 NFE to achieve 0.883 model stability, while EDM takes 1000 NFE to achieve only 0.820. With simple non-adaptive step integration algorithms such as Fulers method and midpoint, the NFE required for convergence is much less than that of baseline models.

---

This indicates that our proposed model has learned a much better vector field, and takes much shorter generation route during generation, this can be justified with visualization Fig. 2(a)
Conclusion and Future Work
We introduce the EquiFM, an innovative molecular geometry generative model that utilizes a simulation-free objective. While flow matching has demonstrated excellent properties in terms of stable training dynamics and efficient sampling in other domains, its application in geometric domains poses significant challenges due the equivariant property and complex data modality. To address these challenges, we propose a hybrid probability path approach in EquiFM. This approach regularizes the probability path coordinates and ensures that the information changes each component of the joint path are appropriately matched. Consequently, EquiFM learns the underly- ing chemical constraints and produces high-quality samples. Through extensive experiments, we demonstrate that the EquiFM not only outperforms existing methods in modeling realistic molecules but also significantly improves sampling speed, achieving speedup of 4.75x compared previous advancements. In future research, as a versatile framework, EquiFM can be extended to various 3D geometric generation applications, such as protein pocket-based generation and antibody design, among others.
Acknowledgments
We thank the Reviewers for the detailed comments and Yanru Qu for the proofreading. This work is supported by National Key R&D) Program of China (2022ZD01 7502), Guoqiang Research Institute General Project, Vanke Special Fund for Public Health and Health Discipline Development, Tsinghua University NN0.20210800553 and Beijing Academy of Artificial Intelligence (BAAT), Tsinghua University (No. 2021 GQ2IGQG1)12). Minkai Xu thanks the generous support of Sequoia Capital Stanford Graduate Fellowship.
References
[1] Michael S Albergo and Eric Vanden-Eijnden. Building normalizing flows with stochastic interpolants. arXiv preprint arXiv:2209. 100.22099.157,,, 2022.
[2] Namrata Anand and Tudor Achim. Protein structure and sequence generation with equivariant denoising diffusion probabilistic models. arXiv preprint urXiv:2205.50019 2022.
[3] Brandon Anderson, Truong Son Hy, and Risi Kondor. Cormorant: Covariant molecular neural networks. Advances neural information processing systems, 32, 2019.
[4] Ricky TQ Chen, Yulia Rubanova, Jesse Bettencourt, and David K Duvenaud. Neural ordinary differential equations. Advances in neural information processing systems, 31, 2018.
[5] Dmitry Chetverikov, Dmitry Svirko, Dmitry Stepanov, and Pavel Krsek. The trimmed iterative closest point algorithm. In 2002 International Conference on Pattern Recognition, volume 3, pages 545-54444 548. IEEE, 2002.
[6] CRC Handbook. CRC Handbook of Chemistry and Physics, 88th Edition. CRC Press, 88 edition, 2007.
[7] David F. Crouse. On implementing 2d rectangular assignment algorithms. IEEE Transactions on Aerospace and Electronic Systems, 52(4):1679-1696 2016.
[8] Prafulla Dhariwal and Alexander Nichol. Diffusion models beat gans on image synthesis. Advances in Neural Information Processing Systems, 34:8780-8794, 2021.
[9] Niklas Gebauer, Michael Gastegger, and Kristof Schütt. Symmetry-adapted generation of 3d point sets for the targeted discovery of molecules. Advances neural information processing systems, 32, 2019.
10

---

[10] Niklas WA Gebauer, Michael Gastegger, Stefaan SP Hessmann, -obur- Müller, and Kristof T Schütt. Inverse design of 3d molecular structures with conditional generative neural networks. arXiv preprint 2 20t:::200.000 (09.04824, 2021.
[11] Jordan Graves, Jacob Byerly, Eduardo Priego, Naren Makkapati, Vince Parish, Brenda Medellin, and Monica Berrondo. A review of deep learning methods for antibodies. Antibodies, 922 2020.
[12] Jonathan Ho. Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistie models. arXiv preprint arXiv:2006. 1239, 2020.
[13] Emiel Hoogeboom, Victor Garcia Satorras, Clément Vignac, and Max Welling. Equivariant diffusion for molecule generation in 3d. In International Conference on Machine Learning, pages 8867-8887. PMLR, 2022.
[14] Wengong Jin, Regina Barzilay, and Tommi Jaakkola. Junction tree variational autoencoder for molecular graph generation. In International conference on machine learning, pages 2323-2332. PMLR, 2018.
[15] Bowen Jing, Stephan Eismann, Patricia Suriana, Raphael John Lamarre Townshend, and Ron Dror. Learning from protein structure with geometric vector perceptrons. In International Conference Learning Representations, 2021.
[16] Jaehyeong Jo, Seul Lee, and Sung Ju Hwang. Score- based generative modeling of graphs via the system of stochastic differential equations. arXiv preprint arXiv: 2022.
[17] 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-925, 1976.
[18] Diederik Kingma, Tim Salimans, Ben Poole, and Jonathan Ho. Variational diffusion models. Advances neural information processing systems, 34:21696-21707, 2021.
[19] Diederik P Kingma and Jimmy Ba. Adam: A method for stochastic optimization. In 3nd International Conference Learning Representations, 2014.
[20] Harold W Kuhn. The hungarian method for the assignment problem. Naval research logistics quarterly, 2(1-2):83-97, 1955.
[21] Landrum. Rdkit: open-source cheminformatics http://www. rdkit. org. 2016.
[22] Jim Lawrence, Javier Bernal, and Christoph Witzgall. A purely algebraic justification of the kabsch-umeyama algorithm. Journal of research of the National Institute of Standards and Technology, 124:1, 2019.
[23] Xiang Lisa Li, John Thickstun, Ishaan Gulrajani, Percy Liang, and Tatsunori Hashimoto. Diffusion-LM improves controllable text generation. In Alice H. Oh, Alekh Agarwal, Danielle Belgrave, and Kyunghyun Cho, editors, Advances in Neural Information Processing Systems, 2022.
[24] Yibo Li, Jianfeng Pei, and Luhua Lai. Structure-based de novo drug design using 3d deep generative models. Chemical science, 12(41):16644-1375 2021.
(25) Haitao Lin, Yufei Huang, Meng Liu, Xuanjing Li, Shuiwang Ji, and Stan LLi. Diffbp: Gener- ative diffusion of 3d molecules for target protein binding, arXiv preprint ar111:2211.11214. 2022.
[26] Yaron Lipman, Ricky TQ Chen, Heli Ben Gam-Hamu, Maximilian Nickel, and Matt Le. Flow matching for generative modeling. arXiv preprint 1rX0v:2210.027777, 2022.
[27] Qi Liu, Miltiadis Allamanis, Mare Brockschmidt, and Alexander Gaunt. Constrained graph variational autoencoders for molecule design. Advances neural information processing systems, 2018
11

---

[28] Qiang Liu. Rectified flow: A marginal preserving approach to optimal transport. arXiv preprint arXiv:2209.14577 2022.
[29] Shitong Luo, Yufeng Su, Xingang Peng, Sheng Wang, Jian Peng, and Jianzhu Ma. Antigen- specific antibody design and optimization with diffusion-basse generative models for protein structures. In Alice H. Oh, Alekh Agarwal, Danielle Belgrave, and Kyunghyun Cho, editors, Advances in Neural Information Processing Systems, 2022.
[30] Youzhi Luo and Shuiwang Ji. An autoregressive flow model for 3d molecular geometry generation from scratch. In International Conference on Learning Representations, 2021.
[31] Tomohide Masuda, Matthew Ragoza, and David Ryan Koes. Generating 3d molecular struc- tures conditional on a receptor binding site with deep generative models. arXiv preprint arXiv:2010. 14442, 2020.
[32] David McAllester and Karl Stratos. Formal limitations the measurement of mutual infor- mation. In International Conference on Artificial Intelligence and Statistics, pages 875-884, PMLR, 2020.
[33] Alexander Quinn Nichol and Prafulla Dhariwal. Improved denoising diffusion probabilistic models. In International Conference on Machine Learning, pages 8162-811 8171. PMLR, 2021.
[34] Adam Paszke, Sam Gross, Soumith Chintala, Gregory Chanan, Edward Yang. Zachary De Vito, Zeming Lin, Alban Desmaison, Luca Antiga, and Adam Lerer. Automatic differentiation in pytorch. In NIPS-W, W, 2017.
[35] Xingang Peng. Shitong Luo, Jiaqi Guan, Qi Xie, Jian Peng, and Jianzhu Ma. Pocket2mol: Efficient molecular sampling based on 3d protein pockets. In International Conference Machine Learning 2022
[36] Janaina Cruz Pereira, Ernesto Raul Caffarena, and Cicero Nogueira Dos Santos. Boosting docking-based virtual screening with deep learning. Journal of chemical information and modeling, 56(12):445-25006 2016.
[37] Aram-Alexandre Pooladian, Heli Ben-Hamu, Carles Domingo-Enrich, Brandon Amos, Yaron Lipman, and Ricky Chen. Multisample flow matching: Straightening flows with minibatch couplings. arXiv preprint arXiv:2304. 144722 2023.
[38] Alexander S. Powers, Helen H. Yu, Patricia Suriana, and Ron O. Dror. Fragment- based ligand generation guided by geometric deep learning on protein-ligand structure. bioRxiv, 2022.
[39] 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, 2014.
[40] Ali Razavi, Aaron Van den Oord, and Oriol Vinyals. Generating diverse high-fidelity images with vq-vae-2. Advances neural information processing systems, 32, 2019.
[41] Robin Rombach, Andreas Blattmann, Dominik Lorenz, Patrick Esser, and Björn Ommer. High- resolution image synthesis with latent diffusion models. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pages 10684-10695 2022.
[42] Victor Garcia Satorras, Emiel Hoogeboom, Fabian B Fuchs, Ingmar Posner, and Max Welling. E (n) equivariant normalizing flows for molecule generation in 3d. arXiv preprint 20X 2021.
[43] Victor Garcia Satorras, Emiel Hoogeboom, and Max Welling. E(n) equivariant graph neural networks. In International conference on machine learning, pages
PMLR, 2021.
[44] Kristof T Schütt, Farhad Arbabzadah, Stefan Chmicla, Klaus R Müller, and Alexandre Tkatchenko. Quantum-chemical insights from deep tensor neural networks. Nature com- munications, 8: 8:130 2017.
[45] Chence Shi, Chuanrui Wang, Jiarui Lu, Bozitao Zhong, and Jian Tang. Protein sequence and structure co-design with equivariant translation. arXiv preprint 2711022210.0871,, 2022.
12

---

[46] Chence Shi, Minkai Xu, Zhaocheng Zhu, Weinan Zhang, Ming Zhang, and Jian Tang. Graphaf: a flow-based autoregressive model for molecular graph generation. arXiv preprint SBKi:22001.033822 2020.
[47] Jascha Sohl-Dickstein, Eric A Weiss, Niru Maheswaranathan, and Surya Ganguli. Deep unsupervised learning using nonequilibrium thermodynamics. arXiv preprint arXiv: 2015.
[48] Yang Song, Conor Durkan, lain Murray, and Stefano Ermon. Maximum likelihood training of score-based diffusion models. Advances in Neural Information Processing Systems, 34:1415- 1428, 2021.
[49] Yang Song and Stefano Ermon. Generative modeling by estimating gradients of the data distribution. Advances in Neural Information Processing Systems, volume 32, pages 1 11918- 1 11930. Curran Associates, Inc., 2019.
[50] Alexander Tong, Nikolay Malkin, Guillaume Huguet, Yanlei Zhang, Jarrid Rector- BrotorBrrooks Kilian Fatras, Guy Wolf, and Yoshua Bengio. Conditional flow matching: Simulation-free dynamic optimal transport. arXiv preprint arXiv:2302.00482. 2023.
[51] Brian Trippe. Jason Yim, Doug Tischer, Tamara Broderick. David Baker, Regina Barzilay, and Tommi Jaakkola. Diffusion probabilistic modeling of protein backbones in 3d for the motif-scaffolding problem. arXiv preprint arXiv:2206.04 19, 2022.
[52] Lemeng Wu, Chengyue Gong, Xingehao Liu, Mao Ye, and qiang liu. Diffusion-based molecule generation with informative prior bridges. In Alice H. Oh, Alekh Agarwal, Danielle Belgrave, and Kyunghyun Cho, editors, Advances in Neural Information Processing Systems, 2022.
[53] Minkai Xu, Alexander Powers, Ron Dror, Stefano Ermon, and Jure Leskovec. Geometric latent diffusion models for 3d molecule generation. arXiv preprint arXiv:2305.01 140, 2023.
[54] Chengxi Zang and Fei Wang. Moflow; an invertible flow model for generating molecular graphs. In Proceedings of the 26th ACM SIGKDD International Conference Knowledge Discovery & Data Mining, pages 617-626, 2020.
13

---

A Sampling and Training Algorithm
We provide a detailed training and sampling pipeline this section. The training algorithm with EOT path the equivariant variable x and VP path on the invariant variable h as example is demonstrated in Algorithm
Algorithm Training Algorithm of EquiFM
1: Input: geometric data distribution Pg. g = (x, h)
2: Initial: vector field network up. minimum constant variance mmin
3: while 0 have not converged do
4: U(0, 1), (O, I). go ~ Pg
5: Subtract center of gravity from Ez in e [E2, En.
6: Obtaining the Equivariant Optimal Transport plan n², RR based on Algorithm 3
7:
(mmin + (1 min (R* + (1 t)zo
[Eq.) 8
8:
h acko + (1 at), a2 = (7()))
[Eq. 5
9: LEqUIFM = 11 ((x,, Nz), t) (-20 + (1--00 o 11 1103 t)
(a,
h
T-of
10:
end while
11:
return Up
The sampling algorithm can be found in Algorithm 2
Algorithm 2 Sampling Algorithm EquiFM
1: Input: vector field model
2: @1 ~ (0, I)
3: Subtract center of gravity from æn in Φ1 [x], hn]
4: go = -DDEoolveegn ug, ug
5: (During the ODE solving, always subtract the center of gravity from vz)
6: sample p( holgo)
7: return (Tg, hu)
Note the p(-lgo) stands for the procedure of transforming the continuous ho into the specific data modality. This is, for the categorical part, p(h go) C(h (ho) and for the integer part p (h (oo
Jn- h+ N (u ho, oo) du.
B Formal Proof of Theorems and Propositions
B.1 Invariant Probability Path: Theorem 4.1
The key properties of the equivariant flow matching model here are the invariant density modeling. For simplicity. here we omit the invariant feature, i.e. h, and focus on the variable x. Here we demonstrated that with invariant prior pi (x) and the equivariant vector field vo (), the marginal distribution implied by the vector field of each time step, (x) is also invariant.
Proof. We are given that Pi (x) is invariant, and that (x(x,t) t) equivariant, i.e. for any ro- tation R. Pi (Rx) = Pi (x), and ve (Rx, t) Rup(x, t). Our target is to prove that Vt e [0, 1], VR, (Rxx = Po. Then specifically, the sampling distribution po (x(x) is invariant.
First recall that the way we generate the distribution (x) is to exert transformation voo. to the prior pı (x). And the definition of the vector field is (x(x, !) = ar (x).
We can derive the equivariance of th, by the following equations:
14

---

{ 0o (x, t)dt
Vi v1(x). v0,((x)
1. Di (Rx, t)dt
3 (Rx) - v)8,((Rx)
f. Reg t)dt
R( (x))
Note that 1/1 (Rx) Ruh (x) since th1 (x) = x, we have voo. (Rx) (x). That to say, Vo, is an equivariant transformation. Thus its inverse Va, is also equivariant, since Vy (x),(x), we have Ry = (x (ex), so dy (Ry) Rx R (y). Also, the
Jacobian matrix (sha.(()) is equivariant, i.e. (v),
l=Rx R Bvs. -ux: which implies that 8u luxx since the det function keeps constant under any rotation.
det (1))
lu=Rx det
(
8u
According to the Change of Variable Theorem,
Tvoo,(u)
(x, = p1 (1 (x))/ det du 1u= (x)
det (v),
pe, (Rx) (1 (Rx))/ (Rx))
Ou uu= (Rx)
Applying the above conclusions together, have
@vo,e(1)
(Rx) P. (R (x))/
det
I (x)
Ou
= =pi(to,( (x))) det @vax(u) Du (x= = (x).
B,2 Explanation of Proposition 4.2
Note for the initial prior distribution, P1 (gl9o). could be the standard distribution, i.e. Pi (glgg) (O,, I). this case, P1 (g (o) = Pi (z(or))Pi (filfi). And for the time step zero, we assume the the gu. i.e., (o(gggu) N T I). And in this case
distribution is Gaussian centralized also have that po(gl go) = (o (ro) (o (h)ph)))
Example 2: For the Gaussian probability path,
Pe (g gi) =N (9 a (o1)²
(10)
where μ J0, Rdd + Rd is the dependent mean of the Gaussian distribution, while [0, 1] x R Rs0 describes time-dependent scalar standard deviation (std).
The Gaussian probability path satisfies that (algg)) Pe (x (ro ) (h//h)) To better clarify, highlight the difference between the conditional probability path and the marginal probability path with the following Remark [
Remark I. With the conditional probability path on x and h being independent each other, the marginal distribution could be correlated, i.e. pr(g) + px())p.hh)
With this property, we could design different paths for modeling complex variables.
15

---

B.3 Proof of Proposition 4.4
In the Proposition 4.4, we claim that the probability path under the EOT map is SE(3) invariant. The invariant property under translations is guaranteed by the training and sampling setting within Zero CoM space, and below we will focus on proving the path is invariant under rotations. The key observation is that R might be different for different
Specifically, for any rotation T acting on the point cloud with N points, the new R corresponding with Tx (we denote it by Rio)) exactly offisetting the impact of TT Formally, let Xy denote the target point cloud, we calculate
(xt R argmin Rx' Rx² Rx - xol12
(r.R)
(Tres R = argmin||
(RRTi (RTx', RTx2 RTxx) X 112
(T.Rt
We claim that not = πT',,, Ris R RT and strict proof follows.
Note that if (R) 33x3 7 R3x3 is reversible map, then for any scalar function f(R). argmin ( = -- (argmin (RR). Here let ((R) RT, and o (R) RT we get
(T R = argmin|| (RTx RTx2
RTx - -xoll2
(r.
[nggin||T(((CC (R)x (R)x - -xoll2
(m,R)
$ (argmin || Rxx Rx² Rx x0l12)
(n,R)
3 ((T"RR") = (TT o (R")) = " RTT
where à is an natural extension of e defined as Ó( R)) (π, of R)), ((T R).
Now we recheck the probability path in Eq. g Since p1 is invariant under rotations, and the transformation EOT satisfies
EOT (x) = (Amin (1 Juin )t) (R*x) + (1 - f)xo 1/
EOT (Tx) = (omin + (1 mmin )t) (t) am R Tx) + (1 f)xo wh
(Amin (1 Gmin) )) n RR x) + (1 -)xo
i.e. v
BOT is invariant. So we conclude that pr = JuP1 is also invariant under rotations.
B.4 Explanation of Proposition 4.6
Here we provide the informal explanation of the Proposition 4.6 The proposition states that for OT path or VP path
t)zo, (amin + (1 mmin)t)²1) or Pa(x mo) = N (x a,zo, (1 af) the mutual information between the marginal variable x, and h monotonically decays following the path from time step to time step 1 where Io (Xo, ho) >> 0. Note that the I (X1, h1) = 0, as p1 (g) = (xx))r(()) (h). Recall the definition of signal-to-nn -o-isee ratio (SNR) as:
or h, i.e. pr(z zo) N (x
122 SNR
For OT-path, SNR "لببب and for VP-path SNR - The key observation is that the SNR along the probability path of both x and on either path will decay monotonically. Intuitively, with t + 1, x; has less information of thus has less of h.
16

---

X0
Xt1
Xt
X1
g
ho
ht
htz
h1
and the direction of arrows indicate the flow corresponding to the target vector field ue running backward (from & = 0 to 1, opposite of the sampling process) random molecule g (not in training data). This can be regarded as process similar to that of a diffusion model, but x and h have conditionally independent paths.
Figure 5:
The dependency relationship the hybrid path. Note that 0 h <
B.5 Proof of Proposition 4.6
Proof. Firstly, we will prove the dependency relationship between the time-dependent variables in the hybrid probability path as shown in Fig.
To this end, demonstrate the relationship between the Xu, hos X Recall the definition of conditional independent probability path as P:(x (Xo, ha)) (x xo). Thus we have the distribution of random variable X Pt, (x)
Pt, (x Xx)p((x)dxx) And hence, there Xx T
bo Xg. Similarly, we could also derive hr, 1 Xn ho.
Next, for any time step 0 < 1 <12 <. We will then demonstrate the dependency relationship between Xo, Xt, Xt.. We denote the target vector field which generates the marginal probability path pa (x) as u. (xx.. then the distribution of random variable Xt, could be then derived Pt, (x)
U,, un (x)df) .P,, (x) I,,,, 2/,, 1-, (x), where ub denotes the transformation corresponding with uy Then we have X 1 Xo X Similarly, h 1 ho h,, could be also derived.
With the above two conclusions, we demonstrate the dependency relationship shown in Fig. 5 holds the hybrid probability path. Therefore, with the dependency relationship we could directly obtain that I(x,, h:,) < I(x,, h,,) If a, in VP path satisfies combination of such OT path and path, we always have [(x,, h:) >
0 when 0 < t < 1, then for any
C Implementation Details
C.1 Solving EOT with variant of iterative closest point (ICP) algorithm.
Problem definition. Given point cloud z 3 and its reference point cloud y R note they are point cloud representations in Euclidean space. The objective is to find an optimal permutation matrix II* E N.N and rotation matrix R R3x3 that minimizes the following objective:
II II"," argmin ((()))
R
(11)
I..
We optimize the objective iteratively with variant of the iterative closest point (ICP) algorithm, where it iteratively obtains II and R.
C.2 Model Architectures and Training Configurations
We use the open-source software RDKIT 21 to preprocess molecules. For QM9 we take atom types (H, C, N, O, F) and integer-valued atom charges as atomic features, while for DRUG we only use atom types.
17

---

Algorithm 3 A variant of iterative closest point (ICP) algorithm.
1: Input: a point cloud z E R3xN and it's reference point cloud y E 3
2: while 7 has not converged do
3: Obtain permutation matrix =I= argmin [[2) T y l|2 with Jonker- Volgenant algorithm
I1
7
4: Obtain rotation matrix R = argmin argmin|[R(II)) [[(zZ)] y 2l2 with Kabsch algorithm 22
5: T [ 1
6: end while
7: return II, R
The vector field network is implemented with EGNNs 14 by PyTorch 34 dimension of latent invariant features atomic feature dimension. For the training of vector field network with layers and 256 hidden features with a batch size 64; and on DRUG, we train EGNNs with layers and 256 hidden features, with batch size 64. The model uses SiLU activations. We train all the modules until convergence, For all the experiments, we choose the Adam optimizer [99 with constant learning rate of 10 our default training configuration. The training QM9 takes approximately 2000 epochs, and on DRUG takes 20 epochs.
package. We set the for QM9 and 2 for DRUG, which extremely reduces the
on QM9, we train EGNNs
D Number of Evaluation (NFE) analysis
We further explore the behavior of adaptive integrators during sampling with Dopri an example. In Fig. 7 we show the average NFE at different time intervals. We could observe that at time intervals 0 the NFE is much larger than at other time intervals. The underlying reason lies in the vector field of h dramatically changes in these steps, which results in the frequent change of the atom type in the last period of sampling. This behavior could be due to the unsmoothness of the categorical manifold which could shed light on several future directions.
E
Scalability
The proposed algorithm has a complexity (n²²) for computing an OT map for single molecule, where is the number of nodes. To understand the computational burden for different molecule sizes, we evaluated the average computing time for OT maps with varying node numbers. The Tab. 4 beloo shows the burden for three datasets/tasks. We also provide the curves of iteration and time needed to solve the EOT map in Fig. 6
iteration node Pumber reletion
lime rodde number ation
0.000
Deeaton Rumber
6.8
8.00%
8.00
0.235
8312
0106
35
3100
202
375
200
Rode
rode rember
(a) iterations towards atom numbers
(b) time towards atom numbers
Figure 6: Time Cost for Solving EOT Maps
Even for the Antibody-CDR data, with an average of 150 atoms, the process time only 18.84 ms, which is acceptable in practice. Additionally, optimize the process further by leveraging its
18

---

Table 4: Scalability Different Molecule Size
Average atom number
EOT mapping time per sample
EOT mapping iteration per sample
QM9
1. 1.10s
4.67
GEOM DRUG
47
1 1.99s
3.89
Antihody-C
150
18.84ms
7.14
NFE Histogram
80
Number Function Evaluations
60
20
0.0
0.2
0.4
0.6
0.8
1.0
t (time)
Figure Number Evaluation analysis of EquiFM generation process with Dopri5 integrator.
parallelizable nature. By Enabling prefetch and multiprocessing, we can minimize the computational overhead further and make virtually inconsequential.
More visualizations
This section presents additional visualizations of molecules generated by our EquiFM method. We include samples from two datasets, QM9 and DRUG, in Fig. [ and Fig. ??, respectively. All examples randomly generated without therry-pickingg but the viewing direction may affeett the visibility of some geometries.
As demonstrated in the figures, our model can generate realistic molecular geometries for small and large molecules alike. However, the model occasionally generates disconnected components, which more common when trained on the large molecule DRUG dataset, shown in the second molecule in Fig. ??. This phenomenon is not unique to our model but is a common issue non-autoregressive molecule generative models [54] [6]. Nevertheless, it is easily solvable by filtering out the smaller components.
We also present a qualitative assessment of controlled molecule generation by EquiFM in Fig, We interpolate the conditioning parameter, polarizability, with different values a, while keeping the prior 9I fixed. Polarizability measures the tendency of matter to acquire an electric dipole moment when subjected to an electric field. In general, less isometric molecular geometries tend correspond to higher a values. This observation is consistent with our results in Fig. g
19

---

Figure 8: Molecules generated from EquiFM trained on QM9,
63.11
66. 76
70.
74.06
77.71
81.36
85.01
88. 66
92.31
95.96
103.27
106.92
110.57
114.22
Figure 9: Molecules generated from conditioned version of EquiFM trained on QM9. We conduct controllable generation with interpolation among different polarizability values with the same prior @1. The given a values are provided at the bottom.
20


Reviews:
[{"review_id": "YrmTSl2w96", "reviewer": "Reviewer_zq9q", "summary": "The paper applies the Flow Matching framework for 3D molecule generation achieving state of the art performance on common benchmarks. The authors introduce several innovations to the FM framework to adjust it for the equivariant data setting with different data modalities within a single sample (e.g., 3D coordinates, atom types, charge).  The two main modifications are (i) applying OT alignment between coordinates to shorten learned paths. (Similar to [1,2]), (ii) using different probability paths for different modalities in the data. \nThe proposed method improves both sampling speed and performance and achieves SOTA in both unconditional and conditional generation.", "questions": "- Regarding the information alignment -  have the authors tried to scale one of the modalities to have the same variance as the other? Can the authors provide some more motivation for using the mutual information?\n\n\n**Missing related work**\nI ask the authors to a add discussion on the following missing related works - \n\n- Two closely related prior/concurrent works that applied OT alignment to improve sampling speed:\n\n  - [[1]](https://arxiv.org/abs/2304.14772) Multisample Flow Matching: Straightening Flows with Minibatch Couplings, Pooladian et. al. (ICML 2023)\n\n  - [[2]](https://arxiv.org/abs/2302.00482) Conditional Flow Matching: Simulation-Free Dynamic Optimal Transport, Tong et. al. (Preprint)\n\n- Another concurrent work pursuing the same applications:\n\n  - [[3]](https://arxiv.org/abs/2305.01140) Geometric Latent Diffusion Models for 3D Molecule Generation, Xu et. al. (ICML 2023)\n\n** Note: These works do not diminish the contribution of this work. It is however necessary to mention and discuss the differences for the completeness of the paper. ", "limitations": "limitations are not discussed.", "rating": 7, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- The paper is well written. The problem setting, the motivation and the proposed method are introduced clearly.\n- The paper showcases an application of FM to equivariant data domain. \n- Achieve SOTA performance on common benchmarks.\n- Experimental section is thorough and presents ablations justifying the algorithmic choices made and demonstrating the benefits in the flexibility of the FM framework.", "weaknesses": "- **Scalability** - in the paper the authors demonstrated the OT alignment on small molecules, how would the method scale to larger sets? How computationally prohibitive it is to compute the OT maps?\n- **$S_n$ invariance** - A proof that the probability path implied by the EOT map is also $S_n$-invariant is missing from proposition 4.4. Meaning that for permutation of the atoms in the molecule the same probability will be returned.\n", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for the insightful comments and the recognization of our work. And we address your concerns in the following:\n### Q1. How would OT maps scales to larger sets and how computationally prohibitive is it? \n\nI first want to clarify that OT maps are not needed during inference. And it is true that OT maps add extra overhead during training. However, this does not limit scalability for larger sets.\n\nThe proposed algorithm has a complexity of $O(n^2)$ for computing OT map for a single molecule, where $n$ is the number of nodes. To understand the computational burden for different molecule sizes, we evaluated the average computing time for OT maps with varying node numbers. The table below shows the burden for three datasets/tasks. **The full curve is in the PDF**.\n\n| - | Average atom number | Average EOT computing time per sample |\n| --- | --- | --- |\n| QM9 | 18 | 1.10ms |\n| GEOM DRUG | 47 | 1.99ms |\n| Antibody-CDR | 150 | 18.84ms |\n\nEven for the Antibody-CDR data, with an average of 150 atoms, the process time is only 18.84 ms, which is acceptable in practice. Additionally, we can optimize the process further by leveraging its parallelizable nature. \nBy Enabling prefetch and multiprocessing, we can minimize the computational overhead further and make it virtually inconsequential.\n\n---\n\n### Q2. The missing proof for the S(n)-invariant probability path of EOT. \n\nThank you for bringing this to our attention. A detailed proof of proposition 4.4 can be found in Appendix B.3 in the submission, and we will include the reference in the updated version. Here we summarize the key steps of the proof and provide a simplified and intuitive explanation:\n\nThe used dynamics model EGNN is permutation-invariant, and Zero CoM guarantees translation invariance. Thus, our goal is to demonstrate the invariance under rotations.\n\nFor any rotation $\\mathbf{T}$ on the point cloud $\\mathbf{x}$ with $N$ points, the new $\\mathbf{R}^*$ corresponding with $\\mathbf{T}\\mathbf{x}$ (we denote it by $\\mathbf{R}^*_\\text{rot}$) exactly offsets the impact of $\\mathbf{T}$. Formally, let $\\mathbf{x}_0$ denotes the target point cloud, we calculate\n\n$(\\pi^*, \\mathbf{R}^*)= \\underset{(\\pi, \\mathbf{R})}{\\operatorname{argmin}} ||\\pi(\\mathbf{R}\\mathbf{x}^1, \\mathbf{R}\\mathbf{x}^2, \\dots, \\mathbf{R}\\mathbf{x}^N) - \\mathbf{x}_0||_2$, and\n\n$(\\pi^*_\\text{rot}, \\mathbf{R}^*_\\text{rot})= \\underset{(\\pi, \\mathbf{R})}{\\operatorname{argmin}} ||\\pi(\\mathbf{R}\\mathbf{T}\\mathbf{x}^1, \\mathbf{R}\\mathbf{T}\\mathbf{x}^2, \\dots, \\mathbf{R}\\mathbf{T}\\mathbf{x}^N) - \\mathbf{x}_0||_2$.\n\nThen $\\pi^*_\\text{rot}=\\pi^*, \\mathbf{R}^*_\\text{rot} = \\mathbf{R}^*\\mathbf{T}^{-1}$. (Refer to Appendix B.3 for strict proof.)\n\nSince $p_1$ is invariant under rotations, and the transformation $\\psi_t^{\\mathrm{EOT}}$ satisfies\n\n$\\psi_t^{\\mathrm{EOT}}(\\mathbf{x}) = (\\sigma_{\\min}+(1-\\sigma_{\\min})t)\\pi^*(\\mathbf{R}^*\\mathbf{x})+(1-t)\\mathbf{x}_0$, and\n\n$\\psi_t^{\\mathrm{EOT}}(\\mathbf{T}\\mathbf{x}) = (\\sigma_{\\min}+(1-\\sigma_{\\min})t)\\pi^*_\\text{rot}(\\mathbf{R}^*_\\text{rot}\\mathbf{T}\\mathbf{x})+(1-t)\\mathbf{x}_0$.\n\ni.e. $\\psi_t^{\\mathrm{EOT}}$ is invariant.\n\n---\n\n### Q3. Can the authors provide some more motivation for using mutual information? Have the authors tried to scale one of the modalities to have the same variance as the other?\n\n1. We explain the motivation for using mutual information. For atom types, the probability path determines when to predict the node type from timestep 0 (noise) to timestep 1 (data). In example 1 of line 214, one corner probability path has no noise from timestep 1 to timestep $1-\\epsilon$ and maximum noise level from $1-\\epsilon$ to timestep 0. When $\\epsilon$ is small, the model tends to fix the node type at around ε step (early in the generation procedure), even if the coordinates are far from reasonable 3D structures. This makes it difficult for the model to learn or generate following such a path. The structure update in the following steps can change the bonded connection, causing a mismatch in valency with the early fixed atom types. Figure 4 shows that as the coordinates approach the ground truth, the predictability of atom types increases significantly. To align with this observation, we use mutual information as a guided bias to describe this tendency.\n \n2. Yes, we have tried different inductive biases for matching the intermediate distributions including scaling different modalities to the same variance as the reviewer suggested. We provide the results for several matching strategies in the following table. \n\n| Metrics | Scaling the variance | Matching entropy in intermediate distribution | Matching Mutual Information(ours) | Without any matching |\n| ------------------- | -------------------- | --------------------------------------------- | --------------------------------- | -------------------- |\n| Mol Stability (QM9) | 81.3 | 83.5  | 88.3 | 77.1 |\n\n---\n\n### Q4. Missing related works:\n\nThanks for introducing useful work to improve the paper. [1] and [2] propose using a joint distribution to replace independent coupling pairs between noise and data. Our method differs from previous research as follows: Previous work focuses on general domains like image generation, without exploring joint distribution design for complex geometries. Additionally, our EOT is conducted at the sample level, instead of batch level, allowing for parallelization.\n[3] advances 3D molecule generation by extending the equivariant diffusion model with a geometry latent space, orthogonal from our research. However, applying EquiFM to a similar latent space may be possible in the future. We will update references and discussions in the next version.\n\n[1] Pooladian et. al. Multisample Flow Matching: Straightening Flows with Minibatch Couplings, ICML 2023\n\n[2] Tong et. al. Conditional Flow Matching: Simulation-Free Dynamic Optimal Transport (Preprint)\n\n[3] Xu et. al. Geometric Latent Diffusion Models for 3D Molecule Generation. ICML 2023\n"}, {"review_id": "Nx2DoOQy2O", "reviewer": "Reviewer_tN9X", "summary": "The contribution is a new 3D generative model for molecules. The model is trained using a novel flow-matching objective. The flow-matching objective for coordinates of atoms is novel in that coordinates of 'source' atoms are permuted and rotated to align with 'target' atoms. There is also some exploration of how to set probability paths for different parts of the molecule description, namely the atom types, charges, and coordinates.", "questions": "Lines 288-289: it looks like some words are missing ('the achieves' does not make grammatical sense) but it seems the authors are claiming that in GEOM-DRUGS, almost every molecule has one or more atoms with incorrect valency. Is that really true? I thought that GEOM-DRUGS was a collection of drug-like molecules in realistic conformations.\n\nIn figure 3, it's surprising that RK4 with 4x the number of evaluations usually does worse than Euler. Why is RK4 so bad?\n\nIn section 4.2 please could the authors spell out what is claimed to be equivariant with respect to what?\n\nEquation (8): this $\\psi_t$ will have discontinuities with respect to $x$ at the values of $x$ where $\\pi$ changes. Is this a problem?\n\nHow are atom types and charges represented in $h$?  Are they one-hot encoded as in Hoogeboom et al's EDM?\n\n", "limitations": "Among the limitations of previous models that the authors say they address is the inability to generate large molecules. However, the paper does not show any good generated molecules with more than 9 heavy atoms.", "rating": 3, "confidence": 4, "soundness": 2, "presentation": 1, "contribution": 2, "strengths": "Flow-matching represents a very promising technique to improve on generative models like Hoogeboom et al.'s EDM, in terms of both training and sampling speed.  The idea of permuting and rotating source atom coordinates to straighten the flow that is matched is interesting, as is the question of how to set the relative corruption rates for different parts of a compound data type (coordinates, atom types, and charges). Figure 2 is beautiful.", "weaknesses": "The paper is not clearly written, and does not seem to have been proof-read (e.g., repetition in lines 26-29).\n\nSection 4.3 describes what one might try to achieve when choosing a probability path for $h$ but does not say how to do it. Is the answer that you calculate and plot the lines in Figure 4 and then pick the path whose line is visually closest to the $I(x_t, h_0)$ line?\n\nThe ablation studies are not thorough. In particular, permuting and rotating atom coordinates as in definition 4.3 is a key novel feature of the proposed model, and the authors should show quantitatively what effect it has on the quality of generated samples.  Line 331 refers to 'Tab 5.2' but I cannot find the table this refers to.\n", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for the valuable feedback and appreciate the effort in reviewing our work. We will address your concerns on the ablation studies, the evaluations, and the method details in the following paragraph, and any further comments are welcome!\n\n### Q1. Presentations and typos in Line26-Line29:\nThanks for pointing it out. We will carefully proofread the draft as suggested in the updated version. \n\n### Q2. Section 4.3 describes what one might try to achieve when choosing a probability path for $h$ but does not say how to do it. \n\nTo clarify, the key ingredient of this section is proposing the use of different probability paths for different modalities. We introduce an inductive bias by aligning the mutual information term $I(h_t,h_0)$ with $I(x_t,h_0)$. This is based on the intuition that molecular/chemical information emerges when the coordinates (x_t) are within a certain range, such as bond distances, to the ground truth. We believe the probability path of atom types $h_t$ should also reflect this emergence of information. Thus, we use mutual information as a measurable quantity to describe this tendency. Matching the tendencies of the two curves is sufficient to achieve appealing results, considering the uncertainty/bias in approximating mutual information $I(x_t,h_0). \n\nFollowing the reviewer's suggestion, we will introduce some quantity metrics to make the approach more rigorous in the next version. We have calculated several quantities to measure the tendencies of different probability trajectories, and these will be incorporated into the updated version.\n\n\n### Q3. In section 4.2 please could the authors spell out what is claimed to be equivariant with respect to what?\nApologies for the confusion caused. In section 4.2, the equivariant optimal transport is defined as a point mapping between two point clouds that is optimal for all E(3) equivariant operations on either point cloud. We understand that this definition can be misleading in terms of the method name (Equivariant Flow Matching), which actually refers to the fact that the modeled vector field is equivariant towards E(3) operations on the input. To avoid any misunderstanding, we will provide additional clarifications to make this distinction clear in the next version.\n\n### Q4. RK4 with 4x the number of evaluations usually does worse than Euler. Why is RK4 so bad?\nFor Euler, one step of integration requires only one NFE, while 4 NFEs are needed for one step of RK4 as it's a fourth-order Runge-Kutta method. For sufficient step time, the RK4 could be better than Euler as in Figure 3.  The performance drop when the step is insufficient could be due to the discrete property of $h$\n\n### Q5. $\\psi_t$ in Eq.8 will have discontinuities with respect to x where $\\pi$ changes. Is it a problem？\nThis is indeed a great question.\nFirstly, the $\\psi_t$ in Eq. 8 is a valid path for training continuous normalizing flows. This is due to the fact that $\\psi_t$ is time-dependent diffeomorphic [1], i.e., there exists a continuous vector field $v_t$ satisfies that $\\frac{d}{d t} \\psi_t(x)=v_t\\left(\\psi_t(x)\\right)$. Here $v_t = \\pi^*\\left(\\mathbf{R}^* \\mathbf{x}_1\\right)- \\mathbf{x}_0$. \nMeanwhile, as the reviewer mentioned, this $\\psi_t$ could have discontinuities with respect to x. This could hurt the generalization/robustness during sampling though the objective is unbiased. And we believe that regularizing the discontinuities could be a reasonable future direction to explore for better performance. \n\n[1] Lipmon et al, Flow Matching for Generative Modeling. ICLR 2023\n\n### Q6. How are atom types and charges represented in $h$? Are they one-hot encoded as in EDM?\nYes, we follow the data representation as in EDM. This, for atom types, we represent it by one-hot encoding and charges are represented as integer variables. \n\nOn the Evaluation and Extra Ablation Studies\n\n### Q7. The ablation studies are not thorough. In particular, the quantitative results of EOT maps. \nThanks for pointing it out, as you mentioned, the 'Tab 5.2' Line 331 is actually Table 3. In Table 3, we actually conduct ablation studies on the effect of EOT maps by fixing the probability path on $h$ as the $VP_{\\text{Linear}}$ path, while enumerating the probability path on atom coordinates as $EOT$, $OT$, $VP$. As we could find, the $EOT$ map could consistently bring performance improvements upon that without the $EOT$ map, e.g. vanilla $OT$ path or $VP$ path. To make the ablation studies more comprehensive, we add extra ablations in the following:\n\n\n| Metrics | $EOT+VP_{cos}$ | $EOT + VP_{sin}$ | $OT+VP_{cos}$ | $OT+VP_{sin}$ |\n| --- | --- | --- | --- | --- |\n| Atom Stability | 98.7 | 98.7 | 97.9 | 97.7 |\n| Mol Stability | 84.7 | 83.4 | 80.1 | 79.8 |\n\n\nWe will update the ablation studies in the updated version. \n\n### Q8. It is claimed that in GEOM-DRUGS, almost every molecule has one or more atoms with incorrect valency. \nAs the reviewer mentioned, the GEOM-DRUGS dataset comprises realistic conformations of drug-like molecules. In this context, it is important to clarify that we are not suggesting any incorrect valency in the real molecules within GEOM-DRUGS. The objective of the benchmarked 3D molecule generation task is to generate atom types and coordinates, followed by the addition of bonds using a predefined module [1,2].\n\nHowever, the process of adding bonds may introduce some bias. Consequently, using the same bond-adding process will result in ground truth data with atom stability lower than 100%. It's worth noting that molecule stability is approximately calculated as the N-th power of atom stability, where N represents the number of atoms in the molecule. Thus, for large molecules in the GEOM-DRUGS dataset, molecule stability is estimated to be close to 0%.\n\n[1]. Hoogeboom, Emiel, et al, Equivariant diffusion for molecule generation in 3d. ICML 2022. \n\n[2]. Wu, et al. Diffusion-based Molecule Generation with Informative Prior Bridges. NeurIPS 2022."}, {"review_id": "Mx3s5lq2dz", "reviewer": "Reviewer_EmYL", "summary": "This paper addresses the molecular generation problem.\nThe authors propose a conditional flow-matching-based method that employs different generative paths for coordinates and node-wise features. \nIn addition, EGNN is used for the SE(3) invariant vector field, and an ICP-like algorithm is used for the equivariant optimal transport map.\nThe proposed method outperforms the conventional method on the molecule generation benchmark tasks.", "questions": "* Do you have any original definitions in Section 3? If you do not mean to argue your originality in the section, it is better to clearly state that the section is a review of [26]. \n* What is the definition of $\\sigma_{min}$?\n* Some contents seem to be missing. For example, the detail of Section 4.3 is not included in Appendix C, although there is a reference in the main text.\n* Does EOT+VP_{linear} in Table 3 mean EOT path on x and VP_{linear} on h?  I cannot find the definition of + in (E)OT+VP_{xxx}.\n* Please explain how you generated atomic types in your model. Did EGNN output a one-hot vector representing atomic types?", "limitations": "The authors did not address the limitations.", "rating": 7, "confidence": 3, "soundness": 3, "presentation": 1, "contribution": 3, "strengths": "* The hybrid probability path is an attractive solution for the generative probability path of different modalities. \n* The equivariant OT map is a reasonable approach for 3D coordinate variables.\n* The proposed method can generate molecules not only with better quality but also faster.\n* The proposed method also outperforms a conventional method in conditional generation leading to many industrial applications.", "weaknesses": "* Since the proposed method is not clearly stated, the ablation study is not easy to understand, such as (E)OT+VP_{xxx} in Table 3.\n* The model architecture is poorly explained, so the generation of discrete variables, such as atomic types, is difficult to understand.\n* Some TeX references in the main text need to be corrected, such as Tab. 5.2 --> Table 3 in Section 5.4,  Fig. 5.4 --> Fig. 3 in Section 5.5, and Appendix C --> B.4 in Section 4.3.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you very much for the detailed and insightful comments! The responses to your concerns are listed below:\n\n### Q1. The ablation study is not easy to understand, such as (E)OT+VP_{xxx} in Table 3\nAs mentioned by the reviewer, in Table 3, the term \"$EOT+VP_{linear}$\" does refer to the method where the EOT path is applied to x and $VP_{linear}$ is applied to h. Thank you for bringing this to our attention. We will ensure that the corresponding sections are carefully proofread in order to eliminate any confusion in the next version.\n\n### Q2. The model architecture is poorly explained.  The generation of discrete variables, such as atomic types, is difficult to understand\n\nWe apologize for the confusion we may have caused. The model architecture/parameterization closely adheres to the EDM [1] paper, in order to ensure a fair comparison of the impact of the new training objective. During the generation of atom types, the EGNN produces continuous vectors at each timestep. The only additional step involved in generating discrete variables is the application of a quantized operation, such as $argmax$, to transform them into discrete vectors. \n\n### Q3.  Missing TEX references and contents, e.g. Details of section 4.3. \nThank you for bringing this to our attention. We acknowledge that the missing details of Section 4.3 can be found in Appendix B.4, as correctly pointed out by the reviewer. We will thoroughly review and address the reference issues mentioned, such as replacing Tab. 5.2 with Table 3 in Section 5.4, updating Fig. 5.4 to Fig. 3 in Section 5.5, and correcting the reference to Appendix C to be B.4 in Section 4.3. These corrections will be diligently made in the updated version.\n\n### Q4. it is better to clearly state that the section is a review of [26]\nThank you for your suggestion, which we greatly appreciate. The content related to flow-matching in Section 3 serves the purpose of providing essential background information and introducing relevant notations. In response to the reviewer's suggestion, we will explicitly clarify in the section that this content is intended to provide foundational knowledge rather than claim originality. By doing so, we hope to eliminate any potential misunderstanding regarding our contributions. Thank you for pointing this out, and we will make the necessary clarification in the revised version.\n\n### Q5. What is the definition of $\\sigma_{min}$\nPlease note that the introduction of $\\sigma_{min}$ serves the purpose of approximating each data point, specifically a delta distribution, by a Gaussian distribution with a narrow peak centered at the data point and a very small variance ($\\sigma_{min}$). This approximation is employed to prevent any corner cases and to facilitate a simpler mathematical representation of the probability path, as demonstrated in [1]. In response to the suggestion, we will enhance clarity by including this explanation in the upcoming version.\n\n[1] Lipmon et al, Flow Matching for Generative Modeling. ICLR 2023"}, {"review_id": "nI3hvo6Zhm", "reviewer": "Reviewer_KyMd", "summary": "The authors introduce a novel method named Equivariant Flow-Matching (EquiFM) for generating 3D molecules, aiming to enhance both categorical features (atom types) and continuous features (atom coordinates). The authors highlight the current limitations of diffusion models, particularly their instability and sampling inefficiencies. EquiFM improves upon these by using a flow-matching objective to stabilize the generative probability path of atom coordinates. Furthermore, it introduces a hybrid generative path to handle different modalities in the atomic feature space. This model utilizes an efficient ODE solver to enhance inference efficiency compared to existing SDE simulations. The authors report an improved performance with EquiFM, showing up to a 7% higher validity rate for large biomolecules and an average speed-up of 4.75x. ", "questions": "1. What is the computational complexity of the EquiFM model, and how does it scale with increasing molecule size which limits the applicability of EquiFM for peptides and antibodies? \n2. How might the performance of EquiFM change when applied to datasets with different distribution characteristics compared to QM9 and GEOM-DRUG?\n3. What steps have been taken to ensure the chemical feasibility and synthetic accessibility of the generated molecules?\n4. How was the model's hyperparameter configuration determined? Would a different configuration change the results significantly?\n5. How would the model performance be affected if evaluated with other key metrics, like novelty, or drug-likeness?\n6. Can the authors provide more detailed information on the hardware used for model training and inference, such as the specific GPU model, amount of RAM, and the number of parallel processes?\n7. On what basis was the 4.75x speedup of the model calculated? Is this in comparison to a specific baseline model or an average across several? Could the authors elaborate on the exact methodology used for this calculation?", "limitations": "While the study introduces an innovative approach to the generation of 3D molecules provides a rigorous comparison with baseline models, there are several notable limitations that warrant further exploration.\n\n1. A significant limitation of the current study is the relatively narrow scope of the references and baseline metrics utilized. Furthermore, the inclusion of more diverse baseline metrics would strengthen the validity of comparisons and aid in determining the true efficacy of the EquiFM model. For instance, discussing chemical similarity metrics between generated molecules and the training dataset could offer a more nuanced view of the model's capabilities.\n\n2. The current evaluation metrics do not include measures such as synthetic accessibility, which assesses the ease with which a generated molecule can be physically synthesized in a laboratory.  Tools like RDKit offer a proxy metric to estimate synthetic accessibility (https://github.com/rdkit/rdkit/blob/master/Contrib/SA_Score/sascorer.py), and the lack of such consideration especially in relation to the DRUG dataset.\n\n3. While the EquiFM model has shown promising results on the datasets used in this study, it remains unclear how it would perform on more diverse or complex molecular data. Expanding tests to include datasets like MOSES (https://github.com/molecularsets/moses) would allow a more comprehensive evaluation of the model's generalization capabilities.\n\nIn conclusion, while the EquiFM model has shown promise in generating 3D molecules, there is room for substantial expansion and improvement in future studies, particularly in the breadth of references and baseline metrics, and the inclusion of practical and diverse evaluation metrics.", "rating": 6, "confidence": 3, "soundness": 2, "presentation": 1, "contribution": 2, "strengths": "1. Equivariant Optimal Transport (EOT) is a novel approach to atom alignment, minimizing straight-line distance between paired atoms across all rotations and alignments.\n2. EOT-based training objective is invariant to initial translations and rotations, increasing the model's robustness against variances in sampled noise and data points.\n3. The iterative algorithm proposed for obtaining the EOT map is grounded in proven techniques (Hungarian and Kabsch algorithms), enhancing the solution's efficacy.\n4. The unique approach of aligning information quantity changes in the probability paths for the variables ensures better modeling of the joint variable.", "weaknesses": "\n1. The application of the Hungarian and Kabsch algorithms for obtaining the EOT map may introduce computational overhead due to their iterative nature, affecting efficiency.\n2. As the variables' probability paths are set independently, there might be cases when certain relationships or interactions between variables aren't accurately captured.\n3. The document assumes prior knowledge of the topic and includes numerous technical terms and mathematical formulas, potentially making it inaccessible to a broader audience. For example, computational chemists and drug discovery scientists. \n4. The use of a single rotation matrix for all possible rotations and alignments could potentially miss nuanced differences in complex data sets.\n5. The model's complexity may lead to challenges in implementation, particularly in situations with less computational resources.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank the reviewer for the constructive and insightful comments. We address all your concerns about the computational complexity, the evaluation, and the details of the methods in the following paragraph, and any further discussions are welcome. \n\n### Q1. Computational Complexity and Applicability for Larger Molecules\nDue to the page limit, the detailed response to the complexity issues could be referred to  Q1 of the response of Reviewer zq9q. And we provide extensive studies on the relationship between complexity and different molecule size. The full curve could be found in the uploaded pdf. \n\n### Q2. Limitations on Capturing the Relationship with Independent Probability Paths\nAccording to the conditional flow matching theorem (Theorem.1 in [3]), one can utilize the conditional probability path $p(x_t,h_t|x_1, h_1)$ to learn the joint marginal flow $p(x_1，h_1)$. It should be emphasized that this statement is applicable regardless of any requirements on the distribution $p(x_1,h_1)$ or the correlation between $x$ and $h$. Additionally, there are no restrictions on the form of the conditional probability path, which allows for a valid and unbiased choice of learning any joint distribution with $p(x_t|x_1,h_1) = p(x_t|x_1)$. To provide further intuition, it is worth noting that the commonly used diffusion path for image generation also fits into the above formula when considering $x$ and $h$ as two pixels within the same image.\n\n### Q3. The limitation of using a single rotation matrix\nAs the reviewer mentioned, the single rotation matrix could lack the capacity to capture details when the point cloud holds more nodes and a complex structure. Fortunately, such limitation could be feasibly addressed by extending the proposed EOT map to a hierarchical version. For example, We could first compute the global EOT map with a rotation matrix to only align the center positions of fragments/scaffolds, And then calculate the local EOT map among the nodes inside each fragment.\n\n### Q4. The performance change of EquiFM on distribution with different characteristics\nEquiFM is likely to excel in datasets with diverse conformation distributions due to the EOT map's ability to reduce learning space, minimizing optimization variance. However, modeling distributions where atom type and 3D structure are not strongly correlated, such as antibody variable regions, may pose challenges for the current version of EquiFM.\n\n\n### Q5. How to ensure chemical feasibility and synthetic accessibility with EquiFM.\nThus for a fair comparison with other fundamental models, we do not explicitly involve the components for chemical feasibility and synthetic accessibility in the current version. However, it could be very flexible for ensuring feasibility and synthetic accessibility with slight extra efforts based on EquiFM. For example, we could add an extra guided component on the vectors field like [2] to optimize the corresponding chemical feasibility and synthetic accessibility. \n\n \n### Q6. Sensitivity to Hyperparameter settings\n\nWe experimented with different settings for several optimization-related hyperparameters, e.g., learning rates and batch size, etc. And we found there is no significant impact of such hyperparameters.\n\n\n### Q7. Detailed configuration of the hardware\nThe configuration of our server is:\nCPU: Intel(R) Xeon(R) Platinum 8362 CPU @ 2.80GHz / 10 Core\nGPU: Nvidia-3090 GPU with 24 GB GPU memory / 1 GPUs\nMemory: 20G\n\n\n### Q8. On what basis was the 4.75x speedup of the model calculated?\nHere the 4.75x speedup is compared to EDM[1], which is an advanced diffusion model for the task with promising performance. Note the EDM uses fixed diffusion steps, e.g., sampling every molecule will use 1000 forward calls of the EGNN in the paper.While for our model, the generation process is essentially solving the neural ODE of the learned vector field. With adaptive ODE solvers(dopri 15), such a process could be accelerated by jumping the repeated evaluation where the vector field changes are not significant. In this way, the adaptive ODE solver could gain acceleration. As the two models share the same EGNN architectures, we calculate the acceleration factor as the forward call number of EDM divided by that of EquiFM. \n\n### Q9. More extensive evaluation metrics \nAs suggested, we add the following extra evaluation metrics： Quantitative estimate of drug-likeness (QED),Retrosynthetic accessibility (RA)， Medicinal chemistry filter (MCF)， Synthetic accessibility score (SAS), Molecular weight(MW),LogP and novelty. Besides, we evaluate the Conformation Energy Distance to evaluate the quality of 3D conformations. The experimental results can be found in the following:\n| Methods | QED（$\\uparrow$） | RA（$\\uparrow$） | MCF （$\\uparrow$） | SAS（$\\downarrow$） | $$\\Delta$$ MW（$\\downarrow$） | $$\\Delta$$ LogP （$\\downarrow$） | Conformation Energy distance（$\\downarrow$） | Novelty $\\uparrow$） |\n| --- | --- | --- | --- | --- | --- | --- | --- | --- |\n| EDM | 0.608 | 0.441 | 0.621 | 4.054 | 0.566 | 23.71 | 0.2180 | 0.791 |\n| EquiFM | 0.627 | 0.519 | 0.693 | 3.893 | 0.478 | 19.54 | 0.2081 | 0.834 |\n\nHere $\\Delta$ stands for the difference compared to ground truth distribution. And here to get valid molecules on GEOM-DRUGs, for both models, we sample Hydrogen coordinates with the help of RDkit following previous work [1]. \n\n\n### Q10. Extended the experimental results on other datasets\n After carefully checking the data and introduction in MOSEs, we find that the MOSEs mainly contain 2D information, e.g., SMILES and atom-bond graphs while the 3D conformation data is missings. Therefore, it would be not suitable to evaluate the proposed EquiFM which is a 3D molecule generative model.\n\n[1]. Hoogeboom, Emiel, et al, Equivariant diffusion for molecule generation in 3d.ICML 2022. \n[2]. Bao et al. Equivariant Energy-Guided SDE for Inverse Molecular Design. ICLR 2023\n[3] Lipmon et al, Flow Matching for Generative Modeling. ICLR 2023"}]

Author rebuttal:
We would like to express our sincere appreciation to the reviewers and the area chairs for the efforts and time spent reviewing our work and the crucial, insightful, and constructive comments. Here we would like to highlight the main concerns of the reviewers  and the corresponding responses here:
1. **Scalability Issues**, we have provided extensive experiments and detailed analysis of the relationship between the computation overhead and molecule size. And we show that this does not leash the scalability of our methods.
2. **Evaluation Metrics**, we have added several evaluation metrics to better compare the chemical property of generated molecules. 
3. **Ablation Study**, we have provided an extra ablation study for comparing the effect of different probability paths and the EOT maps. 
4. **Presentation**, we will address all the typos and presentation issues as suggested by the reviewers. 

Meta-review: The paper adapts the Flow Matching framework for the problem of 3D molecule generation and is able to achieve state of the art performance on common benchmarks. The paper presents interesting contributions including incorporation of equivariance in Flow Matching, Optimal Transport alignment to straighten learned generation paths, and using different probability paths for different data modalities. The weaknesses include issues with the exposition, missing experiments/ablations, scalability to large molecules, and questions about computational overhead of the Hungarian algorithms. The authors already addressed some of these issues in the rebuttal and are requested to revise the paper accordingly in the camera ready version.
Decision: Accept (poster)