SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Leon Klein
leon.klein@fu-berlin.de
andreas.kraemer@fu-berlin.de
Frank Noé
franknoe@microsoft.com
Normalizing flows are a class of deep generative models that are especially interesting for modeling probability distributions in physics, where the exact likelihood of flows allows reweighting to known target energy functions and computing unbiased observables. For instance, Boltzmann generators tackle the long-standing sampling problem in statistical physics by training flows to produce equilibrium samples of many-body systems such as small molecules and proteins. To build effective models for such systems, it is crucial to incorporate the symmetries of the target energy into the model, which can be achieved by equivariant continuous normalizing flows (CNFs). However, CNFs can be computationally expensive to train and generate samples from, which has hampered their scalability and practical application. In this paper, we introduce equivariant flow matching, a new training objective for equivariant CNFs that is based on the recently proposed optimal transport flow matching. Equivariant flow matching exploits the physical symmetries of the target energy for efficient, simulation-free training of equivariant CNFs. We demonstrate the effectiveness of flow matching on rotation and permutation invariant many-particle systems and a small molecule, alanine dipeptide, where for the first time we obtain a Boltzmann generator with significant sampling efficiency without relying on tailored internal coordinate featurization. Our results show that the equivariant flow matching objective yields flows with shorter integration paths, improved sampling efficiency, and higher scalability compared to existing methods.
Generative models have achieved remarkable success across various domains, including images [1, 2, 3, 4], language [5, 6, 7], and applications in the physical sciences [8, 9, 10]. Among the rapidly growing subfields in generative modeling, normalizing flows have garnered significant interest. Normalizing flows [11, 12, 13, 14] are powerful exact-likelihood generative neural networks that transform samples from a simple prior distribution into samples that follow a desired target probability distribution. Previous studies [15, 16, 17, 18, 19] have emphasized the importance of incorporating symmetries of the target system into the flow model. In this work, we focus specifically on many-body systems characterized by configurations  $x \in \mathbb{R}^{N \times D}$  of  $N$  particles in  $D$  spatial dimensions. The symmetries of such systems arise from the invariances of the potential energy function  $U(x)$ . The
Figure 1: Results for the 13 particle Lennard-Jones system (LJ13) for different flow matching training methods. (a, b) Sample pairs generated with the different flow matching objectives during training. (c) Mean transport cost (squared distance) for training batches. (d, e) Integration paths per particle for samples generated by models trained with OT flow matching and equivariant OT flow matching, respectively. (f) Integration path length, i.e. arc length between prior and output sample, distribution compared with the OT path length between the prior and push-forward distribution of the flows.
associated probability distribution, known as the Boltzmann distribution, is given by
$$
\mu (x) \propto \exp \left(- \frac {U (x)}{k _ {B} T}\right), \tag {1}
$$
where  $k_{B}$  represents the Boltzmann constant and  $T$  denotes the temperature.
Generating equilibrium samples from the Boltzmann distribution is a long-standing problem in statistical physics, typically addressed using iterative methods like Markov Chain Monte Carlo or Molecular Dynamics (MD). In contrast, Boltzmann generators (BGs) [8] utilize normalizing flows to directly generate independent samples from the Boltzmann distribution. Moreover, they allow the reweighting of the generated density to match the unbiased target Boltzmann distribution  $\mu$ . When dealing with symmetric densities, which are ubiquitous in physical systems, BGs employ two main approaches: (i) describing the system using internal coordinates [8, 20] or (ii) describing the system in Cartesian coordinates while incorporating the symmetries into the flow through equivariant models [17]. In this work, we focus on the latter approach, which appears more promising as such architectures can in principle generalize across different molecules and do not rely on tailored system-specific featurizations. However, current equivariant Boltzmann generators based on continuous normalizing flows (CNFs) face limitations in scalability due to their high computational cost during both training and inference. Recently, flow matching [21] has emerged as a fast, simulation-free method for training CNFs. A most recent extension, optimal transport (OT) flow matching [22], enables learning optimal transport maps, which facilitates fast inference through simple integration paths. In this work, we apply OT flow matching to train flows for highly symmetric densities and find that the required batch size to approximate the OT map adequately can become prohibitively large (Figure 1c). Thus, the learned flow paths will deviate strongly from the OT paths (Figure 1d), which increases computational cost and numerical errors. We tackle this problem by proposing a novel equivariant OT flow matching objective for training equivariant flows on invariant densities, resulting in optimal paths for faster inference (see Figure 1e,f).
Our main contributions in this work are as follows:
1. We propose a novel flow matching objective designed for invariant densities, yielding nearly optimal integration paths. Concurrent work of [23] proposes nearly the same objective with the same approximation method, which they also call equivariant flow matching.
2. We compare different flow matching objectives to train equivariant continuous normalizing flows. Through our evaluation, we demonstrate that only our proposed equivariant flow matching objective enables close approximations to the optimal transport paths for invariant densities. However, our proposed method is most effective for larger highly symmetric systems, while achieving sometimes inferior results for the smaller systems compared to optimal transport flow matching.
3. We introduce a new invariant dataset of alanine dipeptide and a large Lennard-Jones cluster. These datasets serve as valuable resources for evaluating the performance of flow models on invariant systems.
4. We present the first Boltzmann Generator capable of producing samples from the equilibrium Boltzmann distribution of a molecule in Cartesian coordinates. Additionally, we demonstrate the reliability of our generator by accurately estimating the free energy difference, in close agreement with the results obtained from umbrella sampling. Concurrent work of [24] also introduce a Boltzmann Generator, based on coupling flows instead of CNFs, in Cartesian coordinates for alanine dipeptide. However, they investigate alanine dipeptide at  $T = 800K$  instead of room temperature.
Normalizing flows and Boltzmann generators [8] have been applied to molecular sampling and free energy estimation [25, 20, 26, 27, 28]. Previous flows for molecules only achieved significant sampling efficiency when employing either system-specific featurization such as internal coordinates [8, 29, 25, 20, 30, 31] or prior distributions close to the target distribution [32, 31, 33]. Notably, our equivariant OT flow matching method could also be applied in such scenarios, where the prior distribution is sampled by MD at a different thermodynamic state (e.g., higher temperature or lower level of theory). A flow model for molecules in Cartesian coordinates has been developed in [34] where a transferable coupling flow is used to sample small peptide conformations by proposing iteratively large time steps instead of sampling from the target distribution directly. Equivariant diffusion models [35, 36] learn a score-based model to generate molecular conformations. The score-based model is parameterized by an equivariant function similar to the vector field used in equivariant CNFs. However, they do not target the Boltzmann distribution. As an exception, [37] propose a diffusion model in torsion space and use the underlying probability flow ODE as a Boltzmann generator. Moreover, [38] use score-based models to learn the transition probability for multiple time-resolutions, accurately capturing the dynamics.
To help speed up CNF training and inference, various authors [39, 40, 41] have proposed incorporating regularization terms into the likelihood training objective for CNFs to learn an optimal transport (OT) map. While this approach can yield flows similar to OT flow matching, the training process itself is computationally expensive, posing limitations on its scalability. The general idea underlying flow matching was independently conceived by different groups [21, 42, 43] and soon extended to incorporate OT [22, 44]. OT maps with invariances have been studied previously to map between learned representations [45], but have not yet been applied to generative models. Finally, we clarify that we use the term flow matching exclusively to refer to the training method for flows introduced by Lipman et al. [21] rather than the flow-matching method for training energy-based coarse-grained models from flows introduced at the same time [46].
In this section, we describe the key methodologies used in our study, highlighting important prior work.
3.1 Normalizing flows
Normalizing flows [14, 47] provide a powerful framework for learning complex probability densities  $\mu(x)$  by leveraging the concept of invertible transformations. These transformations, denoted as  $f_{\theta}: \mathbb{R}^n \to \mathbb{R}^n$ , map samples from a simple prior distribution  $q(x_0) = \mathcal{N}(x_0 | 0, I)$  to samples from a more complicated output distribution. This resulting distribution  $p(x_1)$ , known as the push-forward
distribution, is obtained using the change of variable equation:
$$
p \left(x _ {1}\right) = q \left(f _ {\theta} ^ {- 1} \left(x _ {1}\right)\right) \det  \left| \frac {\partial f _ {\theta} ^ {- 1} \left(x _ {1}\right)}{\partial x _ {1}} \right|, \tag {2}
$$
where  $\frac{\partial f_{\theta}^{-1}(x_1)}{\partial x_1}$  denotes the Jacobian of the inverse transformation  $f_{\theta}^{-1}$ .
3.2 Continuous normalizing flows (CNFs)
CNFs [48, 49] are an instance of normalizing flows that are defined via the ordinary differential equation
$$
\frac {d f _ {\theta} ^ {t} (x)}{d t} = v _ {\theta} \left(t, f _ {\theta} ^ {t} (x)\right), \quad f _ {\theta} ^ {0} (x) = x, \tag {3}
$$
where  $v_{\theta}(t,x):\mathbb{R}^n\times [0,1]\to \mathbb{R}^n$  is a time dependent vector field. The time dependent flow  $f_{\theta}^{t}(x)$  generates the probability flow path  $\tilde{p}_t(x):[0,1]\times \mathbb{R}^n\to \mathbb{R}^+$  from the prior distribution  $\tilde{p}_0(x) = q(x_0)$  to the push-forward distribution  $\tilde{p}_1(x)$ . We can generate samples from  $\tilde{p}_t(x)$  by first sampling  $x_0$  from the prior and then integrating
$$
f _ {\theta} ^ {t} (x) = x _ {0} + \int_ {0} ^ {t} d t ^ {\prime} v _ {\theta} \left(t ^ {\prime}, f _ {\theta} ^ {t ^ {\prime}} (x)\right). \tag {4}
$$
The corresponding log density change is given by the continuous change of variable equation
$$
\log \tilde {p} _ {t} (x) = \log q \left(x _ {0}\right) - \int_ {0} ^ {t} d t ^ {\prime} \nabla \cdot v _ {\theta} \left(t ^ {\prime}, f _ {\theta} ^ {t ^ {\prime}} (x)\right), \tag {5}
$$
where  $\nabla \cdot v_{\theta}\left(t^{\prime},f_{\theta}^{t^{\prime}}(x)\right)$  is the trace of the Jacobian.
3.3 Likelihood-based training
CNFs can be trained using a likelihood-based approach. In this training method, the loss function is formulated as the negative log-likelihood
$$
\mathcal {L} _ {\mathrm {N L L}} (\theta) = \mathbb {E} _ {x _ {1} \sim \mu} \left[ - \log q \left(f _ {\theta} ^ {- 1} \left(x _ {1}\right)\right) - \log \left| \det  \frac {\partial f _ {\theta} ^ {- 1} \left(x _ {1}\right)}{\partial x _ {1}} \right| \right]. \tag {6}
$$
Evaluating the inverse map requires propagating the samples through the reverse ODE, which involves numerous evaluations of the vector field for integration in each batch. Additionally, evaluating the change in log probability involves computing the expensive trace of the Jacobian.
3.4 Flow matching
In contrast to likelihood-based training, flow matching allows training CNFs simulation-free, i.e. without integrating the vector field and evaluating the Jacobian, making the training significantly faster (see Appendix A.2). Hence, the flow matching enables scaling to larger models and larger systems with the same computational budget. The flow matching training objective [21, 22] regresses  $v_{\theta}(t,x)$  to some target vector field  $u_{t}(x)$  via
$$
\mathcal {L} _ {\mathrm {F M}} (\theta) = \mathbb {E} _ {t \sim [ 0, 1 ], x \sim p _ {t} (x)} | | v _ {\theta} (t, x) - u _ {t} (x) | | _ {2} ^ {2}. \tag {7}
$$
In practice,  $p_t(x)$  and  $u_t(x)$  are intractable, but it has been shown in [21] that the same gradients can be obtained by using the conditional flow matching loss:
$$
\mathcal {L} _ {\mathrm {C F M}} (\theta) = \mathbb {E} _ {t \sim [ 0, 1 ], x \sim p _ {t} (x | z)} \left| \left| v _ {\theta} (t, x) - u _ {t} (x | z) \right| \right| _ {2} ^ {2}. \tag {8}
$$
The vector field  $u_{t}(x)$  and corresponding probability path  $p_t(x)$  are given in terms of the conditional vector field  $u_{t}(x|z)$  as well as the conditional probability path  $p_t(x|z)$  as
$$
p _ {t} (x) = \int p _ {t} (x | z) p (z) d z \quad \text {a n d} \quad u _ {t} (x) = \int \frac {p _ {t} (x | z) u _ {t} (x | z)}{p _ {t} (x)} p (z) d z, \tag {9}
$$
where  $p(z)$  is some arbitrary conditioning distribution independent of  $x$  and  $t$ . For a derivation see [21] and [22]. There are different ways to efficiently parametrize  $u_{t}(x|z)$  and  $p_t(x|z)$ . We here focus on a parametrization that gives rise to the optimal transport path, as introduced in [22]
$$
z = \left(x _ {0}, x _ {1}\right) \quad \text {a n d} \quad p (z) = \pi \left(x _ {0}, x _ {1}\right) \tag {10}
$$
$$
u _ {t} (x \mid z) = x _ {1} - x _ {0} \quad \text {a n d} \quad p _ {t} (x \mid z) = \mathcal {N} (x \mid t \cdot x _ {1} + (1 - t) \cdot x _ {0}, \sigma^ {2}), \tag {11}
$$
where the conditioning distribution is given by the 2-Wasserstein optimal transport map  $\pi (x_0,x_1)$  between the prior  $q(x_0)$  and the target  $\mu (x_{1})$ . The 2-Wasserstein optimal transport map is defined by the 2-Wasserstein distance
$$
W _ {2} ^ {2} = \inf  _ {\pi \in \Pi} \int c \left(x _ {0}, x _ {1}\right) \pi \left(d x _ {0}, d x _ {1}\right), \tag {12}
$$
where  $\Pi$  is the set of couplings as usual and  $c(x_0,x_1) = ||x_0 - x_1||_2^2$  is the squared Euclidean distance in our case. Following [22], we approximate  $\pi (x_0,x_1)$  by only considering a batch of prior and target samples. Hence, for each batch we generate samples from  $p(z)$  as follows:
1. sample batches of points  $(x_0^1,\ldots ,x_0^B)\sim q$  and  $(x_{1}^{1},\dots,x_{1}^{B})\sim \mu$
2. compute the cost matrix  $M$  for the batch, i.e.  $M_{ij} = ||x_0^i - x_1^j||_2^2$ ,
3. solve the discrete OT problem defined by  $M$
4. generate training pairs  $z^i = (x_0^i, x_1^i)$  according to the OT solution.
We will refer to this training procedure as OT flow matching.
3.5 Equivariant flows
Symmetries can be described in terms of a group  $G$  acting on a finite-dimensional vector space  $V$  via a matrix representation  $\rho(g); g \in G$ . A map  $I: V \to V'$  is called  $G$ -invariant if  $I(\rho(g)x) = I(x)$  for all  $g \in G$  and  $x \in V$ . Similarly, a map  $f: V \to V$  is called  $G$ -equivariant if  $f(\rho(g)x) = \rho(g)f(x)$  for all  $g \in G$  and  $x \in V$ .
In this work, we focus on systems with energies  $U(x)$  that exhibit invariance under the following symmetries: (i) Permutations of interchangeable particles, described by the symmetric group  $S(N')$  for each interchangeable particle group. (ii) Global rotations and reflections, described by the orthogonal group  $O(D)$  (iii) Global translations, described by the translation group  $\mathbb{T}$ . These symmetries are commonly observed in many particle systems, such as molecules or materials.
In [17], it is shown that equivariant CNFs can be constructed using an equivariant vector field  $v_{\theta}$  (Theorem 2). Moreover, [17, 50] show that the push-forward distribution  $\tilde{p}(x_1)$  of a  $G$ -equivariant flow with a  $G$ -invariant prior density is also  $G$ -invariant (Theorem 1). Note that this is only valid for orthogonal maps, and hence not for translations. However, translation invariance can easily be achieved by assuming mean-free systems as proposed in [17]. An additional advantage of mean-free systems is that global rotations are constrained to occur around the origin. By ensuring that the flow does not modify the geometric center and the prior distribution is mean-free, the resulting push-forward distribution will also be mean-free.
Although equivariant flows can be successfully trained with the OT flow matching objective, the trained flows display highly curved vector fields (Figure 1d) that do not match the linear OT paths. Fortunately, the OT training objective can be modified so that it yields better approximations to the OT solution for finite batch sizes.
4 Equivariant optimal transport flow matching
Prior studies [51, 22] indicate that medium to small batch sizes are often sufficient for effective OT flow matching. However, when dealing with highly symmetric densities, accurately approximating the OT map may require a prohibitively large batch size. This is particularly evident in cases involving permutations, where even for small system sizes, it is unlikely for any pair of target-prior samples  $(x_0, x_1)$  to share the exact same permutation. This is because the number of possible permutations scales with the number of interchangeable particles as  $N!$ , while the number of combinations of the sample pairs scales only with squared batch size (Figure 1). To address this challenge, we propose using a cost function
$$
\tilde {c} \left(x _ {0}, x _ {1}\right) = \min  _ {g \in G} \left| \left| x _ {0} - \rho (g) x _ {1} \right| \right| _ {2} ^ {2}, \tag {13}
$$
that accurately accounts for the underlying symmetries of the problem in the OT flow matching algorithm. Hence, instead of solely reordering the batch, we instead also align samples along their orbits.
We summarize or main theoretical findings in the following theorem.
Theorem 1. Let  $G$  be a compact group that acts on an Euclidean  $n$ -space by isometries. Let  $T: x \mapsto y$  be an OT map between  $G$ -invariant measures  $\nu_{1}$  and  $\nu_{2}$ , using the cost function  $c$ . Then
1.  $T$  is  $G$ -equivariant and the corresponding OT plan  $\pi(\nu_1, \nu_2)$  is  $G$ -invariant.
2. For all pairs  $(x, T(x))$  and  $y \in G \cdot T(x)$ :
$$
c (x, T (x)) = \int_ {G} c (g \cdot x, g \cdot T (x)) d \mu (g) = \min  _ {g \in G} c (x, g \cdot y) \tag {14}
$$
3.  $T$  is also an OT map for the cost function  $\tilde{c}$ .
Refer to Appendix B.1 for an extensive derivation and discussion. The key insights can be summarized as follows: (i) Given the  $G$ -equivariance of the target OT map  $T$ , it is natural to employ an  $G$ -equivariant flow model for its learning. (ii) From 2. and 3., we can follow that our proposed cost function  $\tilde{c}$ , effectively aligns pairs of samples in a manner consistent with how they are aligned under the  $G$ -equivariant OT map  $T$ .
In this work we focus on  $O(D)$ - and  $S(N)$ -invariant distributions, which are common for molecules and multi-particle systems. The minimal squared distance for a pair of points  $(x_0, x_1)$ , taking into account these symmetries, can be obtained by minimizing the squared Euclidean distance over all possible combinations of rotations, reflections, and permutations
$$
\tilde {c} \left(x _ {0}, x _ {1}\right) = \min  _ {r \in O (D), s \in S (N)} \left\| x _ {0} - \rho (r s) x _ {1} \right\| _ {2} ^ {2}. \tag {15}
$$
However, computing the exact minimal squared distance is computationally infeasible in practice due to the need to search over all possible combinations of rotations and permutations. Therefore, we approximate the minimizer by performing a sequential search
$$
\tilde {c} \left(x _ {0}, x _ {1}\right) = \min  _ {r \in S O (D)} | | x _ {0} - \rho (r \tilde {s}) x _ {1} | | _ {2} ^ {2}, \quad \tilde {s} = \underset {s \in S (N)} {\arg \min } | | x _ {0} - \rho (s) x _ {1} | | _ {2} ^ {2}. \tag {16}
$$
We demonstrate in Section 6 that this approximation results in nearly OT integration paths for equivariant flows, even for small batch sizes. While we also tested other approximation strategies in Appendix A.9, they did not yield significant changes in our results, but come at with additional computational overhead. We hence alter the OT flow matching procedure as follows: For each element of the cost matrix  $M$ , we first compute the optimal permutation with the Hungarian algorithm [52] and then align the two samples through rotation with the Kabsch algorithm [53]. The other steps of the OT flow matching algorithm remain unchanged. We will refer to this loss as Equivariant OT flow matching. Although aligning samples in that way comes with a significant computational overhead, this reordering can be done in parallel before or during training (see Appendix C.4). It is worth noting that in the limit of infinite batch sizes, the Euclidean cost will still yield the correct OT map for invariant densities.
5 Architecture
The vector field  $v_{\theta}(t,x)$  is parametrized by an  $O(D)$ - and  $S(N)$ -equivariant graph neural network, similar to the one used in [18] and introduced in [54]. The graph neural network consists of  $L$  consecutive layers. The update for the  $i$ -th particle is computed as follows
$$
h _ {i} ^ {0} = \left(t, a _ {i}\right), \quad m _ {i j} ^ {l} = \phi_ {e} \left(h _ {i} ^ {l}, h _ {j} ^ {l}, d _ {i j} ^ {2}\right), \tag {17}
$$
$$
x _ {i} ^ {l + 1} = x _ {i} ^ {l} + \sum_ {j \neq i} \frac {\left(x _ {i} ^ {l} - x _ {j} ^ {l}\right)}{d _ {i j} + 1} \phi_ {d} \left(m _ {i j} ^ {l}\right), \tag {18}
$$
$$
h _ {i} ^ {l + 1} = \phi_ {h} \left(h _ {i} ^ {l}, m _ {i} ^ {l}\right), \quad m _ {i} ^ {l} = \sum_ {j \neq i} \phi_ {m} \left(m _ {i j} ^ {l}\right) m _ {i j} ^ {l}, \tag {19}
$$
$$
v _ {\theta} \left(t, x ^ {0}\right) _ {i} = x _ {i} ^ {L} - x _ {i} ^ {0}, \tag {20}
$$
Table 1: Comparison of flows trained with different training objectives. Errors are computed over three runs.
Training type	NLL (↓)	ESS (↑)	Path length (↓)
Lacklhood [18]	DW4
1.72 ± 0.01	86.87 ± 0.19%	3.11 ± 0.04
OT flow matching	1.70 ± 0.02	92.37 ± 0.89%	2.94 ± 0.02
Equivariant OT flow matching	1.68 ± 0.01	88.71 ± 0.40%	2.92 ± 0.01
LJ13
Lacklhood [18]	-15.83 ± 0.07	39.78 ± 6.19%	5.08 ± 0.22
OT flow matching	-16.09 ± 0.03	54.36 ± 5.43%	2.84 ± 0.01
Equivariant OT flow matching	-16.07 ± 0.02	57.98 ± 2.19%	2.15 ± 0.01
LJ55
OT flow matching	-88.45 ± 0.04	3.74 ± 1.06%	7.53 ± 0.02
Equivariant OT flow matching	-89.27 ± 0.04	4.42 ± 0.35%	3.52 ± 0.01
Alanine dipeptide
OT flow matching	-107.54 ± 0.03	1.41 ± 0.04%	10.19 ± 0.03
Equivariant OT flow matching	-106.78 ± 0.02	0.69 ± 0.05%	9.46 ± 0.01
where  $\phi_{\alpha}$  are neural networks,  $d_{ij}$  is the Euclidean distance between particle  $i$  and  $j$ , and  $a_{i}$  is an embedding for the particle type. Notably, the update conserves the geometric center if all particles are of the same type (see Appendix B.2), otherwise we subtract the geometric center after the last layer. This ensures that the resulting equivariant vector field  $v_{\theta}(t,x)$  conserves the geometric center. When combined with a symmetric mean-free prior distribution, the push-forward distribution of the CNF will be  $O(D)$ - and  $S(N)$ -invariant.