Title: Equivariant Flow Matching with Hybrid Probability Transport for 3D Molecule Generation
Venue: NeurIPS 2023
Accepted: yes
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

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 × speed up of sampling on average. 2

1 Introduction

Geometric generative models aim to approximate the distribution of complex geometries and emerge as an important research direction in various scientific domains. A 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 [44] and the proteins could be seen as 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 [11] and protein engineering [45].

With the advancements of deep generative modeling, there has been a series of fruitful research progresses achieved in geometric generative modeling, especially molecular structures. For example, [9, 30] and [42] proposed data-driven methods to 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 as 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 [13] on applying DMs in this field have also demonstrated superior performance. However, existing

t=1

0=4

h

N

N

N

N

N

C

C

Sample

Noise

X

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 FlowMatching (EquiFM), for geometric generative modeling. Our method is inspired by the recent advancement of flow matching [26], a simulation-free objective for training CNFs that has demonstrated appealing generation performance with stable training and efficient sampling. Nevertheless, designing suitable geometric flow-matching objectives for molecular generation is non-trivial:

(1) the 3D skeleton modeling is sensitive, i.e. , a slight difference in the atom coordinates could affect the formulation of 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 as follows:

For stabling the 3D skeleton modeling, we introduce an Equivariant Optimal-Transport to 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, we proposed to differ the generative probability path of different components based on the information quantity and thus introduce a hybrid generative path. The hybrid-path techniques distinguish different modalities without adding extra modeling complexity or computational load.

The proposed model lies in the scope of continuous normalizing flow, which is parameterized by an ODE. We can use an efficient ODE solver during the molecule generation process to improve the inference efficiency upon the SDE simulation required in 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 on multiple benchmarks, including both unconditional and property-conditioned molecule generation. Results demonstrate that EquiFM can consistently achieve superior generation performance on all the metrics, and 4.75 × speed up on average. Empirical studies also show a significant improvement in controllable generation. All the empirical results demonstrate that the EquiFM enjoys a significantly higher modeling capacity and inference efficiency.

2 Related Work

Flow Matching and Diffusion Models Diffusion models have been studied in various research works such as [47, 12, 49], and have recently shown success in fields like high-dimensional statistics [41], language modeling [23], and equivariant representations [13]. 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 [33, 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, 1, 28] have proposed this simulation-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 [14, 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 structurebased drug design [24, 35, 38]. 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 a separate fitting algorithm.

In the past year, the attention has shifted towards using DMs for 3D molecule generation [13, 52, 53], with successful applications in target drug generation [25], antibody design [29], and protein design [2, 51]. However, our method is 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 q , x 0 represents a data point from q and x 1 represents a sample from the prior distribution p 1 . The time-dependent probability path is defined as p t ∈ [0 , 1] : R d → R > 0 , and the time-dependent vector field is defined as v t ∈ [0 , 1] : R d → R d . The vector field uniquely defines time-dependent flow ψ t ∈ [0 , 1] : R d → R d by the following ordinary differential equation (ODE):

[4] proposed to train the parameterized flow model ψ t called a continuous normalizing flow (CNF) with black-box ODE solvers. Such a model could reshape a simple prior distribution p 1 to the complex real-world distribution q . CNFs are difficult to train due to the need for numerical ODE simulations. [26] introduced flow matching, a simulation-free objective, by regressing the neural network v θ ( x, t ) to some target vector field u t ( x ) :

The objective L FM requires access to the vector field u t ( x ) and the corresponding probability path p t ( x ) . However, these entities are difficult to define in practice. Conversely, the conditional vector field u t ( x | x 0 ) and the corresponding conditional probability path p t ( x | x 0 ) are readily definable.

The probability path can be marginalized from a mixture of conditional probability path p t ( x ) = ∫ p t ( x | x 0 ) q ( x 0 )d x 0 , and the vector field u t ( x ) can be marginalized from conditional vector field as u t ( x ) = E x 0 ∼ q u t ( x | x 0 ) p t ( x | x 0 ) p t ( x ) . This illustrates how u t ( x ) and p t ( x ) are related to their conditional form, and [26] further proved that with the conditional vector field u t ( x | x 0 ) generating the conditional probability path p t ( x | x 0 ) , the marginal vector field u t ( x ) will generate the marginal probability path p t ( x ) . The observation inspires the new conditional flow matching (CFM) objective:

The CFM objective enjoys the tractability for optimization, and optimizing the CFM objective is equivalent to optimizing Eq. 2, i.e. , ∇ θ L FM ( θ ) = ∇ θ L CFM ( θ ) . For the inference phase, ODE solvers could be applied to solve the Eq. 1, e.g. , x 0 = ODESolve( x 1 , v θ , 1 , 0) . In this paper, we consider using the Gaussian conditional probability path, which lies in the form of p t ( x | x 0 ) = N ( x | µ t ( x 0 ) , σ t ( x 0 ) 2 I ) . We introduce two probability paths utilized in the following of paper:

Conditional Optimal Transport Path With the prior distribution p 1 defined as a standard Gaussian distribution, the empirical data distribution p 0 ( x | x 0 ) is approximated with a peaked Gaussian centered in x 0 with a small variance σ min as N ( x | x 0 , σ 2 min I ) . The probability path is p t ( x | x 0 ) = N ( x | (1 -t ) x 0 , ( σ min +(1 -σ min ) t ) 2 I ) and the corresponding flow is ψ t ( x ) = ( σ min + (1 -σ min ) t ) x + (1 -t ) x 0 . Then the vector field could be obtained by Eq. 1 as: u t ( ψ t ( x ) | x 0 ) = d d t ψ t ( x ) = -x 0 +(1 -σ min ) x .

Put the above terms into Eq. 3, the reparameterized objective is as:

Intuitively, the conditional optimal transport objective tends to learn the transformation direction from noise to data sample in a straight line which could hold appealing geometric properties.

Variance Preserving Path The variance-preserving (VP) path is defined as p t ( x | x 0 ) = N ( x | α t x 0 , (1 -α 2 t ) I ) , where α t = e -1 2 T ( t ) and T ( t ) = ∫ t 0 β ( s ) ds . Here β is some noise schedule function. Following the Theorem. 3 in [26], the target conditional vector field of VP path could be derived as u t ( x | x 0 ) = α ′ t 1 -α 2 t ( α t x -x 0 ) . α ′ t denotes the derivative with respect to time. And the objective for VP conditional flow matching is as:

The VP path is flexible to control the information dynamics, e.g. correlation changes towards x 0 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 we elaborate on the design details of the hybrid probability path of the equivariant variable and invariant variable in Section. 4.2 and Section. 4.3 respectively. A high-level schematic is provided in Figure. 1.

4.1 Equivariant Flow Matching

Recall molecule could be presented as the tuple g = ⟨ x , h ⟩ , where x = ( x 1 , . . . , x N ) ∈ X is the atom coordinates matrix and h = ( h 1 , . . . , h N ) ∈ R N × d is the node feature matrix, such as atomic type and charges. Here X = { x ∈ R N × 3 : 1 N ∑ N i =1 x i = 0 } is the Zero Center-of-Mass (Zero CoM) space, which means the average of the N elements should be 0 . We introduce the general form of equivariant flow matching in the following.

SE(3) Invariant Probability Path For modeling the density function in the geometric domains, it 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 the prior distribution and vector field in the Zero CoM space, i.e. 1 N ∑ N i =1 v ( x, t ) i = 0 . Formally, the rotational invariance could be satisfied by making the parameterized vector field equivariant and the prior p 1 invariant to the rotational transformations as shown in the following statement:

Theorem 4.1. Let ( v x θ ( g , t ) , v h θ ( g , t )) = v θ ( g , t ) , where v x θ ( g , t ) and v h θ ( g , t ) are the parameterized vector field for x and h . If the vector field is equivariant to any rotational transformation R , i.e. , v θ ( ⟨ Rx , h ⟩ , t ) = ( R ( v x θ ( g , t )) , v h θ ( g , t )) . With an rotational invariant prior function p 1 ( x , h ) , i.e. , p 1 ( Rx , h ) = p 1 ( x , h ) , then the probability path p θ,t generated by the vector field v θ ( · ) is also rotational invariant.

To make the vector field satisfy the equivariance constraint, we parameterize it with an Equivariant Graph Neural Network (EGNN) [43]. And more details can be found in Appendix C.

(a) EquiFM

(b) FM

(c) EDM

Hybrid Probability Modeling We refer to the target conditional vector field on each part as u x t ( g | g 0 ) and u h t ( g | g 0 ) correspondingly, then we could get the objective in the following formulation:

Proposition 4.2. There could be joint probability path p t ( g | g 0 ) which satisfies that p t ( g | g 0 ) = p t ( x | x 0 ) p t ( h | h 0 ) , and the conditional vector field on x and h is independent: u x t ( g | g 0 ) = u t ( x | x 0 ) , u h t ( g | g 0 ) = u t ( h | h 0 ) .

The above Proposition 4.2 states a special property of conditional flow matching, i.e. , in the multivariable setting the probability path of each variable could be designed independently. Such property is appealing in our setting, as x and h hold different data types and come from different manifolds, thus it is intuitive to use different probability paths for modeling and generating the two variables.

4.2 Coordinates Matching with Equivariant Optimal Transport

We focus on the generation of coordinates variable x . 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 a straight line. However, directly applying the objective could be problematic in 3D molecule generation. With x 0 as the point cloud from molecule distribution and x 1 from the prior distribution, the objective in Eq. 6 tends to move the atom based on a random alignment between the atoms. Optimizing the vector field toward such a direction could involve extra variance for training and lead to a twisted and unstable generation procedure as shown in molecule generation visualization Fig. 2(b) and Fig. 2(c).

To address the above-mentioned issue, we first introduce the concept of equivariant optimal transport (EOT) between two geometries as follows:

Definition 4.3. Given two point clouds, z = ( z 1 , . . . , z N ) ∈ R N × 3 and y = ( y 1 , . . . , y N ) ∈ R N × 3 . We define the equivariant optimal transport plan as

Here π is a permutation of N elements and R ∈ R 3 × 3 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 the point clouds. Therefore, the mappings are referred to as equivariant optimal transport(EOT).

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 p 0 and a sampled point cloud from p 1 as:

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:

Agood property of the objective with EOT is that the training characteristics are invariant to translation and rotation of initial x 1 , and equivariant with respect to both sampled noise x 1 and data point x 0 , which empirically contributes to more effective training.

Solving Equivariant Optimal Transport We propose an iterative algorithm to obtain the equivariant optimal transport map. The algorithm first conducts the Hungarian algorithm [20] to align the atoms between the initial geometry from p 1 and the ground truth geometry from p 0 ; 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 as 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: p t ( x | x 0 ) = p 0 ( x | x 0 ) , ∀ t < ϵ x and p t ( x | x 0 ) = p 1 ( x | x 0 ) , ∀ t ≥ ϵ x .

We define the p t ( h | h 0 ) similarly with a different parameter ϵ h .Here we consider the corner case that ϵ x → 0 and ϵ h → 1 , i.e. no noise for atom types from timestep 0 to timestep ϵ h 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 ϵ h 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-theoretic inspired quantity as the measurement to identify probability paths for learning the flow matching model on 3D molecules.

Definition 4.5. For distribution p 0 on the joint space g , and two corresponding conditional probability path p t ( x | g 0 ) and p t ( h | g 0 ) , we denote the I ( x t , h t ) as the mutual information for x t with distribution ∫ p t ( x | g 0 ) p 0 ( g 0 )d g 0 and h t with distribution ∫ p t ( h | g 0 ) p 0 ( g 0 )d g 0 .

Proposition 4.6. For the independent conditional probability path p t ( g | g 0 ) = p t ( x | x 0 ) p t ( h | h 0 ) , when the conditional probability path of x and h lies in OT path or VP path, if I ( x 0 , h 0 ) > 0 , then ∀ t ∈ (0 , 1) , I ( x t , h t ) > 0 and I ( x t i , h t i ) > I ( x t j , h t j ) , ∀ t i < t j .

We use the quantity I t ( x t , h t ) as the key property to distinguish different probability paths. Given the conditional probability path p t ( x | x 0 ) , it implies an information quantity change trajectory from I ( x 1 , h 1 ) = 0 to I ( x 0 , h 0 ) following I ( x t , h 0 ) . Thus, one well-motivated probability path on h is to align the information quantity changes by setting I ( h t , h 0 ) = I ( x t , h 0 ) . 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 β which could naturally adjust the information change by choosing different noise schedules, so we explore the probability path mainly on the VP path. For I ( h t , h 0 ) , we decompose it as I ( h t , h 0 ) = H ( h 0 ) -H ( h 0 | h t ) where H ( h 0 ) is constant and H ( h 0 | h t ) is the conditional entropy towards h 0 with h t as the logits. Similarly, the difficulty of estimation I ( x t , h 0 ) lies in I ( h 0 | x t ) . Following the difference of entropy estimator in [32], we build prediction model p ϕ ( h 0 | x t ) to estimate I ( h 0 | x t ) for selected time t . More details can be found in Appendix C. We demonstrate 20 time steps for I ( x t , h 0 ) , and I ( h t , h 0 ) for vanilla VP path

with the linear schedule (VPlinear) on β , VP path with cosine schedules(VPcos) [33] and polynomial schedules(VPpoly) [13], and the OT path in Fig. 4.

We observe that the information quantity of I ( x t , h 0 ) 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 a 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 I ( x t , h 0 ) also implies a generation procedure where the coordinates x transform first and the atom types h is then determined when x are relatively stable.

5 Experiments

In this section, we justify the advantages of the proposed EquiFM with comprehensive experiments. The experimental setup is introduced in Section 5.1. Then we report and analyze the evaluation results for the unconditional and conditional settings in Section 5.2 and 5.3. We provide detailed ablation studies in Section 5.4 to 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, 42, 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 on 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 QM9 dataset [39], which has been widely adopted in previous 3D molecule generation studies [9, 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, 13].

5.2 Molecular Modeling and Generation

Evaluation Metrics. The model performance is 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 [13].

Next, we evaluate the quality of our predicted molecular graph by calculating both atom stability and molecule stability metrics. The atom stability metric measures the proportion of atoms that have a correct valency, while the molecule stability metric quantifies the percentage of generated molecules in which all atoms are stable. Additionally, we report validity and uniqueness metrics that indicate the

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-Schnet [9] 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 objective is simulation-based. Equivariant Graph Diffusion Models ( EDM ) with its non-equivariant variant ( GDM ) [13] are recent progress on diffusion models for molecule generation. Most recently, [52] proposed an improved version of EDM ( EDM-Bridge ), which improves upon the performance of EDM by incorporating well-designed informative prior bridges. To yield a fair comparison, all the model-agnostic configurations are set as the same as described in Sec. 5.1.

Results and Analysis. We generate 10 , 000 samples from each method to calculate the above metrics, and the results are reported in Table 1. 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 are subsequently added using a 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 a 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 GEOM-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 F and the efficiency study in Appendix E.

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 α , orbital energies ε HOMO , ε LUMO and their gap ∆ ε , Dipole moment µ , and heat capacity C v . For evaluating the model's capacity to conduct property-conditioned generation, we follow the [42] to first split the QM9 training set into two halves with 50 K samples in each. Then we train a property prediction network ω on the first half and train conditional models in the second half. Afterward, given a range of property values s , we conditionally draw samples from the generative models and then use ω to calculate their property values as ˆ s . The Mean Absolute Error (MAE) between s and ˆ s is reported to measure whether generated molecules are 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 ω . A smaller gap with QM9 numbers indicates a better property-conditioning performance.

Baselines. We incorporate existing EDM as our baseline model. In addition, we follow [13] to also list two baselines agnostic to ground-truth property s , named Random and N atoms . Random means we simply do random shuffling of the property labels in the dataset and then evaluate ω on it. This operation removes any relation between molecule and property, which can be viewed as an upper bound of MAE metric. N atoms 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

0

200

400

600

800

1

,

000

0

.

4

0

.

5

0

.

6

0

.

7

0

.

8

0

.

9

Number of Function Evaluations (NFE)

Molecule Stability

EquiFM-Dopri5

EquiFM-Euler

EquiFM-MidPoint

EquiFM-Rk4

EDM

0

0

.

2

0

.

4

0

.

6

0

.

8

1

0

0

.

2

0

.

4

0

.

6

0

.

8

1

t

relative

I

(

h

t

, h

0

)

or

I

(

x

t

, h

0

)

OT

VPlinear

VPpoly

VPcos

I

(

x

t

,

h

0

)

property information into the generated molecules. And overcoming N atoms further indicates the model can incorporate conditioning into molecular structures beyond the number of atoms.

Results and Analysis. The visualizations of conditioned generation can be found in Appendix F. As shown in Table 2, for all the conditional generation tasks, our proposed EquiFM outperforms other comparable models with a 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 on the coordinate variable x and the categorical variable h ? 2) how does the equivariant optimal transport path boost the generation?

To answer these questions, we apply several different probability paths and compare them on the QM9 dataset, including the variance-preserving (VPLinear) path(Eq. 5), vanilla optimal transport (OT) path(Eq. 4), and the equivariant optimal(Eq. 9) transport path(EOT) on the coordinate variable x ; And variance-preserving (VPLinear) path(Eq. 5), vanilla optimal transport (OT) (Eq. 4), the variancepreserving path with polynomial decay (VPPoly), variance preserving path with cosine schedule (VPCos). The result is illustrated in Tab. 3. We notice that OT-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-based prior, the EOT path could further boost the performance by a large margin. To gain a more intuitive understanding, we further provide the generation dynamic comparison in Fig. 2(b). As shown, the generation procedure trained with vanilla OT path, though more stable than the EDM generation procedure, also exits some twisted phenomenon, i.e. , all atoms tend to first contract together and then expand; Such phenomenon disappears in the generation procedure of EOT path due to that the generation direction is well constrained. For the probability path on the categorical variable, we find the VP path, holds the superior performance due to the closest alignment with the information quantity changes. If there is a significant discrepancy in the information quantity dynamics, e.g. , OT path, it may result in a substantial decline in performance.

5.5 Sampling Efficiency

We also evaluate the sampling efficiency of our model, as shown in Fig. 3, the results of EquiFM with 4 different integrating algorithms converge to state-of-the-art results in much less NFE compared to baseline model EDM. Remarkably, the red triangle is the result of EquiFM with the 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 Eulers 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 a much shorter generation route during generation, this can be justified with visualization Fig. 2(a).

6 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 to 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 on coordinates and ensures that the information changes on each component of the joint path are appropriately matched. Consequently, EquiFM learns the underlying 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 a speedup of 4.75 × compared to 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 (2022ZD0117502), Guoqiang Research Institute General Project, Vanke Special Fund for Public Health and Health Discipline Development, Tsinghua University (NO.20221080053), and Beijing Academy of Artificial Intelligence (BAAI), Tsinghua University (No. 2021GQG1012). Minkai Xu thanks the generous support of Sequoia Capital Stanford Graduate Fellowship.

References

Michael S Albergo and Eric Vanden-Eijnden. Building normalizing flows with stochastic interpolants. arXiv preprint arXiv:2209.15571 , 2022.

Namrata Anand and Tudor Achim. Protein structure and sequence generation with equivariant denoising diffusion probabilistic models. arXiv preprint arXiv:2205.15019 , 2022.

Brandon Anderson, Truong Son Hy, and Risi Kondor. Cormorant: Covariant molecular neural networks. Advances in neural information processing systems , 32, 2019.

Ricky TQ Chen, Yulia Rubanova, Jesse Bettencourt, and David K Duvenaud. Neural ordinary differential equations. Advances in neural information processing systems , 31, 2018.

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-548. IEEE, 2002.

CRC Handbook. CRC Handbook of Chemistry and Physics, 88th Edition . CRC Press, 88 edition, 2007.

David F. Crouse. On implementing 2d rectangular assignment algorithms. IEEE Transactions on Aerospace and Electronic Systems , 52(4):1679-1696, 2016.

Prafulla Dhariwal and Alexander Nichol. Diffusion models beat gans on image synthesis. Advances in Neural Information Processing Systems , 34:8780-8794, 2021.

Niklas Gebauer, Michael Gastegger, and Kristof Schütt. Symmetry-adapted generation of 3d point sets for the targeted discovery of molecules. Advances in neural information processing systems , 32, 2019.

Niklas WA Gebauer, Michael Gastegger, Stefaan SP Hessmann, Klaus-Robert Müller, and Kristof T Schütt. Inverse design of 3d molecular structures with conditional generative neural networks. arXiv preprint arXiv:2109.04824 , 2021.

Jordan Graves, Jacob Byerly, Eduardo Priego, Naren Makkapati, S Vince Parish, Brenda Medellin, and Monica Berrondo. A review of deep learning methods for antibodies. Antibodies , 9(2):12, 2020.

Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models. arXiv preprint arXiv:2006.11239 , 2020.

Emiel Hoogeboom, Vıctor 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.

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.

Bowen Jing, Stephan Eismann, Patricia Suriana, Raphael John Lamarre Townshend, and Ron Dror. Learning from protein structure with geometric vector perceptrons. In International Conference on Learning Representations , 2021.

Jaehyeong Jo, Seul Lee, and Sung Ju Hwang. Score-based generative modeling of graphs via the system of stochastic differential equations. arXiv preprint arXiv:2202.02514 , 2022.

Wolfgang Kabsch. A solution for the best rotation to relate two sets of vectors. Acta Crystallographica Section A: Crystal Physics, Diffraction, Theoretical and General Crystallography , 32(5):922-923, 1976.

Diederik Kingma, Tim Salimans, Ben Poole, and Jonathan Ho. Variational diffusion models. Advances in neural information processing systems , 34:21696-21707, 2021.

Diederik P Kingma and Jimmy Ba. Adam: A method for stochastic optimization. In 3nd International Conference on Learning Representations , 2014.

Harold W Kuhn. The hungarian method for the assignment problem. Naval research logistics quarterly , 2(1-2):83-97, 1955.

G Landrum. Rdkit: open-source cheminformatics http://www. rdkit. org. 2016.

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.

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.

Yibo Li, Jianfeng Pei, and Luhua Lai. Structure-based de novo drug design using 3d deep generative models. Chemical science , 12(41):13664-13675, 2021.

Haitao Lin, Yufei Huang, Meng Liu, Xuanjing Li, Shuiwang Ji, and Stan Z Li. Diffbp: Generative diffusion of 3d molecules for target protein binding. arXiv preprint arXiv:2211.11214 , 2022.

Yaron Lipman, Ricky TQ Chen, Heli Ben-Hamu, Maximilian Nickel, and Matt Le. Flow matching for generative modeling. arXiv preprint arXiv:2210.02747 , 2022.

Qi Liu, Miltiadis Allamanis, Marc Brockschmidt, and Alexander Gaunt. Constrained graph variational autoencoders for molecule design. In Advances in neural information processing systems , 2018.

Qiang Liu. Rectified flow: A marginal preserving approach to optimal transport. arXiv preprint arXiv:2209.14577 , 2022.

Shitong Luo, Yufeng Su, Xingang Peng, Sheng Wang, Jian Peng, and Jianzhu Ma. Antigenspecific antibody design and optimization with diffusion-based generative models for protein structures. In Alice H. Oh, Alekh Agarwal, Danielle Belgrave, and Kyunghyun Cho, editors, Advances in Neural Information Processing Systems , 2022.

Youzhi Luo and Shuiwang Ji. An autoregressive flow model for 3d molecular geometry generation from scratch. In International Conference on Learning Representations , 2021.

Tomohide Masuda, Matthew Ragoza, and David Ryan Koes. Generating 3d molecular structures conditional on a receptor binding site with deep generative models. arXiv preprint arXiv:2010.14442 , 2020.

David McAllester and Karl Stratos. Formal limitations on the measurement of mutual information. In International Conference on Artificial Intelligence and Statistics , pages 875-884. PMLR, 2020.

Alexander Quinn Nichol and Prafulla Dhariwal. Improved denoising diffusion probabilistic models. In International Conference on Machine Learning , pages 8162-8171. PMLR, 2021.

Adam Paszke, Sam Gross, Soumith Chintala, Gregory Chanan, Edward Yang, Zachary DeVito, Zeming Lin, Alban Desmaison, Luca Antiga, and Adam Lerer. Automatic differentiation in pytorch. In NIPS-W , 2017.

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 on Machine Learning , 2022.

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):2495-2506, 2016.

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.14772 , 2023.

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.

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.

Ali Razavi, Aaron Van den Oord, and Oriol Vinyals. Generating diverse high-fidelity images with vq-vae-2. Advances in neural information processing systems , 32, 2019.

Robin Rombach, Andreas Blattmann, Dominik Lorenz, Patrick Esser, and Björn Ommer. Highresolution image synthesis with latent diffusion models. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition , pages 10684-10695, 2022.

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 arXiv:2105.09016 , 2021.

Vıctor Garcia Satorras, Emiel Hoogeboom, and Max Welling. E(n) equivariant graph neural networks. In International conference on machine learning , pages 9323-9332. PMLR, 2021.

Kristof T Schütt, Farhad Arbabzadah, Stefan Chmiela, Klaus R Müller, and Alexandre Tkatchenko. Quantum-chemical insights from deep tensor neural networks. Nature communications , 8:13890, 2017.

Chence Shi, Chuanrui Wang, Jiarui Lu, Bozitao Zhong, and Jian Tang. Protein sequence and structure co-design with equivariant translation. arXiv preprint arXiv:2210.08761 , 2022.

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 arXiv:2001.09382 , 2020.

Jascha Sohl-Dickstein, Eric A Weiss, Niru Maheswaranathan, and Surya Ganguli. Deep unsupervised learning using nonequilibrium thermodynamics. arXiv preprint arXiv:1503.03585 , 2015.

Yang Song, Conor Durkan, Iain Murray, and Stefano Ermon. Maximum likelihood training of score-based diffusion models. Advances in Neural Information Processing Systems , 34:14151428, 2021.

Yang Song and Stefano Ermon. Generative modeling by estimating gradients of the data distribution. In Advances in Neural Information Processing Systems , volume 32, pages 1191811930. Curran Associates, Inc., 2019.

Alexander Tong, Nikolay Malkin, Guillaume Huguet, Yanlei Zhang, Jarrid Rector-Brooks, Kilian Fatras, Guy Wolf, and Yoshua Bengio. Conditional flow matching: Simulation-free dynamic optimal transport. arXiv preprint arXiv:2302.00482 , 2023.

Brian L 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.04119 , 2022.

Lemeng Wu, Chengyue Gong, Xingchao 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.

Minkai Xu, Alexander Powers, Ron Dror, Stefano Ermon, and Jure Leskovec. Geometric latent diffusion models for 3d molecule generation. arXiv preprint arXiv:2305.01140 , 2023.

Chengxi Zang and Fei Wang. Moflow: an invertible flow model for generating molecular graphs. In Proceedings of the 26th ACM SIGKDD International Conference on Knowledge Discovery &Data Mining , pages 617-626, 2020.

A Sampling and Training Algorithm

We provide a detailed training and sampling pipeline in this section. The training algorithm with EOT path on the equivariant variable x and VP path on the invariant variable h as an example is demonstrated in Algorithm 1.

The sampling algorithm can be found in Algorithm 2.

Algorithm 2 Sampling Algorithm of EquiFM

1: Input: vector field model v θ

2: g 1 ∼ N ( 0 , I )

3: Subtract center of gravity from x 1 in g 1 = [ x 1 , h 1 ]

4: g 0 = ODESolve( g 1 , v θ , 1 , 0)

6: sample ˆ p ( x 0 , h 0 | g 0 )

5: {During the ODE solving, always subtract the center of gravity from v x θ }

7: return ⟨ x 0 , h 0 ⟩

Note the ˆ p ( ·| g 0 ) stands for the procedure of transforming the continuous h 0 into the specific data modality. This is, for the categorical part, ˆ p ( h | g 0 ) = C ( h | h 0 ) and for the integer part p ( h | g 0 ) = ∫ h + 1 2 h -1 N ( u | h 0 , σ 0 ) d u .

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 an invariant prior p 1 ( x ) and the equivariant vector field v θ ( x , t ) , the marginal distribution implied by the vector field of each time step, p θ,t ( x ) is also invariant.

Proof. We are given that p 1 ( x ) is invariant, and that v θ ( x , t ) is equivariant, i.e. for any rotation R , p 1 ( Rx ) = p 1 ( x ) , and v θ ( Rx , t ) = R v θ ( x , t ) . Our target is to prove that ∀ t ∈ [0 , 1] , ∀ R , p θ,t ( Rx ) = p θ,t ( x ) . Then specifically, the sampling distribution p 0 ( x ) is invariant.

First recall that the way we generate the distribution p θ,t ( x ) is to exert a transformation ψ θ,t to the prior p 1 ( x ) . And the definition of the vector field is v θ ( x , t ) = d d t ψ θ,t ( x ) .

We can derive the equivariance of ψ θ,t by the following equations:

Note that ψ 1 ( Rx ) = R ψ 1 ( x ) since ψ 1 ( x ) = x , we have ψ θ,t ( Rx ) = R ψ θ,t ( x ) . That is to say, ψ θ,t is an equivariant transformation. Thus its inverse ψ -1 θ,t is also equivariant, since ∀ y = ψ θ,t ( x ) , we have Ry = R ψ θ,t ( x ) = ψ θ,t ( Rx ) , so ψ -1 θ,t ( R y ) = Rx = R ψ -1 θ,t ( y ) . Also, the Jacobian matrix ∂ψ θ,t ( x ) ∂ x is equivariant, i.e. ∂ψ θ,t ( u ) ∂ u | u = Rx = R ∂ψ θ,t ( u ) ∂ u | u = x , which implies that det ∂ψ θ,t ( u ) ∂ u | u = Rx = det ∂ψ θ,t ( u ) ∂ u | u = x , since the det function keeps constant under any rotation. According to the Change of Variable Theorem,

Applying the above conclusions together, we have

B.2 Explanation of Proposition 4.2

Note for the initial prior distribution, p 1 ( g | g 0 ) , could be the standard distribution, i.e. p 1 ( g | g 0 ) = N ( 0 , I ) . In this case, p 1 ( g | g 0 ) = p 1 ( x | x 0 ) p 1 ( h | h 0 ) . And for the time step zero, if we assume the distribution is a Gaussian centralized on the g 0 , i.e. , p 0 ( g | g 0 ) = N ( g 0 , σ min I ) . And in this case we also have that p 0 ( g | g 0 ) = p 0 ( x | x 0 ) p 0 ( h | h 0 ) .

Example 2: For the Gaussian probability path,

where µ : [0 , 1] × R d → R d is the time-dependent mean of the Gaussian distribution, while σ : [0 , 1] × R → R > 0 describes a time-dependent scalar standard deviation (std).

The Gaussian probability path satisfies that p t ( g | g 0 ) = p t ( x | x 0 ) p t ( h | h 0 ) . To better clarify, we highlight the difference between the conditional probability path and the marginal probability path with the following Remark B.1.

̸

Remark B.1 . With the conditional probability path on x and h being independent of each other, the marginal distribution could be correlated, i.e. p t ( g ) = p t ( x ) p t ( h )

With this property, we could design different paths for modeling complex variables.

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 x .

Specifically, for any rotation T acting on the point cloud x with N points, the new R ∗ corresponding with Tx (we denote it by R ∗ rot ) exactly offsetting the impact of T . Formally, let x 0 denote the target point cloud, we calculate

We claim that π ∗ rot = π ∗ , R ∗ rot = R ∗ T -1 , and a strict proof follows.

Note that if ϕ ( R ) : R 3 × 3 → R 3 × 3 is a reversible map, then for any scalar function f ( R ) , argmin R f ( ϕ ( R )) = ϕ -1 (argmin R f ( R )) . Here let ϕ ( R ) = RT , and ϕ -1 ( R ) = RT -1 , we get

where ˆ ϕ is an natural extension of ϕ defined as ˆ ϕ (( π, R )) = ( π, ϕ ( R )) , ∀ ( π, R ) .

Now we recheck the probability path p t in Eq. 8. Since p 1 is invariant under rotations, and the transformation ψ EOT t satisfies

i.e. ψ EOT t is invariant. So we conclude that p t = [ ψ EOT t ] ∗ p 1 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 on x or h , i.e. p t ( x | x 0 ) = N ( x | (1 -t ) x 0 , ( σ min +(1 -σ min ) t ) 2 I ) or p t ( x | x 0 ) = N ( x | α t x 0 , (1 -α 2 t ) I ) , the mutual information between the marginal variable x t and h t monotonically decays following the path from time step 0 to time step 1 where I 0 ( x 0 , h 0 ) > 0 . Note that the I 1 ( x 1 , h 1 ) = 0 , as p 1 ( g ) = p 1 ( x ) p 1 ( h ) . Recall the definition of signal-to-noise ratio (SNR) as:

For OT-path, SNR = (1 -t ) 2 t 2 ; and for VP-path SNR = α 2 t 1 -α 2 t . The key observation is that the SNR along the probability path of both x and h on either path will decay monotonically. Intuitively, with t → 1 , x t has less information of x 0 thus has less of 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. 5.

To this end, we demonstrate the relationship between the x 0 , h 0 , x t 1 . Recall the definition of conditional independent probability path as p t ( x | ⟨ x 0 , h 0 ⟩ ) = p t ( x | x 0 ) . Thus we have the distribution of random variable x t 1 , p t 1 ( x ) = ∫ p t 1 ( x | x 0 ) p ( x 0 )d x 0 . And hence, there is x t 1 ⊥ h 0 | x 0 . Similarly, we could also derive h t 1 ⊥ x 0 | h 0 .

Next, for any time step 0 < t 1 < t 2 < 1 . We will then demonstrate the dependency relationship between x 0 , x t 1 , x t 2 . We denote the target vector field which generates the marginal probability path p t ( x ) as u t ( x ) , then the distribution of random variable x t 2 could be then derived as p t 2 ( x ) = [ ∫ t 2 t 1 u t ( x )d t ] ∗ p t 1 ( x ) = [ ψ u t 2 -ψ u t 1 ] ∗ p t 1 ( x ) , where ψ u t denotes the transformation corresponding with u t . Then we have x t 2 ⊥ x 0 | x t 1 . Similarly, h t 2 ⊥ h 0 | h t 1 could be also derived.

With the above two conclusions, we demonstrate the dependency relationship shown in Fig. 5 holds in the hybrid probability path. Therefore, with the dependency relationship we could directly obtain that I ( x t 2 , h t 2 ) ≤ I ( x t 1 , h t 1 ) . If α t in VP path satisfies α t > 0 when 0 < t < 1 , then for any combination of such OT path and VP path, we always have I ( x t , h t ) > 0 .

C Implementation Details

C.1 Solving EOT with a variant of iterative closest point (ICP) algorithm.

Problem definition. Given a point cloud z ∈ R 3 × N and its reference point cloud y ∈ R 3 × N , note they are point cloud representations in Euclidean space. The objective is to find an optimal permutation matrix Π ∗ ∈ R N × N and a rotation matrix R ∗ ∈ R 3 × 3 that minimizes the following objective:

We optimize the objective iteratively with a variant of the iterative closest point (ICP) algorithm, where it iteratively obtains Π 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.

Algorithm 3 A variant of iterative closest point (ICP) algorithm. 1: Input: a point cloud z ∈ R 3 × N and it's reference point cloud y ∈ R 3 × N . 2: while τ has not converged do 3: Obtain permutation matrix Π = argmin Π ∥ Π ( Rz ) ⊤ -y ⊤ ∥ 2 with Jonker-Volgenant algorithm [7] 4: Obtain rotation matrix R = argmin R ∥ R ( Πz ) ⊤ -y ⊤ ∥ 2 with Kabsch algorithm [22] 5: τ = ∥ Π ( Rz ⊤ ) ⊤ -y ⊤ ∥ 2 6: end while 7: return Π , R

The vector field network is implemented with EGNNs [43] by PyTorch [34] package. We set the dimension of latent invariant features k to 1 for QM9 and 2 for DRUG, which extremely reduces the atomic feature dimension. For the training of vector field network v θ : on QM9, we train EGNNs with 9 layers and 256 hidden features with a batch size 64 ; and on DRUG, we train EGNNs with 4 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 [19] with a constant learning rate of 10 -4 as our default training configuration. The training on QM9 takes approximately 2000 epochs, and on DRUG takes 20 epochs.

D Number of Evaluation (NFE) analysis

We further explore the behavior of adaptive integrators during sampling with Dopri15 as an example. In Fig. 7, we show the average NFE at different time intervals. We could observe that at time intervals near 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 of O ( n 2 ) for computing an 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 Tab. 4 below 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.

(a) iterations towards atom numbers

(b) time towards atom numbers

iterationnode_numberrelation

7.5

time-nodenumberrelation

7.0

0.035

6.5

0.030

iterationnumber

6.0

(seconds)

0.025

5.5

0.020

5.0

time

0.015

4.5

0.010

4.0

0.005

3.5

0.000

0

25

50

75

100

125

150

175

200

0

25

50

75

100

125

150

175

200

nodenumber

nodenumber

Even 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

NFE Histogram

80

Number of Function Evaluations

60

40

20

0.0

0.2

0.4

0.6

0.8

1.0

t(time)

parallelizable nature. By Enabling prefetch and multiprocessing, we can minimize the computational overhead further and make it virtually inconsequential.

F 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. 8 and Fig. ?? , respectively. All examples are randomly generated without cherry-picking, but the viewing direction may affect 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 is more common when trained on the large molecule DRUG dataset, as shown in the second molecule in Fig. ?? . This phenomenon is not unique to our model but is a common issue in non-autoregressive molecule generative models [54, 16]. 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. 9. We interpolate the conditioning parameter, polarizability, with different values of α , while keeping the prior g 1 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 to correspond to higher α values. This observation is consistent with our results in Fig. 9.

63.11

66.76

70.41

74.06

77.71

81.36

85.01

88.66

92.31

95.96

99.62

103.27

106.92

110.57

114.22

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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.\n\nSTRENGTHS:\n1. 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.\n\nWEAKNESSES:\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.\n\nQUESTIONS:\n1. 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?\n\nLIMITATIONS:\nWhile 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.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 2, "DATE": "2023-07-01"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS:\n* 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.\n\nWEAKNESSES:\n* 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.\n\nQUESTIONS:\n* 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?\n\nLIMITATIONS:\nThe authors did not address the limitations.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-02"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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.\n\nSTRENGTHS:\nFlow-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.\n\nWEAKNESSES:\nThe 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\nQUESTIONS:\nLines 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\nLIMITATIONS:\nAmong 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.", "TITLE": "Official Review", "RECOMMENDATION": 3, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 2, "DATE": "2023-07-06"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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.\n\nSTRENGTHS:\n- 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.\n\nWEAKNESSES:\n- **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\nQUESTIONS:\n- 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.\n\nLIMITATIONS:\nlimitations are not discussed.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-06"}