SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Yuxuan Song $^{1,*}$ , Jingjing Gong $^{1,*}$ , Minkai Xu $^{2}$ , Ziyao Cao $^{1}$ , Yanyan Lan $^{1,3}$ , Stefano Ermon $^{2}$ , Hao Zhou $^{1}$ , Wei-Ying Ma $^{1}$
$^{1}$  Institute of AI Industry Research(AIR), Tsinghua University
$^{2}$  Computer Science Department, Stanford University
$^{3}$ Beijing Academy of Artificial Intelligence
{songyuxuan,gongjingjing}@air.tsinghua.edu.cn
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. $^2$
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
*Equal Contribution, Hao Zhou (zhouhao@air.tsinghua.edu.cn) is the corresponding author.
2The code is available at https://github.com/AlgoMole/MolFM
Figure 1: Illustration of EquiFM. We define a hybrid path for generating molecules  $\mathbf{g} = \langle \mathbf{x},\mathbf{h}\rangle$  where  $\mathbf{x}$  is trained on an equivariant optimal transport path and  $\mathbf{h}$  is trained on a path whose information quantity is aligned with  $\mathbf{x}$ 's path. The sampling is conducted by solving an ODE, i.e.  $\mathbf{g}_0 = \mathrm{ODESolve}(\mathbf{g}_1,v_\theta ,1,0)$ .
DM-based methods typically suffer from unstable probability dynamics which could lead to an inefficient sampling speed also limit the validity rate of generated molecules.
In this work, we propose a novel and principled flow-matching objective, termed Equivariant Flow-Matching (EquiFM), for geometric generative modeling. Our method is inspired by the recent advancement of flow matching [26], a simulation-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 stabilizing 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 \times$  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.
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 structure-based 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\in [0,1]}:\mathbb{R}^d\to \mathbb{R}_{>0}$ , and the time-dependent vector field is defined as  $v_{t\in [0,1]}\colon \mathbb{R}^d\to \mathbb{R}^d$ . The vector field uniquely defines time-dependent flow  $\psi_{t\in [0,1]}\colon \mathbb{R}^d\to \mathbb{R}^d$  by the following ordinary differential equation (ODE):
$$
\frac {\mathrm {d}}{\mathrm {d} t} \psi_ {t} (x) = v _ {t} \left(\psi_ {t} (x)\right), \psi_ {1} (x) = x \tag {1}
$$
[4] proposed to train the parameterized flow model  $\psi_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_{\theta}(x,t)$  to some target vector field  $u_t(x)$ :
$$
\mathcal {L} _ {\mathrm {F M}} (\theta) = \mathbb {E} _ {t, p _ {t} (x)} \| v _ {\theta} (x, t) - u _ {t} (x) \| ^ {2} \tag {2}
$$
The objective  $\mathcal{L}_{\mathrm{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\mid x_{0})$  and the corresponding conditional probability path  $p_t(x\mid x_0)$  are readily definable.
The probability path can be marginalized from a mixture of conditional probability path  $p_t(x) = \int p_t(x \mid x_0) q(x_0) \, \mathrm{d}x_0$ , and the vector field  $u_t(x)$  can be marginalized from conditional vector field as  $u_t(x) = \mathbb{E}_{x_0 \sim q} \frac{u_t(x \mid x_0) p_t(x \mid 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 \mid x_0)$  generating the conditional probability path  $p_t(x \mid 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:
$$
\mathcal {L} _ {\mathrm {C F M}} (\theta) = \mathbb {E} _ {t, q \left(x _ {0}\right), p _ {t} \left(x \mid x _ {0}\right)} \| v _ {\theta} (x, t) - u _ {t} (x \mid x _ {0}) \| _ {2} ^ {2} \tag {3}
$$
The CFM objective enjoys the tractability for optimization, and optimizing the CFM objective is equivalent to optimizing Eq. 2, i.e.,  $\nabla_{\theta}\mathcal{L}_{\mathrm{FM}}(\theta) = \nabla_{\theta}\mathcal{L}_{\mathrm{CFM}}(\theta)$ . For the inference phase, ODE solvers could be applied to solve the Eq. 1, e.g.,  $x_0 = \mathrm{ODESolve}(x_1,v_\theta ,1,0)$ . In this paper, we consider using the Gaussian conditional probability path, which lies in the form of  $p_t(x\mid x_0) = \mathcal{N}\left(x\mid \mu_t(x_0),\sigma_t(x_0)^2 I\right)$ . 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 \mid x_0)$  is approximated with a peaked Gaussian centered in  $x_0$  with a small variance  $\sigma_{\mathrm{min}}$  as  $\mathcal{N}\left(x \mid x_0, \sigma_{\mathrm{min}}^2 I\right)$ . The probability path is  $p_t(x \mid x_0) = \mathcal{N}\left(x \mid (1 - t)x_0, (\sigma_{\mathrm{min}} + (1 - \sigma_{\mathrm{min}})t)^2 I\right)$  and the corresponding flow is  $\psi_t(x) = (\sigma_{\mathrm{min}} + (1 - \sigma_{\mathrm{min}})t)x + (1 - t)x_0$ . Then the vector field could be obtained by Eq. 1 as:  $u_t(\psi_t(x) \mid x_0) = \frac{\mathrm{d}}{\mathrm{d}t} \psi_t(x) = -x_0 + (1 - \sigma_{\mathrm{min}})x$ .
Put the above terms into Eq. 3, the reparameterized objective is as:
$$
\mathcal {L} _ {\mathrm {C F M}} ^ {\mathrm {O T}} (\theta) = \mathbb {E} _ {t, q \left(x _ {0}\right), p _ {1} \left(x _ {1}\right)} \left\| v _ {\theta} \left(\psi_ {t} \left(x _ {1}\right), t\right) - \left(- x _ {0} + \left(1 - \sigma_ {\min }\right) x _ {1}\right) \right\| ^ {2} \tag {4}
$$
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 \mid x_0) = \mathcal{N}(x \mid \alpha_t x_0, (1 - \alpha_t^2) I)$ , where  $\alpha_t = e^{-\frac{1}{2} T(t)}$  and  $T(t) = \int_0^t \beta(s) ds$ . Here  $\beta$  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 \mid x_0) = \frac{\alpha_t'}{1 - \alpha_t^2} (\alpha_t x - x_0)$ .  $\alpha_t'$  denotes the derivative with respect to time. And the objective for VP conditional flow matching is as:
$$
\mathcal {L} _ {\mathrm {C F M}} ^ {\mathrm {V P}} (\theta) = \mathbb {E} _ {t, q (x _ {0}), p _ {t} (x | x _ {0})} \left\| v _ {\theta} (x, t) - \frac {\alpha_ {t} ^ {\prime}}{1 - \alpha_ {t} ^ {2}} \left(\alpha_ {t} x - x _ {0}\right) \right\| ^ {2} \tag {5}
$$
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.
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  $\mathbf{g} = \langle \mathbf{x},\mathbf{h}\rangle$  , where  $\mathbf{x} = (\mathbf{x}^{1},\dots ,\mathbf{x}^{N})\in X$  is the atom coordinates matrix and  $\mathbf{h} = (\mathbf{h}^1,\ldots ,\mathbf{h}^N)\in \mathbb{R}^{N\times d}$  is the node feature matrix, such as atomic type and charges. Here  $X = \left\{\mathbf{x}\in \mathbb{R}^{N\times 3}: \frac{1}{N}\sum_{i = 1}^{N}\mathbf{x}^i = \mathbf{0}\right\}$  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  $\mathbf{x}$  invariant to the translation by setting the prior distribution and vector field in the Zero CoM space, i.e.  $\frac{1}{N} \sum_{i=1}^{N} 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_{\theta}^{\mathbf{x}}(\mathbf{g},t),v_{\theta}^{\mathbf{h}}(\mathbf{g},t)) = v_{\theta}(\mathbf{g},t)$ , where  $v_{\theta}^{\mathbf{x}}(\mathbf{g},t)$  and  $v_{\theta}^{\mathbf{h}}(\mathbf{g},t)$  are the parameterized vector field for  $\mathbf{x}$  and  $\mathbf{h}$ . If the vector field is equivariant to any rotational transformation  $\mathbf{R}$ , i.e.,  $v_{\theta}(\langle \mathbf{Rx},\mathbf{h}\rangle ,t) = (\mathbf{R}(v_{\theta}^{\mathbf{x}}(\mathbf{g},t)),v_{\theta}^{\mathbf{h}}(\mathbf{g},t))$ . With an rotational invariant prior function  $p_1(\mathbf{x},\mathbf{h})$ , i.e.,  $p_1(\mathbf{Rx},\mathbf{h}) = p_1(\mathbf{x},\mathbf{h})$ , then the probability path  $p_{\theta ,t}$  generated by the vector field  $v_{\theta}(\cdot)$  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
Figure 2: Generation route visualization of different models. Note that a lighter color indicates an earlier step of an atom and a denser color corresponds to a later step. A change of base color indicates a change of atom type. EquiFM generates molecules in a straightforward route as shown in 2(a). Vanilla flow matching method 2(b) on the other hand, takes a detour while generating molecules, resulting in a route inward then outward before converging to a molecule. The generation process in EDM 2(c) is rather chaotic until the last few steps before converging into a molecule.
(c) EDM
Hybrid Probability Modeling We refer to the target conditional vector field on each part as  $u_{t}^{\mathbf{x}}(\mathbf{g} \mid \mathbf{g}_{0})$  and  $u_{t}^{\mathbf{h}}(\mathbf{g} \mid \mathbf{g}_{0})$  correspondingly, then we could get the objective in the following formulation:
$$
\mathcal {L} _ {\mathrm {C F M}} (\theta) = \mathbb {E} _ {t, q \left(\mathbf {g} _ {0}\right), p _ {t} \left(\mathbf {g} \mid \mathbf {g} _ {0}\right)} \left[ \| v _ {\theta} ^ {\mathbf {x}} (\mathbf {g}, t) - u _ {t} ^ {\mathbf {x}} (\mathbf {g} \mid \mathbf {g} _ {0}) \| _ {2} ^ {2} + \left\| v _ {\theta} ^ {\mathbf {h}} (\mathbf {g}, t) - u _ {t} ^ {\mathbf {h}} (\mathbf {g} \mid \mathbf {g} _ {0}) \right\| _ {2} ^ {2} \right] \tag {6}
$$
Proposition 4.2. There could be joint probability path  $p_t(\mathbf{g}|\mathbf{g}_0)$  which satisfies that  $p_t(\mathbf{g}|\mathbf{g}_0) = p_t(\mathbf{x}|\mathbf{x}_0)p_t(\mathbf{h}|\mathbf{h}_0)$ , and the conditional vector field on  $\mathbf{x}$  and  $\mathbf{h}$  is independent:  $u_t^{\mathbf{x}}(\mathbf{g}|\mathbf{g}_0) = u_t(\mathbf{x}|\mathbf{x}_0), u_t^{\mathbf{h}}(\mathbf{g}|\mathbf{g}_0) = u_t(\mathbf{h}|\mathbf{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  $\mathbf{x}$  and  $\mathbf{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  $\mathbf{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  $\mathbf{x}_0$  as the point cloud from molecule distribution and  $\mathbf{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,  $\mathbf{z} = (\mathbf{z}^1,\dots ,\mathbf{z}^N)\in \mathbb{R}^{N\times 3}$  and  $\mathbf{y} = (\mathbf{y}^1,\dots ,\mathbf{y}^N)\in \mathbb{R}^{N\times 3}$ . We define the equivariant optimal transport plan as
$$
\pi^ {*}, \mathbf {R} ^ {*} = \underset {\pi , \mathbf {R}} {\operatorname {a r g m i n}} \| \pi \left(\mathbf {R} \mathbf {z} ^ {1}, \mathbf {R} \mathbf {z} ^ {2}, \dots , \mathbf {R} \mathbf {z} ^ {N}\right) - \left(\mathbf {y} ^ {1}, \mathbf {y} ^ {2}, \dots , \mathbf {y} ^ {N}\right) \| _ {2} \tag {7}
$$
Here  $\pi$  is a permutation of  $N$  elements and  $\mathbf{R} \in \mathbb{R}^{3 \times 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:
$$
p _ {t} = \left[ \psi_ {t} ^ {\mathrm {E O T}} \right] _ {*} p _ {1}, \text {w h e r e} \psi_ {t} ^ {\mathrm {E O T}} (\mathbf {x}) = \left(\sigma_ {\min } + (1 - \sigma_ {\min }) t\right) \pi^ {*} \left(\mathbf {R} ^ {*} \mathbf {x}\right) + (1 - t) \mathbf {x} _ {0} \tag {8}
$$
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:
$$
\mathcal {L} _ {\mathrm {C F M}} ^ {\mathrm {E O T}} (\theta) = \mathbb {E} _ {t, q (\mathbf {x} _ {0}), p _ {1} (\mathbf {x} _ {1})} \left\| v _ {\theta} \left(\psi_ {t} ^ {\mathrm {E O T}} (\mathbf {x} _ {1}), t\right) - (- \mathbf {x} _ {0} + (1 - \sigma_ {\min }) \pi^ {*} (\mathbf {R} ^ {*} \mathbf {x} _ {1})) \right\| ^ {2} \tag {9}
$$
A good property of the objective with EOT is that the training characteristics are invariant to translation and rotation of initial  $\mathbf{x}_1$ , and equivariant with respect to both sampled noise  $\mathbf{x}_1$  and data point  $\mathbf{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