SlowGuess/NeurIPS_2023_OCR huggingface dataset:
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(\mathbf{x}|\mathbf{x}_0) = p_0(\mathbf{x}|\mathbf{x}_0), \forall t < \epsilon_{\mathbf{x}}$  and  $p_t(\mathbf{x}|\mathbf{x}_0) = p_1(\mathbf{x}|\mathbf{x}_0), \forall t \geq \epsilon_{\mathbf{x}}$ .
We define the  $p_t(\mathbf{h}|\mathbf{h}_0)$  similarly with a different parameter  $\epsilon_{\mathrm{h}}$ . Here we consider the corner case that  $\epsilon_{\mathbf{x}} \rightarrow 0$  and  $\epsilon_{\mathbf{h}} \rightarrow 1$ , i.e. no noise for atom types from timestep 0 to timestep  $\epsilon_{\mathbf{h}}$  and max noise level from  $\epsilon_{\mathbf{h}}$  to timestep 1. (Reversely for  $\epsilon_{\mathbf{x}}$ ) Under such a probability path, the model will be encouraged to determine and fix the node type at around  $\epsilon_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  $\mathbf{g}$ , and two corresponding conditional probability path  $p_t(\mathbf{x}|\mathbf{g}_0)$  and  $p_t(\mathbf{h}|\mathbf{g}_0)$ , we denote the  $I(\mathbf{x}_t, \mathbf{h}_t)$  as the mutual information for  $\mathbf{x}_t$  with distribution  $\int p_t(\mathbf{x}|\mathbf{g}_0)p_0(\mathbf{g}_0)\mathrm{d}\mathbf{g}_0$  and  $\mathbf{h}_t$  with distribution  $\int p_t(\mathbf{h}|\mathbf{g}_0)p_0(\mathbf{g}_0)\mathrm{d}\mathbf{g}_0$ .
Proposition 4.6. For the independent conditional probability path  $p_t(\mathbf{g}|\mathbf{g}_0) = p_t(\mathbf{x}|\mathbf{x}_0)p_t(\mathbf{h}|\mathbf{h}_0)$ , when the conditional probability path of  $\mathbf{x}$  and  $\mathbf{h}$  lies in  $OT$  path or  $VP$  path, if  $I(\mathbf{x}_0,\mathbf{h}_0) > 0$ , then  $\forall t\in (0,1),I(\mathbf{x}_t,\mathbf{h}_t) > 0$  and  $I(\mathbf{x}_{t_i},\mathbf{h}_{t_i}) > I(\mathbf{x}_{t_j},\mathbf{h}_{t_j}),\forall t_i < t_j$ .
We use the quantity  $I_{t}(\mathbf{x}_{t},\mathbf{h}_{t})$  as the key property to distinguish different probability paths. Given the conditional probability path  $p_t(\mathbf{x}|\mathbf{x}_0)$ , it implies an information quantity change trajectory from  $I(\mathbf{x}_1,\mathbf{h}_1) = 0$  to  $I(\mathbf{x}_0,\mathbf{h}_0)$  following  $I(\mathbf{x}_t,\mathbf{h}_0)$ . Thus, one well-motivated probability path on  $\mathbf{h}$  is to align the information quantity changes by setting  $I(\mathbf{h}_t,\mathbf{h}_0) = I(\mathbf{x}_t,\mathbf{h}_0)$ . Based on such intuition, we design our probability path on  $\mathbf{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  $\beta$  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(\mathbf{h}_t,\mathbf{h}_0)$ , we decompose it as  $I(\mathbf{h}_t,\mathbf{h}_0) = H(\mathbf{h}_0) - H(\mathbf{h}_0|\mathbf{h}_t)$  where  $H(\mathbf{h}_0)$  is constant and  $H(\mathbf{h}_0|\mathbf{h}_t)$  is the conditional entropy towards  $\mathbf{h}_0$  with  $\mathbf{h}_t$  as the logits. Similarly, the difficulty of estimation  $I(\mathbf{x}_t,\mathbf{h}_0)$  lies in  $I(\mathbf{h}_0|\mathbf{x}_t)$ . Following the difference of entropy estimator in [32], we build prediction model  $p_{\phi}(\mathbf{h}_0|\mathbf{x}_t)$  to estimate  $I(\mathbf{h}_0|\mathbf{x}_t)$  for selected time  $t$ . More details can be found in Appendix C. We demonstrate 20 time steps for  $I(\mathbf{x}_t,\mathbf{h}_0)$ , and  $I(\mathbf{h}_t,\mathbf{h}_0)$  for vanilla VP path
Table 1: Results of atom stability, molecule stability, validity, and validity × uniqueness. A higher number indicates a better generation quality. The results marked with an asterisk were obtained from our own tests.
# Metrics	QM9	DRUG
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-Bridge	98.8	84.6	92.0*	90.7	82.4	92.8*
EQUIFM	98.9 ± 0.1	88.3 ± 0.3	94.7 ± 0.4	93.5 ± 0.3	84.1	98.9
with the linear schedule  $(\mathrm{VP}_{\mathrm{linear}})$  on  $\beta$ , VP path with cosine schedules  $(\mathrm{VP}_{\mathrm{cos}})$  [33] and polynomial schedules  $(\mathrm{VP}_{\mathrm{poly}})$  [13], and the OT path in Fig. 4.
We observe that the information quantity of  $I(\mathbf{x}_t, \mathbf{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  $\mathbf{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  $\mathbf{x}$  then loses the intrinsic chemical information. Reversely, the dynamics  $I(\mathbf{x}_t, \mathbf{h}_0)$  also implies a generation procedure where the coordinates  $\mathbf{x}$  transform first and the atom types  $\mathbf{h}$  is then determined when  $\mathbf{x}$  are relatively stable.
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.
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
Table 2: Mean Absolute Error for molecular property prediction. A lower number indicates a better controllable generation result.
Property Units	α Bohr3	Δε meV	εHOMO meV	εLUMO meV	μ D	Cv mol K
QM9*	0.10	64	39	36	0.043	0.040
Random*	9.01	1470	645	1457	1.616	6.857
Natoms	3.86	866	426	813	1.053	1.971
EDM	2.76	655	356	584	1.111	1.101
EQUIFM	2.41	591	337	530	1.106	1.033
Table 3: Ablation study, EquiFM models trained with different probability path, the effect of EOT is also evaluated.
Method	Atom Stable (%)	Mol Stable (%)
EquiFMEOT+VPLinear	98.9±0.1	88.3±0.3
EquiFMOT+VPLinear	98.7±0.1	84.9±0.4
EquiFMPV+VPLinear	98.4±0.1	81.6±0.3
EquiFMEOT+VPcos	98.7±0.1	84.7±0.2
EquiFMEOT+VPpoly	98.7±0.1	83.4±0.5
EquiFMEOT+OT	97.3±0.1	77.1±0.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-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  $\alpha$ , orbital energies  $\varepsilon_{\mathrm{HOMO}}$ ,  $\varepsilon_{\mathrm{LUMO}}$  and their gap  $\Delta \varepsilon$ , Dipole moment  $\mu$ , 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  $50K$  samples in each. Then we train a property prediction network  $\omega$  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  $\omega$  to calculate their property values as  $\hat{s}$ . The Mean Absolute Error (MAE) between  $s$  and  $\hat{s}$  is reported to measure whether generated molecules are close to their conditioned property. We also test the MAE of directly running  $\omega$  on the second half QM9, named QM9 in Table 2, which measures the bias of  $\omega$ . 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  $\omega$  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
Figure 3: Resulting performance in molecule stability w.r.t. change of NFEs for each integration algorithm applied in the sampling process. EquiFM is better than our 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 be observed that  $\mathrm{VP}_{\mathrm{linear}}$  path holds the closest tendency to that of  $I(\mathbf{x}_t,\mathbf{h}_0)$ , and OT path has the largest discrepancy.
property information into the generated molecules. And overcoming  $N_{\mathrm{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  $\mathbf{x}$  and the categorical variable  $\mathbf{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  $\mathrm{(VP_{Linear})}$  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  $\mathrm{(VP_{Linear})}$  path(Eq. 5), vanilla optimal transport (OT) (Eq. 4), the variance-preserving path with polynomial decay  $\mathrm{(VP_{Poly})}$ , variance preserving path with cosine schedule  $\mathrm{(VP_{Cos})}$ . 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 Euler 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).
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 \times$  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.
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.
[1] Michael S Albergo and Eric Vanden-Eijnden. Building normalizing flows with stochastic interpolants. arXiv preprint arXiv:2209.15571, 2022.
[2] Namrata Anand and Tudor Achim. Protein structure and sequence generation with equivariant denoising diffusion probabilistic models. arXiv preprint arXiv:2205.15019, 2022.
[3] Brandon Anderson, Truong Son Hy, and Risi Kondor. Cormorant: Covariant molecular neural networks. Advances in 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-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 in neural information processing systems, 32, 2019.
[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 arXiv:2001.09382, 2020.
[47] Jascha Sohl-Dickstein, Eric A Weiss, Niru Maheswaranathan, and Surya Ganguli. Deep unsupervised learning using nonequilibrium thermodynamics. arXiv preprint arXiv:1503.03585, 2015.
[48] Yang Song, Conor Durkan, Iain 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. In Advances in Neural Information Processing Systems, volume 32, pages 11918-11930. Curran Associates, Inc., 2019.
[50] 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.
[51] 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.
[52] 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.
[53] 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.
[54] 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  $\mathbf{x}$  and VP path on the invariant variable  $\mathbf{h}$  as an example is demonstrated in Algorithm 1.
Algorithm 1 Training Algorithm of EquiFM
1: Input: geometric data distribution  $p_{\mathbf{g}}$ ,  $\mathbf{g} = \langle \mathbf{x}, \mathbf{h} \rangle$   
2: Initial: vector field network  $v_{\theta}$ , minimum constant variance  $\sigma_{\mathrm{min}}$   
3: while  $\theta$  have not converged do  
4:  $t \sim \mathbf{U}(0, 1)$ ,  $\epsilon \sim \mathcal{N}(\mathbf{0}, \mathbf{I})$ ,  $g_0 \sim p_{\mathbf{g}}$   
5: Subtract center of gravity from  $\epsilon_x$  in  $\epsilon = [\epsilon_x, \epsilon_h]$   
6: Obtaining the Equivariant Optimal Transport plan  $\pi^*$ ,  $\mathbf{R}^*$  based on Algorithm 3  
7:  $x_t = (\sigma_{\mathrm{min}} + (1 - \sigma_{\mathrm{min}})t)\pi^*(\mathbf{R}^*\epsilon_x) + (1 - t)x_0$  {Eq. 8}  
8:  $h_t = \alpha_th_0 + (1 - \alpha_t^2)\epsilon_h, \alpha_t = e^{-\frac{1}{2}T(t)}$  {Eq. 5}  
9:  $\mathcal{L}_{\mathrm{EquiFM}} = ||v_\theta^x(\langle x_t, h_t\rangle, t) - (-x_0 + (1 - \sigma_{\mathrm{min}})\epsilon_x)||^2 + ||v_\theta^h(\langle x_t, h_t\rangle, t) - \frac{\alpha'_t}{1 - \alpha'_t^2}(\alpha_th_t - h_0)||^2$   
10: end while  
11: return  $v_\theta$
The sampling algorithm can be found in Algorithm 2.