Equivariant flow matching
Authors: Leon Klein
Andreas Krämer
Frank Noe
Venue: NeurIPS 2023 poster (2023)
Keywords: Normalizing Flows
Flow Matching
Equivariance
Boltzmann Generators
Molecular Dynamics
Optimal Transport
TL;DR: We propose a novel flow matching loss for training equivariant continuous normalizing flows

Abstract: 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.

Equivariant flow matching
Leon Klein
Andreas Krämer
Freie
Universität Berlin
Freie Universität Berlin andreas. krramerr@uuberllli
leon. cein@fu-berrin...
Frank Noé
Microsoft Research A14Science
Freie Universität Berlin
Rice University franknoe@@microooff..oo
Abstract
Normalizing flows are class of deep generative models that are especially inter- esting for modeling probability distributions in physics, where the exact likelihood of flows allows reweighting to known target energy functions and computing un- biased observables. For instance, Boltzmann generators tackle the Long-standing sampling problem in statistical physies by training flows to produce equilibrium samples of many-body systems such small molecules and proteins. To build effective models for such systems, it is crucial to incorporate the symmetries 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. this paper, we introduce equivariant flow matching, a new training objective for equivariant CNFs that is based the recently proposed optimal trans- port 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
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 existing methods.
Introduction
Generative models have achieved remarkable success across various domains, including images 4. language 6 71, and applications in the physical sciences [8] [0]. Among the rapidly growing subfields in generative modeling, normalizing flows have garnered significant interest.
Normalizing flows [
12
3 [4] are powerful exact-likelihood generative neural networks that
transform samples from a simple prior distribution into samples that follow desired target probability distribution. Previous studies
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 æ E NDD of N particles in 1) spatial dimensions. The symmetries of such systems arise from the invariances of the potential energy function (xx). The
15 16
37th Conference on Neural Information Processing Systems (NcurIPS 2023).

---

OT flow matching
Equivariant OT flww matching
Mean batch transport cost
b)
Mcan cost
Training
poth target
250
300 Batch size
406
d)
U13 Sampled path lengths
gr
Sampling
Iatent path
sample
Path lengthh
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.c. are 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
(- ((²)
m(x) exp
kgT (-ker)
where kg represents the Boltzmann constant and T' denotes the temperature.
Generating equilibrium samples from the Boltzmann distribution is 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 μ. When dealing with symmetric densities, which are ubiquitous 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 the latter approach, which appears more promising such architectures can in principle generalize across different molecules and do not rely on tailored system- specific featurizations. However, current equivariant Boltzmann generators based continuous normalizing flows (CNFs) face limitations in scalability due to their high computational cost during both training and inference. Recently, flow matching (2I) has emerged as a fast, simulation-fiee 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 Thus, the learned flow paths will deviate strongly from the OT paths (Figure Itt,, which increases computational cost and numerical errors. We tackle this problem by proposing novel equivariant OT flow matching objective for training equivariant flows on invariant densities, resulting in optimal paths for faster inference (see Figure T..f).
Our main contributions in this work are as follows:
1. We propose 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 new invariant dataset of alanine dipeptide and large Lennard-Jooe 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 molecule in Cartesian coordinates. Additionally, 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.
2
Related work
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
30, [1] 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 different thermodynamic state (c.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 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, 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 rime-resolutions, accurately capturing the dynamics.
R 29 25 20
[37 propose
To help speed up CNF training and inference, various authors (3 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 [1
42 433 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-mater fow-matching method for training energy- based coarse-grained models from flows introduced at the same time [46).
3 Method
In this section, we describe the key methodologies used in our study, highlighting important prior work.
3.1 Normalizing flows
Normalizing flows 14 provide a powerful framework for learning complex probability densities μ(x)) by leveraging the concept of invertible transformations. These transformations, denoted as So RD - TR.", map samples from a simple prior distribution g( = A(zaa 7) to samples from more complicated output distribution. This resulting distribution p(z1). 1), known as the push-forward
3

---

distribution, is obtained using the change of variable equation:
p(x1) = (z)) det O5 (x1)
(21))
dr
(ar) den
where
denotes the Jacobian of the inverse transformation fg"
3.2
Continuous normalizing flows (CNFs)
CNFs
48
49
are an instance of normalizing flows that are defined via the ordinary differential
equation
dfb(x) (x()),
(x)
(3)
dt where (t R" [0, R" is time dependent vector field. The time dependent flow (z) generates the probability flow path pr(x) [0, 1] x R + R from the prior distribution Po( q(xo) the push-forward distribution pi (2r). We generate samples from (x by first sampling zp from the prior and then integrating
fz(x) = '. dl' (1.56(o))
(4)
The corresponding log density change is given by the continuous change of variable equation
log (x) = log g(0)) L dt'v vo ((,)(())))
(5)
where V Uo (6 (z))
is the trace of the Jacobian.
3.3
Likelihood-based training
CNFs can trained using a likelihood-basee approach. this training method, the loss function is formulated the negative log- lieelhholho
[ log (z log det 0f (x1) (x1)
CNLL
Exsr
(6)
Ox
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 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.2p. Hence, the flow matching enables scaling to larger models and larger systems with the same computational budget. The flow matching training objective [1 22 regresses wo (t, r) to some target vector field (xr) via
(0) Ee- (xx)).)) ((t 2x(x)112-
practice, (ar) and (x(x) are intractable, it has been shown in 21 that the same gradients can be obtained by using the conditional flow matching loss:
LCFM Er- (...m. [lvo(t, ( (
(8)
The vector field we (2r) and corresponding probability path (z(x) are given in terms the conditional vector field ue (x) z) as well the conditional probability path p=(x12)
Pr (1(2)))u((1p) p(z)dz,
(x))) po(xl2)o(2)dz and ux(x))
(9)
(r(x)

---

where p(2) is some arbitrary conditioning distribution independent of x and t. For a derivation see 22 and [22]. There are different ways to efficiently parametrize u, (x12) and (z(z)2). We here focus on a parametrization that gives rise to the optimal transport path, as introduced in [2
2 (1) and p(z) x (1)
(10)
uz [x|zz) = 20 and (e(a)z)) (a)t x x (1- zo, o²),
(11)
where the conditioning distribution is given by the 2- Wasserstein optimal transport map x <2o, x1) between the prior g(2r)) and the target ((z1). The 2- 2-Wasserstein optimal transport map is defined by the 2- E-Wassesstein distance
W - inf / c(zoo z. ( dx1);
-inf xell
(12)
where II is the set of couplings as usual and c( = 1|z0 - x x1112 is the squared Euclidean distance in our case. Following [22], we approximate (2o, x1) by only considering batch of prior and target samples. Hence, for each batch we generate samples from p(z) as follows:
1. sample batches of points (z)....,. q and (zl..... x1 ~ j4,
2. compute the cost matrix M for the batch, i.e. M, |1 2 11
3. solve the discrete OT problem defined by M,
4. generate training pairs z' (#)) 21) according to the OT solution.
We will refer to this training procedure OT flow matching.
3.5 Equivariant flows
Symmetries can be described in terms of group G acting on finite-dimensional vector space via matrix representation p(g); E G. A map I: V - V is called G-invariant if T(p(g)=) (xx) for all g E and x E V. Similarly, map f V - V is called G-equivariant if f(p(g) = f())
for all g E G and x 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 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 vo (Theorem 2). Moreover, [17 50) show that the posh-forward distribution p(21) of 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 [d)) 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 [5] [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 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 (ro, 21) 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 To address this challenge, we propose sample pairs scales only with squared batch size (Figure [I.
using a cost function
č(zo, x1) min 1120 p(g)x1113 gEG
(13)
5

---

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 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 y be an OT map between G-invariant measures in and V2, using the cost function Then
1. T is G-equivariant and the corresponding OT plan xtv,, 12) is G-invariant.
2. For all pairs (x, [x)) and y G- (T(()))
(x,T(x)) [ clg x.g (()))du(g) min c(x, -.) v) gEG
(14)
3. T is also an OT map for the cost function غ
Refer to Appendix B. 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 é effectively aligns pairs samples in a manner consistent with how they are aligned under the C-equivariant OT map T'
In this work focus on O and (invaiinar distributions, which are common for molecules and multi-particle systems. The minimal squared distance for pair of points (mo, z1). taking into account these symmetries, can be obtained by minimizing the squared Euclidean distance over all possible combinations of rotations, reflections, and permutations
Č(zo, .1) = min rEO(D),
I p(rs)=1112
(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, approximate the minimizer by performing a sequential search
Z(zo, 21) min rESO(D)
|| - p(r8)=1 112, 8 arg min [00 (1))
(16)
*ES(N)
We demonstrate in Section 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 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 this loss as Equivariant OT flow matching. Although aligning samples in that way comes with significant computational overhead. this reordering can be done in parallel before or during training (see Appendix C.4p. 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 (t, is parametrized by an O(D)- and -(Nu-vquivariant graph neural network, similar to the one used in [8] 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 a;),
mly de (h4, Rd..h..o df))
(17)
a = 2 +
dij + pe(mm)))
(18)
h1 oi (b/.,,) m!)
m Σ
Di (mly))m
(19)
3#1
06 (t, = x x
(20)

---

Table I: Comparison of flows trained with different training objectives. Errors are computed over three runs.
Training type
NLL (4)
ESS (t)
Path length (4)
DW4
Likelihood [8]
1.72 + 0.01
86.87 + 0.19%
3.11 + 0.04
OT flow matching Equivariant OT flow matching
1.70 + 0.02
92.37 0.89%
2.94 + 0.02
1.68 = 0.01
88.71 + 0.40%
2.92 +
0.01 0.01
LJ13
Likelihood 18
15.83 + 0.07
39.78 6.19%
5.08 0.22
OT flow matching Equivariant OT flow matching
- 16.09 + 0.03
54.36 + 5.43%
2.84 0.01
16.07 = 0.02
57.98 + 2.19%
2.15
LJ55
OT flow matching
88.45 = 0.04
3.74 + 1.06% 4.42 - 0.35%
7.53 + 0.02 3.52 = 0.01
Equivariant OT flow matching
89.27 + 0.04
Alanine dipeptide
OT flow matching Equivariant OT flow matching
107.54 + 0.03
1.41 0.04% 0.69 + 0.05%
10.19 0.03
106.78 + 0.02
9.46 + 0.01
where 6a are neural networks, dy is the Euclidean distance between particle i and j. and oj is an embedding for the particle type. Notably, the update conserves the geometric center if all particles are of the same type (see Appendix [..2). otherwise we subtract the geometric center after the last layer. This ensures that the resulting equivariant vector field up x) conserves the geometric center. When combined with symmetric mean-free prior distribution, the push-forward distribution of the CNF will be OD)- and ((NN)-ivvaria
Experiments
In this section, we demonstrate the advantages of equivariant OT flow matching over existing training methods using four datasets characterized by invariant energies. We explore three different training objectives: (i) likelihood- based training, (ii) OT flow matching, and (iii) equivariant OT flow matching. For comprehensive overview of experimental details, including dataset parameters, error bars, learning rate schedules, computing infrastructure, and additional experiments, please refer to Appendix|| A and Appendix iin the supplementary material. For all experiments, we employ mean-free Gaussian prior distribution. The number of layers and parameters the equivariant CNF vary across datasets while remaining consistent within each dataset (see Appendix C). We provide naïve flow matching as an additional baseline Appendix IT However, the results are very similar to the ones obtained with OT flow matching, although the integration paths generally slightly longer.
6.1 DW4 and IJJ3
We first evaluate the different loss functions many-particle systems, DW4 and LJ13, that were specifically designed for benchmarking equivariant flows as described in [17]. These systems feature pair-wise double-well and Lennard-Jones interactions with 4 and 13 particles, respectively (see Appendix Cece for more details). While state-of-the-art results have been reported in [18), their evaluations were performed
very small test sets and were biased for the LJ13 system. To provide a fair comparison, we retrain their model using likelihood- bbased training as well as the two flow matching losses on resampled training and sets for both systems. Additionally, we demonstrate improved likelihood performance of flow matching results, presented Table | show that the two flow matching objectives outperform likelihood-based training while being computationally more efficient. The effective sample sizes (ESS) and negative log likelihood (NLL.) are comparable for the flow matching runs. However, for the IJ13 system, the equivariant OT flow matching objective significantly reduces the integration path length compared
their biased test set in Appendix A.3 The
7

---

b)
batch transport
Sampled
CTree
metohin
Nwe naenting
Eiuaaeterr maatiing
Batch size
Path length
d)
y$5 Integration Error
455 effective
f)
Energy distribution
Ri64 integration Error
taa narche
Effective samoles
crroww waciag
6AM
wort ara negroed
Equ Trasstatt matovan
Stavstanof mochieg stseeto
645
0.00
0.00
Integration steps
Function evaluations
Energy / FAT
Figure 2: Results for the LJ55 system (a) Integration paths per particle for OT flow matching (left) and equivariant OT flow matching (right). (b) Mean transport cost (squared distance) for training batches. (c) Integration path length distribution. (d) Integration error for fixed step size integrator (rk4) with respect to a reference solution generated by an adaptive solver (dropi5). (e) Effective samples vs number of function evaluations, i.e. evaluations of the vector field, for sampling batch size of 1000. (f) Energy histograms for a flow trained with equivariant OT flow matching.
to other methods due to the large number of possible permutations (refer to Figure || for a visual illustration).
6.2
LJ55
The effectiveness of equivariant OT flow matching becomes even more pronounced when training on larger systems, where likelihood training infeasible (see Appendix A.2 To this end, we investigate a large Lennard-Jones cluster with 55 particles (LJ55).
We observe that the mean batch transport cost of training pairs is about 10 times larger for OT flow matching compared to equivariant OT flow matching (Figure[pp) resulting in curved and twice as long integration paths during inference (Figure (a.c). However, the integration paths for OT flow matching are shorter than those seen during training, see Pp.pendi/AA.]fooo a more detailed discussion. We compare the integration solver (dropi5 (55)). As the integration paths follow nearly straight lines for the equivariant OT flow matching, the resulting integration error is minimal, while the error is significantly larger for OT flow matching (Figure 1). Hence, we can use a fixed step integrator, with e.g. 20 steps, for sampling for equivariant OT flow matching, resulting in a three times speed-up over OT flow matching (Figure 2+), emphasizing the importance of accounting for symmetries in the loss for large, highly symmetric systems. Moreover, the equivariant OT flow matching objective outperforms OT flow matching on all evaluation metries (Table []] To ensure that the flow samples all states, we compute the energy distribution (Figure 2 and perform deterministic structure minimization of the samples, similar to 17]. in Appendix A.7
caused by using fixed step integrator (rk4) instead of adaptive
6.3 Alanine dipeptide
In our final experiment, we focus on the small molecule alanine dipeptide (Figure 3a) in Cartesian coordinates. The objective is to train Boltzmann Generator capable of sampling from the equilibrium Boltzmann distribution defined by the semi-empiricas GFN2-xTB force-field [56). This emmiriaiii potential energy is invariant under permutations of atoms of the same chemical element and global rotations and reflections. However, here we consider the five backbone atoms defining the yp and v dihedral angles as distinguishable to facilitate analysis. Since running MD simulations with semi-empirical force-fields is computationally expensive, we employ a surrogate training set, further challenging the learning task.
8

---

b)
Mean batch transport cost
c)
Sampled path lengths
200
fow matching
of mus
Eap vertet Tow malching
66
fix marting
Mean cost
Ecutramtent nlw meiching
190
64
200
130
203 size
Path length
d)
Ramachandran plot
Energy distribution
Free energy projection
Frite weneennaan gy ke"
HD reeced smpple
Butzmarn porertoor
wel glead samp en
Free energy kgT
Uittralla sampling
Boltzmann gerenator
646
604
660
-90 Energy / Kg
-2
9
Figure 3: Results for the alanine dipeptide system (a) Alanine dipeptide molecule. (b) Mean transport
cost (squared distance) for training batches. (c) Integration path length distribution. (d) Ramachandran plot depicting the generated joint marginal distribution over the backbone dihedral angles ip and ei after filtering out samples with right-handed chirality and high energies.
(e) Energy histograms for samples generated by a flow trained with OT flow matching. (f) Free energy distribution along the slowest transition (jp dihedral angle) computed with umbrella sampling and the equivariant flow.
Alanine dipeptide data set generation The alamine dipeptide training data set is generated through two steps: (i) Firstly, we perform an MD simulation using the classical Amber ff99SBildn force- field for a duration of ms (20). (ii) Secondly, we relax 10 randomly selected states from the MD simulation using the semi-empirical GFN2-xTB force-field for 100 fs each. For more detailed information, refer to Appendix C.2 in the supplementary material. This training data generation significantly cheaper than performing long MD simulation with the semi-empirical force field.
Although the Boltzmann generator trained on biased training set, we can generate asymp- totically unbiased samples from the semi empirical target distribution by employing reweighting (Appendix B33) as demonstrated in Figure 33. In this experiment, OT flow matching outperforms equivariant OT flow matching in terms of effective sample size (ESS) and negative log likelihood (NLL) as shown in Table I However, the integration path lengths are still longer for OT flow matching compared to equivariant OT flow matching, as depicted in Figure 3c. Since the energy of alanine dipeptide is invariant under global reflections, the flow generates samples for both chiral states. While the presence of the mirror state reflects the symmetries of the energy, in practical scenarios, molecules typically do not change their chirality spontaneously. Therefore, it may be undesirable to have samples from both chiral states. However, it is straightforward to identify samples of the undesired chirality and apply a mirroring operation to correct them. Alternatively, one may use a SO(3) equivariant flow without reflection equivariance [54] to prevent the generation of mirror states.
Alanine dipeptide free energy difference The computation of free energy differences is a common challenge in statistical physics it determines the relative stability of metastable states. In the specific case of alanine dipeptide, the transition between negative and positive yp dihedral angle is the slowest process (Figure 3d), and equilibrating the free energy difference between these two states from MD simulation requires simulating numerous transitions, i.e. millions of consecutive MD) steps. which expensive for the semi-empirical force field. In contrast, we can train Boltzmann generator from data that is not in global equilibrium, i.c. using our biased training data, which is significantly cheaper to generate. Moreover, as the positive &p state is much less likely, we can even bias our training data to have nearly equal density in both states, which helps compute more precise free energy estimate (see Appendix C.5). The equivariant Boltzmann generator is trained using the OT flow matching loss, which exhibited slightly better performance for alanine dipeptide. We obtain similar results for the equivariant OT flow matching loss (see Appendix (A.4). To obtain an accurate estimation of the free energy difference, five umbrella sampling simulations are conducted along the yo dihedral angle using the semi-empirical force-field (see Appendix C.2). The free energy difference estimated by the Boltzmann generator demonstrates good agreement with the results

---

Table 2: Dimensionless free energy differences for the slowest transition of alanine dipeptide estimated from various methods. Umbrella sampling yields converged reference solution. Errors over five runs.
MD
relaxed MD
Umbrella sampling
Boltzmann generator
Free energy difference KgT
5.31
5.00
4.10 + 0.26
4.10 + 0.08
of these simulations (Table 2 and Figure 3D
whereas both the relaxed training data and classical Molecular Dynamics simulation overestimate the free energy difference.
7 Discussion
We have introduced novel flow matching objective for training equivariant continuous normalizing flows on invariant densities, leading to optimal transport integration paths even for small training batch sizes. By leveraging flow matching objectives, we successfully extended the applicability of equivariant flows to significantly larger systems, including the large Lennard- -annadd-Jones cluster (LJ55). We conducted experiments comparing different training objectives four symmetric datasets, and our results demonstrate that as the system size increases, the importance of accounting for symmetries within the flow matching objective becomes more pronounced. This highlights the critical role of our proposed flow matching objective in scaling equivariant CNFs to even larger systems.
Another notable contribution of this work is the first successful application of Boltzmann generator to model alanine dipeptide in Cartesian coordinates. By accurately estimating free energy differences using a semi-empirical force-field, our approach of applying OT flow matching and equivariant OT flow matching to equivariant flows demonstrates its potential for reliable simulations of complex molecular systems.
8 Limitations / Future work
While did not conduct experiments to demonstrate the transferability of our approach, the architecture and proposed loss function can potentially be used to train transferable models. We leave this for future research. Although training with flow matching is faster and computationally cheaper than likelihood training for CNFs, the inference process still requires the complete integration of the vector field, which be computationally expensive. However, if the model is trained with the equivariant OT flow matching objective, faster fixed-step integrators can be employed during inference. Our suggested approximation of Equation (15) includes the Hungarian algorithm, which has computational complexity of O(N*). To improve efficiency, this step could be replaced by heuristics relying on approximations to the Hungarian algorithm (57). Moreover, flow matching does not allow for energy based training, as this requires integration similar to NLL training. A potential alternative approach to initially train a CNF using flow matching with a small set of samples. Subsequent sample generation through the CNF, followed by reweighting to the target distribution, allows these samples to be added iteratively the training set, similarly as in 27 377. Notably, in our experiments, we exclusively examined Gaussian prior distributions. However, flow matching allows to transform arbitrary distributions. Therefore, our equivariant OT flow matching method holds promise for application in scenarios where the prior distribution is sampled through MD at distinct thermodynamic state, such as a higher temperature or a different level of theory 33 33 In these cases, where both the prior and target distributions are close, with samples alignable through rotations and permutations, we expect that the advantages of equivariant OT flow matching will become even pronounced than what observed in our experiments.
Building upon the success of training equivariant CNFs for larger systems using our flow matching objective, future work should explore different architectures for the vector field. Promising candidates, such as [58] [59 600 61 62). could be investigated to improve the modeling capabilities of equivariant CNFs. While our focus in this work has been on symmetric physical systems, it is worth noting that equivariant flows have applications in other domains that also exhibit symmetries, such as traffic data
generation 63 point cloud and set modeling [64] 655 as well as invariant distributions on arbitrary
manifolds [66] Our equivariant OT flow matching approach can readily applied to these areas of application without modification.
10

---

Acknowledgements
We gratefully acknowledge support by the Deutsche Forschungsgemmennscatt (SFB1114, Projects No. CO3, No. A04, and No. B08), the European Research Council (ERC CoG 772230 "ScaleCell"), the Berlin Mathematics center MATH+ (AA1-6), and the German Ministry for Education and Research (BIFOLD- - Berlin Institute for the Foundations of Learning and Data). We gratefully acknowledge the computing time granted by the Resource Allocation Board and provided on the supercomputer Lise at NHR @ ZIB and NHR @ Göttingen part of the NHR infrastructure. We thank Maaike Galama, Michele Invernizzi, Jonas Köhler, and Max Schebek for insightful discussions. Moreover, we would like to thank Lea Zimmermann and Felix Draxler, who tested our code and brought to our attention critical bug.
References
[1] Durk Kingma, Shakir Mohamed, Danilo Jimenez Rezende, and Max Welling. Semi-supervised learning with deep generative models. Advances in neural information processing systems, 27, 2014.
[2] Laurent Dinh, Jascha Sohl-Dickstein, and Samy Bengio. Density estimation using real NVP. In International Conference Learning Representations, 2017.
[3] Durk P Kingma and Prafulla Dhariwal. Glow: Generative flow with invertible 1x1 convolutions. In S. Bengio, H. Wallach, H. Larochelle, K. Grauman, N. Cesa-Bianchi, and R. Garnett, editors, Advances Neural Information Processing Systems, volume 31. Curran Associates, Inc., 2018.
[4] Tero Karras, Samuli Laine, Miika Aittala, Janne Hellsten, Jaakko Lehtinen, and Timo Aila. Analyzing and improving the image quality of stylegan. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition, pages 8110-8118 19, 2020.
[5] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. Advances in Neural Information Processing Systems, 30, 2017.
[6] Jacob Devlin, Ming-Wes Wei Chang, Kenton Lee, and Kristina Toutanova. BERT: Pre-training of deep bidirectional transformers for language understanding. In Proceedings of the 2019 Confer- ence of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies. Volume (Long and Short Papers), pages 4171 Minneapolis, Minnesota, June 2019. Association for Computational Linguistics.
[7] Luciano Floridi and Massimo Chiriatti. Gpt-3: Its nature, scope, limits, and consequences. Minds and Machines, 30:681-698 -64,, 2020.
[8] Frank Noé, Simon Olsson, Jonas Köhler, and Hao Wu. Boltzmann generators - sampling equilibrium states of many-body systems with deep learning. Science, 365:eaaw 1 147, 2019.
[9] John Jumper, Richard Evans, Alexander Pritzel, Tim Green, Michael Figurnov, Olaf Ron- neberger, Kathryn Tunyasuvunakool, Russ Bates, Augustin Zidek, Anna Potapenko, et al. Highly accurate protein structure prediction with alphafold. Nature, 596(7873):583-58, 2021.
[10] Joseph L. Watson, David Juergens, Nathaniel R. Bennett, Brian L. Trippe. Jason Yim, Helen E. Eisenach, Woody Ahem, Andrew J. Borst, Robert J. Ragotte, Lukas F. Milles, Basile I. M. Wicky, Nikita Hanikel, Samuel J. Pellock, Alexis Courbet, William Sheffler, Jue Wang, Preetham Venkatesh, Isaac Sappington, Susana Vázquez Torres, Anna Lauko, Valentin De Bortoli, Emile Mathieu, Regina Barzilay, Tommi S. Jaakkola, Frank DiMaio, Minkyung Baek, and David Baker. Broadly applicable and accurate protein design by integrating structure prediction networks and diffusion generative models. bioRxiv, 2022.
[11] Esteban G Tabak, Eric Vanden-Eijndee et al. Density estimation by dual ascent of the log- likelihood. Communications in Mathematical Sciences, 8(1):217-233 2010.
[12] Esteban G Tabak and Cristina V Turner. family of nonparametric density estimation algo- rithms. Communications on Pure and Applied Mathematics, 66(2).145-164, 2013.
11

---

[13] George Papamakarios, Eric Nalisnick, Danilo Jimenez Rezende, Shakir Mohamed, and Balaji Lakshminarayanan. Normalizing flows for probabilistic modeling and inference. CoRR, 2019.
[14] Danilo Rezende and Shakir Mohamed. Variational inference with normalizing flows. In International conference on machine learning, pages 1530-1538. PMLR, 2015.
[15] Gurtej Kanwar, Michael S. Albergo, Denis Boyda, Kyle Cranmer, Daniel C. Hackett, Sébastien Racanière, Danilo Jimenez Rezende, and Phiala E. Shanahan. Equivariant flow-based sampling for lattice gauge theory. Phys. Rev. Lett., 125:121601, Sep 2020.
[16] Isay Katsman, Aaron Lou. Derek Lim, Qingxuan Jiang, Ser Nam Lim, and Christopher M De Sa. Equivariant manifold flows. In M. Ranzato, A. Beygelzimer, Y. Dauphin, P.S. Liang, and Wortman Vaughan, editors, Advances in Neural Information Processing Systems, volume 34, pages 10600-10612. Curran Associates, Inc., 2021.
[17] Jonas Köhler, Leon Klein, and Frank Noé. Equivariant flows: exact likelihood generative learning for symmetric densities. In International conference on machine learning, pages 5361-5370. PMLR, 2020.
[18] Victor Garcia Satorras, Emiel Hoogeboom, Fabian Fuchs, Ingmar Posner, and Max Welling. E(n) equivariant normalizing flows. In M. Ranzato, A. Beygelzimer, Y. Dauphin, P.S. Liang, and J. Wortman Vaughan, editors, Advances in Neural Information Processing Systems, volume 34, pages 4181-4192. Curran Associates, Inc., 2021.
[19] Jonas Köhler, Michele Invernizzi, Pim de Haan, and Frank Noé. Rigid body flows for sampling molecular crystal structures. In International Conference on Machine Learning. ICML 2023, volume 202 of Proceedings of Machine Learning Research, pages 17301-17326. PMLR, 2023.
[20] Jonas Köhler, Andreas Krämer, and Frank Noé. Smooth normalizing flows. In M. Ranzato, A. Beygelzimer, Y. Dauphin, P.S. Liang, and J. Wortman Vaughan, editors, Advances in Neural Information Processing Systems, volume 34, pages 2796-2809. Curran Associates, Inc., 2021.
[21] Yaron Lipman, Ricky T. Q. Chen, Heli Ben-Hamu, Maximilian Nickel, and Matthew Le. Flow matching for generative modeling. In The Eleventh International Conference on Learning Representations, 2023.
[22] Alexander Tong, Nikolay Malkin, Guillaume Huguet, Yanlei Zhang, Jarrid Brotos, Kilian Fatras, Guy Wolf, and Yoshua Bengio. Conditional flow matching: Simulation-free dynamic optimal transport. arXiv preprint arXiv:2302.00482. 2023.
[23] Yuxuan Song, Jingjing Gong, Minkai Xu, Ziyao Cao, Yanyan Lan, Stefano Ermon, Hao Zhou, and Wei-Ying Ma. Equivariant flow matching with hybrid probability transport for molecule generation. In Thirty-seventh Conference on Neural Information Processing Systems, 2023.
[24] Laurence Illing Midgley, Vincent Stimper, Javier Antoran, Emile Mathieu, Bernhard Schölkopf, and José Miguel Hernández-Lobato. SE(3) equivariant augmented coupling flows. In Thirty- seventh Conference on Neural Information Processing Systems, 2023.
[25] Manuel Dibak, Leon Klein, Andreas Krämer, and Frank Noé. Temperature steerable flows and Boltzmann generators. Phys. Rev. Res., 4:1042005, Oct 2022.
[26) Peter Wirnsberger, Andrew J Ballard, George Papamakarios, Stuart Abercrombie, Sébastien Racanière, Alexander Pritzel, Danilo Jimenez Rezende, and Charles Blundell. Targeted free energy estimation via learned mappings. The Journal of Chemical Physics, 153(14):144112, 2020.
[27] Laurence Illing Midgley, Vincent Stimper, Gregor N. C. Simm, Bernhard Schölkopf, and José Miguel Hernández-LLoatoo Flow annealed importance sampling bootstrap. In The Fleventh International Conference Learning Representations, 2023.
[28] Xinqiang Ding and Bin Zhang. Deepbar: fast and exact method for binding free energy computation. Journal of Physical Chemistry Letters, 12:2509-255 2515, 3 2021.
12

---

[29] Hao Wu, Jonas Köhler, and Frank Noe. Stochastic normalizing flows. In H. Larochelle, M. Ranzato, R. Hadsell, M. F. Balcan, and H. Lin, editors, Advances Neural Information Processing Systems, volume 33, pages 3933-5944 Curran Associates, Inc., 2020.
[30] Xinqiang Ding and Bin Zhang. Computing absolute free energy with deep generative models. Biophysical Journal, 120())195a, 2021.
[31] Andrea Rizzi, Paolo Carloni, and Michele Parrinello. Multimap targeted free energy estimation, 2023.
[32] Andrea Rizzi, Paolo Carloni, and Michele Parrinello. Targeted free energy perturbation revisited: Accurate free energies from mapped reference potentials. Journal of Physical Chemistry Letters, 12:9449-944 9454, 2021.
[33] Michele Invernizzi, Andreas Krämer, Cecilia Clementi, and Frank Noé. Skipping the replica exchange ladder with normalizing flows. The Journal of Physical Chemistry Letters, 13:1 1643-
1 11649, 2022.
[34] Leon Klein, Andrew Y. K. Foong, Tor Erlend Fjelde, Bruno Kacper Mlodozeniec, Marc Brockschmidt, Sebastian Nowozin, Frank Noe, and Ryota Tomioka. Timewarp: Transferable acceleration of molecular dynamics by learning time-coarsened dynamics. In Thirty-seventh Conference on Neural Information Processing Systems, 2023.
[35] Emiel Hoogeboom, Víctor Garcia Satorras, Clément Vignac, and Max Welling. Equivariant diffusion for molecule generation in 3D. In Kamalika Chaudhuri, Stefanie Jegelka, Le Song, Csaba Szepesvari. Gang Niu, and Sivan Sabato, editors, Proceedings of the 39th International Conference Machine Learning, volume 162 of Proceedings of Machine Learning Research, pages 8867-8887. PMLR, 17-23 Jul 2022.
[36] Minkai Xu, Lantao Yu, Yang Song, Chence Shi, Stefano Ermon, and Jian Tang. Geodiff: A geometric diffusion model for molecular conformation generation. In International Conference on Learning Representations, 2022.
[37] Bowen Jing, Gabriele Corso, Jeffrey Chang, Regina Barzilay, and Tommi Jaakkola. Tor- sional diffusion for molecular conformer generation. In S. Koyejo, S. Mohamed, A. Agarwal, D. Belgrave, K. Cho, and A. Oh, editors, Advances in Neural Information Processing Systems, volume 35, pages 24240-24253 Curran Associates, Inc., 2022.
[38] Mathias Schreiner, Ole Winther, and Simon Olsson. Implicit transfer operator learning: Multiple time-resolution models for molecular dynamics. In Thirty-seventh Conference on Neural Information Processing Systems, 2023.
[39] Chris Finlay, Jömn-enrik Jacobsen, Levon Nurbekyan, and Adam Oberman. How to train your neural ode: the world of jacobian and kinetic regularization. In International conference on machine learning, pages 3164-3164. PMLR, 2020.
[40] Alexander Tong, Jessie Huang, Guy Wolf, David Van Dijk, and Smita Krishnaswamy. Trajecto- rynet: A dynamic optimal transport network for modeling cellular dynamics. In International conference on machine learning, pages 9526-9536. PMLR, 2020.
[41] Derek Onken, Samy Wu Fung, Xingjian Li, and Lars Ruthotto. Ot-flow: Fast and accurate continuous normalizing flows via optimal transport. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 35, pages 9223-9232, 2021.
[42] Michael Samuel Albergo and Eric Vanden-Eijnden. Building normalizing flows with stochastic interpolants. In The Eleventh International Conference on Learning Representations, 2023.
[43] Xingchao Liu, Chengyue Gong, and qiang liu. Flow straight and fast: Learning to generate and transfer data with rectified flow. In The Eleventh International Conference on Learning Representations, 2023.
[44] 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 27X1v:2304.14772
13

---

[45] David Alvarez-Melis, Stefanie Jegelka, and Tommi S. Jaakkola. Towards optimal transport with global invariances. In Kamalika Chaudhuri and Masashi Sugiyama, editors, Proceedings of the Swenty-Second International Conference on Artificial Intelligence and Statistics, volume 89 of Proceedings of Machine Learning Research, pages 1870-1879. PMIR, 16-18 Apr 2019.
[46] Jonas Köhler, Yaoyi Chen, Andreas Krämer, Cecilia Clementi, and Frank Noé. Flow-matching: Efficient coarse-graining of molecular dynamics without forces. Journal of Chemical Theory and Computation, 19())942-952, 2023.
[47] George Papamakarios, Erie T Nalisnick, Danilo Jimenez Rezende, Shakir Mohamed, and Balaji Lakshminarayanann Normalizing flows for probabilistic modeling and inference. Mach. Learn. Res., 22(57):1-64, 2021.
[48] Tian Qi Chen, Yulia Rubanova, Jesse Bettencourt, and David K Duvenaud. Neural ordinary differential equations. In Advances in neural information processing systems, pages 6571-6583, 2018.
[49] Will Grathwohl, Ricky T. Q. Chen, Jesse Bettencourt, and David Duvenaud. Scalable re- versible generative models with free-form continuous dynamics. In International Conference on Learning Reprexentations, 2019.
[50] Danilo Jimenez Rezende, Sébastien Racanière, Irina Higgins, and Peter Toth. Equivariant Hamiltonian flows. arXiv preprint
2019.
[51] Kilian Fatras, Younes Zine, Szymon Majewski, Rémi Flamary, Rémi Gribonval, and Nicolas Courty. Minibatch optimal transport distances; analysis and applications. arXiv preprint arXiv:2101.0192,, 2021.
[52] Harold W Kuhn. The hungarian method for the assignment problem. Naval research logisties quarterly, 2(1-2):83- 97, 1955.
[53] Wolfgang Kabsch. A solution for the best rotation to relate two sets of vectors. Acta Crystal- lographica Section A: Crystal Physics, Diffruction, Theoretical and General Crystallography, 32(5):922-923, 1976.
[54] Victor Garcia Satorras, Emiel Hoogeboom, and Max Welling. E (n) equivariant graph neural networks. In International conference on machine learning, pages 9323-9332. PMLR, 2021.
[55] John R Dormand and Peter J Prince. A family of embedded runge-kutta formulae. Journal of computational and applied mathematics, 6(1):19-26, 1980.
[56] Christoph Bannwarth, Sebastian Ehlert, and Stefan Grimme. Gfn2-xtb-as an accurate and broadly parametrized self-consistent tight-binding quantum chemical method with multipole electrostatics and density-deeenddet dispersion contributions. Journal of Chemical Theory and Computation, 15(3): 1671, 2019. PMID: 30741547.
[57] Jerome M Kurtzberg. On approximation methods for the assignment problem. Journal of the ACM (JACM), 94):419-439 1962.
[58] 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.
[59] Kristof Schütt, Oliver Unke, and Michael Gastegger. Equivariant message passing for the prediction of tensorial properties and molecular spectra. In International Conference on Machine Learning, pages 9377-9388. PMLR, 2021.
[60] Yi-Lun Liao and Tess Smidt. Equiformer: Equivariant graph attention transformer for 3d atomistic graphs. The Fleventh International Conference on Learning Representations, 2023.
[61] Johannes Gasteiger, Florian Becker, and Stephan Günnemann. Gemnet: Universal directional graph neural networks for molecules. Advances Neural Information Processing Systems, 006670-6600,,00 2021.
14

---

[62] Weitao Du, Yuanqi Du, Limei Wang, Dieqiao Feng, Guifeng Wang. Shuiwang Ji, Carla Gomes, and Zhi-Ming Ma. A new perspective on building efficient and expressive 3d equivariant graph neural networks. arXiv preprint arXiv:2304.04755 2023.
[63] Berend Zwartsenberg, Adam Šcibior. Matthew Niedoba, Vasileios Lioutas, Yunpeng Liu, Justice Sefas, Setareh Dabiri, Jonathan Wilder Lavington, Trevor Campbell, and Frank Wood. Conditional permutation invariant flows. arXiv preprint NXNv:2206.0902/ 2022.
[64] Marin Biloš and Stephan Günnemann. Equivariant normalizing flows for point processes and sets, 2021.
[65] Yang Li, Haidong Yi, Christopher Bender, Siyuan Shan, and Junier B Oliva. Exchangeable neural ode for set modeling. Advances in Neural Information Processing Systems, 33:6936 6946, 2020.
[66] Isay Katsman, Aaron Lou, Derek Lim, Qingxuan Jiang, Ser Nam Lim, and Christopher M De Sa. Equivariant manifold flows. Advances in Neural Information Processing Systems, 34:10600-1006 2021.
[67] Pierre Monteiller, Sebastian Claici, Edward Chien, Farzaneh Mirzazadeh, Justin M Solomon, and Mikhail Yurochkin. Alleviating label switching with optimal transport. Advances in Neural Information Processing Systems, 32, 2019.
[68] Yann Brenier. Polar factorization and monotone rearrangement of vector-valued functions. Communications pure and applied mathematics, 44(4))375-417, 1991.
[69] Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, Alban Desmaison, Andreas Kopf, Edward Yang. Zachary DeVito, Martin Raison, Alykhan Tejani. Sasank Chilamkurthy, Benoit Steiner, Lu Fang, Junjie Bai, and Soumith Chintala. Pytorch: An imperative style, high-performance deep learning library. In Advances in Neural Information Processing Systems 32, pages 8024-8035. Curran Associates, Inc., 2019.
[70] Michael Poli, Stefano Massaroli, Atsushi Yamashita, Hajime Asama, Jinkyoo Park, and Stefano Ermon. Torchdyn: Implicit models and neural numerical methods in pytorch. In Neural Information Processing Systems, Workshop Physical Reasoning and Inductive Biases for the Real World, volume 2, 2021.
[71] Rémi Flamary, Nicolas Courty, Alexandre Gramfort, Mokhtar Z. Alaya, Aurélic Boisbunon, Stanislas Chambon, Laetitia Chapel, Adrien Corenflos, Kilian Fatras, Nemo Fournier, Léo Gautheron, Nathalie T.H. Gayraud, Hicham Janati, Alain Rakotomamonjy, levgen Redko, Antoine Rolet, Antony Schutz, Vivien Seguy, Danica J. Sutherland, Romain Tavenard, Alexander Tong, and Titouan Vayer. Pot: Python optimal transport. Journal of Machine Learning Research, 22(78):1-8, 2021.
[72] Peter Eastman, Jason Swails, John D Chodera, Robert T McGibbon, Yutong Zhao, Kyle Beauchamp, Lee-Ping Wang, Andrew C Simmonett, Matthew P Harrigan, Chaya D Stern, et al. Openmm 7: Rapid development of high performance algorithms for molecular dynamics. PLoS computational biology, 10056699 2017.
[73] Ask Hjorth Larsen, Jens Jørgen Mortensen, Jakob Blomqvist, Ivano E Castelli, Rune Christensen, Marcin Dułak, Jesper Fris,, Michael N Groves, Bjørk Hammer, Cory Hargus, Eric D Hermes, Paul C Jennings, Peter Bjerre Jensen, James Kermode, John R Kitchin, Esben Leonhard Kolsbjerg, Joseph Kubal, Kristen Kaasbjerg, Steen Lysgaard, Jón Bergmann Maronsson, Tristan Maxson, Thomas Olsen, Lars Pastewka, Andrew Peterson, Carsten Rostgaard, Jakob Schiøtz, Ole Schütt, Mikkel Strange, Kristian S Thygesen, Tejs Vegge, Lasse Vilhelmsen, Michael Walter, Zhenhua Zeng, and Karsten W Jacobsen. The atomic simulation enviromment- python library for working with atoms. Journal of Physics: Condensed Matter, 29327)2273002, 2017.
[74] H Wiegand. Kish, I.: Survey sampling. john wiley & sons, inc., new york, london 1965, ix+ 643 31 abb., 56 tab., preis 83 s., 1968.
15

---

2)
U13 OT flow matching
b) Alanine dipeptide - OT flow matching
4)
USS - OT flow matching
230
Samplad urgite
Samples lengtha
175
Sumples path langtta
130
Path length
Path length
Path length
Figure 4: Path lengths distributions for models trained with OT flow matching. a) LJ13 system b) Alanine dipeptide system c) LJ55 system.
Supplementary Material
A
Additional experiments
A.1 Equivariant flows can achieve better paths than they are trained for
We make interesting observation regarding the performance of equivariant flows trained with the flow matching objective. As shown in Figure --c these flows generate integration paths that are significantly shorter compared to their training counterparts. While they do not learn the OT path and exhibit curved trajectories with notable directional changes (as seen in Figure T and Figure 2b) an interesting finding is the absence of path crossings. Any permutation of the samples would increase the distance between the latent sample and the push-forward sample. To quantify the length of these paths, we compute the Euclidean distance between the latent sample and the forwad sample. The relatively short path lengths can be attributed, in part, to the architecture the flow model. As the flow is equivariant with respect to permutations, it prevents the crossing of particle paths and forces the particles to change the direction instead, as an average vector field is learned.
A.2 Memory requirement
As the system size increases, the memory requirements for likelihood training become prohibitively large. In the case of the LJ55 system, even batch size of 5 necessitates more than 24 GB of memory, making effective training infeasible. In contrast, utilizing flow matching enables training the same model with batch size of 256 while requiring less than 12 GB of memory.
For the alanine dipeptide system, likelihood training with batch size of 8 demands more than 12 GB of memory, and batch size of 16 exceeds GB. On the other hand, employing flow matching allows training the same model with batch size of 256, consuming less than 3 GB of memory.
The same holds for energy based training, as need evaluate Equation (3) and Equation (4 as well during training. Therefore, we opt to solely utilize flow matching for these two datasets.
A.3
Comparison with prior work on equivariant flows
Both 17 and [18 evaluate their models the DW4 and LJ13 systems. However, 11 do not provide the numerical values for the negative log likelihood (NLL), only presenting them visually. On the other hand, [8] use small test set for both experiments. Additionally, their test set for LJ13 is biased as it originates from the burn-in phase of the sampling.
Although the test sets for both systems insufficient, we still compare OT flow matching and cquivariant OT flow matching to these previous works in Table should noted that the value range differs significantly from the values reported in Table I as the likelihood in this case includes the normalization constant of the prior distribution.
A.4 Alanine dipeptide free energy difference
Here we present additional results for the free energy computation of alanine dipeptide, as discussed in Section 6.3. The effective sample size for the trained models is 0.50 0.13%, which is lower due to training on the biased dataset (compare to Table We illustrate the Ramachandran plot in
16

---

Table 3: Comparison of different training methods for the DW4 and (biased) LJ13 dataset as used in for 10 and 10s training points, respectively.
18
Dataset
Likelihood 18
OT flow matching Equivariant OT flow matching
DW4 NLL (4)
7.48
0.05
7.21 0.03
7.18 0.02
LJ13 NLL (L)
30.41 0.16
30.34 + 0.09
30.30 = 0.03
Ramachandran plat
b)
Energy distribution
Free energy projection
relaed samples
Uniérilla campiing
Free energy KGT
6/0
606
eoc
320
42 Energy KgT
-2
Figure 5: Alanine dipeptide- free energy experiments. The model is trained on the biased dataset. (a) Ramachandran plot depicting the generated joint marginal distribution over the backbone dihedral angles yp and v after filtering out samples with right- right-handed chirality and high energies. (b) Energy histograms for samples generated by flow trained with OT flow matching. (c) Free energy distribu- tion along the slowest transition (fp dihedral angle) computed with umbrella sampling, the equivariant flow, and relaxed MD trajectory.
Figure Sp. the energy distribution in Figure 5p. and the free energy projections for Umbrella sampling, model samples, and relaxed MD samples in Figure
We obtain similar results with equivariant OT flow matching, as shown in Table 4
A.5 Alanine dipeptide integration paths
We compare the integration paths of the models trained on the alanine dipeptide dataset in b. Despite the higher number of atoms in the alanine dipeptide system (22 atoms), we observe comparable difference in the integration paths as observed in the LJ13 system. This similarity arises from the presence of different particle types in the alanine dipeptide, resulting in a similar number of possible permutations compared to the LJ13 system.
A.6 Integration error
The integration error, which arises from the use of fixed-step integration method (rk4) instead of adaptive solver (dopri5), is quantified by calculating the of absolute differences in sample positions and changes in log probability (dlogp). The discrepancies in position are depicted in Figure while the difference in log probability change are presented in Figure Gc. Both of these crucial in ensuring the generation of unbiased samples that accurately represent the target Boltzmann distribution of interest.
A.7 Structure minimization
Following similar approach the deterministic structure minimization conducted for the LJ13 system in [77 we extend our investigation to the sampled minima of the LJ55 system. We minimize the energy of samples generated by a model trained with equivariant OT flow matching, as well as samples obtained from the test data. The resulting structures represent numerous local minima within the potential energy distribution.
In Figure 6d, we observe good agreement between the distributions obtained from both methods. Furthermore, it is noteworthy that both approaches yield samples corresponding to the global min- imum. The energies of the samples at 300K are considerably higher than those of the minimized structures.
17

---

Table 4: Dimensionless free energy differences for the slowest transition of alanine dipeptide estimated with a Boltzmann Generator trained with the equivariant OT flow matching objective. Umbrella sampling yields converged reference solution. Errors over five runs.
Umbrella sampling 4.10 0.26
Equivariant OT flow matching
Free energy difference / KgT
3.95 0.19
Integration
Minimized structure distribution
Euuverant maicing
Deiomen
latent
patn
RC4
Integration steps
Energy kg!
Figure 6: (a.b) Integration paths for alanine dipeptide for models trained with OT flow matching and equivariant OT flow matching, respectively. The color code for the atoms is as usual: gray H. black C, blue N, red O. (c) Integration error for a fixed step integrator for the LJ55 system. (d) Minimized samples from a model trained on the LJ55 system with equivariant OT flow matching.
A.8 Log weight distribution
Samples generated with a trained Boltzmann generators can be reweighted to the target distribution via Equation (29). Generally, the lower the variance of the log weight distribution, the higher the ESS. Although large outliers might reduce the ESS significantly. We present the log weight distribution for the LJ55 system in Figure 7b showing that the weight distribution for the BG trained with equivariant OT flow matching has lower variance.
A.9 Approximation methods for equivariant OT flow matching
The approximation given in Equation (T6) is in practice performed using the Hungarian algorithm for permutations and the Kabsch algorithm for rotations. However, one could also think of other approximations, such as performing the rotation first or alternating multiple permutations and rotations. We compare different approximation strategies in Figure P.. The baseline reference is computed with an expensive search over the approximation given in Equation (T6). Namely, we evaluate Equation (T6) for 100 random rotations combined with the global reflection, denoted as O200, for each sample, i.e.
a(zo, x1)
min C(zo, p(0)x1)) oEO2on(D)
where ö(ro, 21) given by Equation 16). Hence, this 200 times more expensive than our approach. This baseline should be much closer to the true batch OT solution. Applying our approximation multiple times reduced the transportation cost slightly. Performing the rotations first, lead to inferior results. We observe the same behavior for the other systems. Therefore, used the approximation given Equation (16) throughout the experiments, showing that this sufficient to learn a close approximation of the OT map.
A.10 Different data set sizes
investigates different training set sizes to train equivariant flows, showing that for simple datasets already quite few samples are sufficient achieve good NLL. We test if this also holds for flow matching for the more complex alanine dipeptide system. To this end, we run our flow matching experiments also for smaller data set sizes of 10000 and 1000. We report our findings in T In agreement with prior work, the NLL becomes worse the smaller the dataset. The same is true for the ESS. However, the sampled integration path lengths are nearly independent on the dataset size.
Prior work 17 18
18

---

a)
b)
C)
0.200
LISS log weight distribtiin
GT Bow matching
260
AD2 Mean batch transport cost
120
AD2 Mean batch transport cost
eq OT now manthing
115
0.33
0
110
0.100
160
189
1.55
0.015
140
Nelve mulhhes
feterence
100
0.29
0.000
223
100
200
500
300
Batch
Batch
Figure 7: (a) Log weight distribution for samples generated for the LJ55 system. (b, c) Comparison of different approximation methods to approximate Equation (16) for alanine dipeptide.
Table 5: Comparison of flows trained with different training set sizes. Errors over three runs.
Training type
NLL (4)
ESS (t)
Path length (4)
Alanine dipeptide-
1000 training samples
OT flow matching Eq OT flow matching
104.97 + 0.05
0.25 + 0.04%
10.14 0.05
103.38 +:2 0.03
0.39 + 0.08%
9.27 0.02
Alanine dipeptide
10000 training samples
OT flow matching Eq OT flow matching
-105.16 0.17
0.33 + 0.07% 10.13 + 0.02
-00.33 + 0.15
0.40 = 0.06%
9.32 + 0.01
A.11 Naïve flow matching
We provide naive flow matching, i.e. flow matching without the OT reordering, as an additional baseline. Naïve flow matching results in even longer integration paths, while the ESS and likelihoods close the results of OT flow matching, as shown in Table 6 Flow matching with non equivariant architecture, e.g. fully connected neural network, failed for all systems but DW4 and is hence not reported.
B Proofs and derivations
B.1 Equivariant OT flow matching
We show that our equivariant OT flow matching converges to the OT solution, loosely following the theory in (67] and extending it to infinite groups. Let G be compact (topological) group that acts on an Euclidean n-space X by isometries, i.e., d(g x. g y) d(x, y), where is the Fuclidean distance. This definition includes the symmetric group Sn, the orthogonal group O(n), and all their
subgroups. Consider the OT problem of minimizing the cost function c(x, y) = d(I, y) 2 1 between G-invariant measures V1 and V/2
Let us further define the pseudometric d( y)" = mingeg d(x,g g y) We want to show that any OT map T: : X + X for cost functions c(x, y) = d(æ, 2 is also an OT map for ē(z, v) = d(x, u)P. Lemma 1. Let V1 and U2 be G-invariant measures, then there exists an G-invariant OT plan Tg x. 9 y) (x. y)-
Table 6: Results for naive flow matching, Same hyperparameters as for OT flow matching. Errors over three runs
Dataset
NLL (L)
ESS (t)
Path length (1)
DW4
1.68 + 0.01
93.01 + 0.12%
3.41 + 0.02
LJ13
16. 10 + 0.01
57.55 + 2.20%
3.77 + 0.01
LJ55
88.44 + 0.03
3.15 + 1.04%
8.66 + 0.03
Alanine dipeptide
107.56 + 0.09
1.42 - 0.51%
11.00 0.02
19

---

Let x' E {I[(,,, be (not necessarily invariant) OT plan. We average the OT plan over the group
Proof.
J/= l'gg - v)du((),
A(2, v)
(21)
where y is the Haar measure G.
An important property the Haar integral is that
Je ((ha)dulo) - f(g)d((g)
(22)
for any measurable function f and h E G. Therefore the average plan is G-invariant:
(h n. h. 1 <'gg bg v)dp(s) = Ic 7' 1g v)dp(e) = =
x,
Next we compute the marginals
J ((, dy) / x [ 7'lg x, dy)du(() 1 1 7'(9 x.g- dy)d)(())
= / vilg 11(g- x)dp(()p)))))))))))))))))))))x) = 1 (zz)du(()) = v2(
and analogously Jx (dx, y) vz()). Hence, the average plan is also coupling, [[( 12).
The cost of the average plan
f c(x, v)(((r dy)
/ c(z.8)
/ l'gg dx, g- dy)dp(g)
(Definition of T)
x2
[aJe c( (x,,))=( dx, g dy)dp(g)
(Fubini's theorem)
[ [ clg-t clg-t.2, x.---
म) (dx, dp))
cellume-preservinn substitution)
( { dx,u)m (dr, dy dg)dp(g)
(Isometric group action)
xx c(a, c(x,y) (dæ, dy),
and hence the average plan is optimal.
Lemma 2. Let T: x0 + y be an G-equivariant OT map, x E X, and g" = arg argmin c(x, g T(x)).
gEG
Then g" = id.
Proof. T' maps the orbit of x, denoted G {g- xlg G}. the corresponding orbit of (zx).
due to equivariance.
If there exists an xx E X such that g" 4 id then this property holds true for the entire orbit of x².
Now, let us define an Geeuivarriant map T" x + y such that TG = "TT T(x) everywhere else. However, in this case, T would not be the OT map, as c(x", T' (x'x)))) )) o(x", T)x*)). Therefore, follows that g = id holds for the G-equivariant OT map.
x and ('
Discrete optimal transport Consider finite number of iid samples from the two G-invariant distributions V1, P2, denoted as {m,) {ye OT plan. The discrete optimal transport problem seeks the solution to
1 respectively. Let further π (v1, 22) be G G-invariant
arg min
Yig Mij.
(23)
PCP
20

---

where My c(2s, Bj) represents the cost matrix and P is the set of probability matrices defined as p= 1y E (xt) 1,/n, 7² 1, = 1n/nj. The total transportation cost C' is given by
= min Σ Mij. 7EP
(24)
Approximately generating sample pairs from the OT plan (11, 22), required for flow matching (Section 3.4), can be achieved by utilizing the probabilities assigned to sample pairs by y Let further
arg min 7EP
Σ Yay Mi
225)
denote the optimal transport matrix for the cost function ö(v, y) and È the corresponding transporta- tion cost. Sample pairs (#i. v/}) generated according to 7th are orientated along their orbits to have minimal cost, i.e.
(or, v/) (24, 9 V;), g" == arg min c(2 9 Wj).
(26)
gEG
Note that the group action g does not change the probability under the marginals, as and 1 both G-invariant.
respectively. Then, sample pairs generated based on 5 have lower averuge cost than sample pairs generated according to y
Lemma Let 7 and 5" defined as in Equation (3) and Equation 25
Proof. The total transportation cost Ò is always less than or equal to C' since the cost function a(sa, Wj) & U5) for all and by definition. Note that the number of samples approaches infinity, the inequality between Ö and C becomes an equality.
Note that in practice, samples generated wrt Moreover, by construction, sample pairs generated according to 7" satisfy arg min mingego(1,, 9 yj) id, as indicated in Equation (26). This property, which is also true for samples from the G-invariant OT plan according to Lemma 2 ensures that the generated pairs are correctly matched within their respective orbits. However, sample pairs generated according to y* do not generally possess this property.
much better approximation, shown in Section
We now combine findings in the following Theorem.
Theorem 1. Ler T: z + y be an OT map between G-invariant measures and V2, using the cost function e. Then
1. T is G-equivariant and the corresponding OT plan v2) is G-invariant.
2. For all pairs (x, (x)) and G. GTT()
clx, (zz)) [ elg x.g. ((g)((z)) = min c(z, vv)
(27)
gEG
3. T also an OT map for the cost cost function
Proof. We prove the parts of Theorem in order.
1. Follows from Lemma and the uniqueness of the OT plan for convex cost functions 68]]
2. Follows from Lemmal
3. Follows directly from 2., Lemma and
/ min cla, g y) dy) Sor x2
x-
c(x, ()x,y))(de,d dy).
(28)
sEG
However, this OT map is not unique.
21

---

Hence, we can generate better sample pairs, approximating the OT plan (V1, v2), by using č as cost function instead of c. Training pairs (I,, vf) are orientated along their orbits to have minimal cost (Equation (26).
These findings provide the motivation for our proposed equivariant flow matching objective, as described in Section 4 as well the utilization of G-equivariant normalizing flows, discussed in Section 5
B.2 Mean free update for identical particles
The update, described in Section 5 conserves the geometric center if all particles are of the same type.
Proof. The update of layer I does not alter the geometric center center mass if all particles are of the same type as ml, m); and then
(2f 2/)
r + Σ
)a(m))))
dij +1
(r)
xt)) + (x² x!t)
(bl(m))))
dip +1
B.3 Selective reweighting for efficient computation of observables
As normalizing flows are exact likelihood models, we can reweight the push-forward distribution to the target distribution of interest. This allows the unbiased evaluation of expectation values of observables O through importance sampling
Ex, (p.()) [w( (0(1)]
m(x1) P1 (21)
00) p
Eg, (pn (as) (x1)) w(<1)
(29)
C Technical details
C.1 Code libraries
Flow models and training are implemented in Pytorch 69
using the following code libraries: bgflow
g
17,, torchdyn [70), Pot: Python optimal transport
and the code corresponding to [8] The
MD simulations are run using OpenMM (72), ASE 73 and xtb-python [56].
The code will be integrated in the bgflow otheyy/tttteggggghhb...ooooogogp
C.2 Benchmark systems
For the DW4 and LJ13 system, we choose the same parameters as in
17 18
DW4 The energy Uxx) for the DW4 system is given by
Dd (x)
Σ a (dy do) + b - do) + c(dg do)
(30)
2r
1.
where dey the Fuclidean distance between partiele and j. the parameters are chosen as 0, 4, 0.9, do 4, and T 1, which is dimensionless temperature factor.
22

---

Table 7: Model hyperparameters
Dataset
L
Ahidden
Num. of parameters
DW4, LJ13
3
32
22468
LJ55
204936
alanine dipeptide
5
64
147599
LJ13
The energy (x) for the LJ13 system is given by
(-)
Σ 2
Ym
- 2
Pm
(31)
2T
dag
dgs
with parameters Tm = 1, 1, and 7 1.
LJ55
We choose the same parameters as for the LJ13 system.
Both the LJ13 and the LJ55 system were generated with MCMC with 1000 parallel chains, where each chain is run for 10000 steps after a long burn-in phase of 200000 steps starting from a random generated initial state.
Alanine dipeptide The alanine dipeptide training data temperature 7 300K set is generated through two steps: (i) Firstly, we perform an MD simulation, using the classical Amber ff99sBildn force-field, at 300K for implicit solvent for a duration of ms [20) using the openMM library.
(ii) Secondly, we relax 10a randomly selected states from the MD simulation for 100 fs each, using the semi-empirical GFN2-xTB force-field [56] and the ASE library 73 with a friction constant of 0.5 a.u.
Both simulations a time step of fs. We create test set in the same way.
Alanine dipeptide umbrella sampling To accurately estimate the free energy difference, we performed five umbrella sampling simulations along the yp dihedral angle using the semi-empirical GFN2-xTB force-fiele field. The simulations were carried out at 25 equispaced ço angles. For each angle, the simulation was initiated by relaxing the initial state for 2000 steps with high friction value of 0.5 a.u. Subsequently, molecular dynamics (MD) simulations were performed for total of 10th steps using the GFN2-xTB foree field, with a friction constant 0.02 a.u. The system states were recorded every 1000 steps during the second half of the simulation. A timestep of 1 fs was utilized for the simulations.
The datasets (Bdeeaa08454446bgd022255555
are
available
at
https://osf. 10/. /t//rrgg//vveewwonll//
C.3 Hyperparameters
Depending on the dataset, different model sizes were used, as reported in Table All neural networks da have one hidden layer with Whidden neurons and SiLU activation functions. The embedding a, given by single linear layer with Phidden neurons.
We report the used training schedules in Table 8 Note that 5e-4/5e-5 in the second row means that the training was started with learning rate of 5e-4 for 200 epochs and then continued with a learning rate of 5e-5 for another 200 epochs. All batches were reordered prior to training the model. The only exception is alanine dipeptide trained with equivariant OT flow matching.
C.4 Parallel OT batch generation
As Equation (66) needs to be evaluated for each possible sample pair in a given batch, the computation cost of equivariant OT flow matching is quite high. However, simple way to speed the equivariant OT training to generate the batches beforehand or parallel to the training process. This process highly parallelizable and can be performed CPUs, which are in practice usually more available
23

---

Table 8: Training schedules
Training type
Batch size
Learning rate
Epochs
Training time
DW4
Likelihood 18
256
5e-4
20
3h
OT flow matching Equivariant OT flow matching
256
5e-4/5e-5
200/200
0.5h
256
5e-4/5e-5
200/200
0.5h
LJ13
Likelihood [88
64
5c-4
13h
OT flow matching
256
5e-4/5e-5
1000/1000
3h
Equivariant OT flow matching
256
5e-4/5e-5
1000/1000
LJ55
OT flow matching Equivariant OT flow matching
256
5e-4/5e-5
600/400
17h
256
5e-4/5e-5
600/400
17h
Alanine dipeptide
flow matching Equivariant OT flow matching (Batch reordering during training)
256
5e-4/5e-5
1000/1000
6.5h
32
5e-4/5e-5
200/200
25h
Table 9: Wall-clock time to reorder single batch for the two OT flow matching methods. The times for the DW4 and IJ13 system are below 0.01 seconds for OT flow matching and are therefore not reported.
Training type
Batch size
Wall-clock time
DW4
Equivariant OT flow matching
256
3.6s
LJ13
Equivariant OT flow matching
256
4.5s
LJ55
OT flow matching
256
0.01s
Equivariant OT flow matching
256
20.2s
Alanine dipeptide
OT flow matching
256
0.01s
Equivariant OT flow matching
256
22.4s
Equivariant OT flow matching
32
0.4s
than GPUs and also in higher numbers. This also allows for larger batch sizes for the equivariant OT model and comes at little additional cost. Hence, scaling equivariant flow matching to even larger systems should not be an issue. The wall-clock time for the reordering of single batch on single CPU reported in Table 9
C.5 Biasing target samples
In the case of alanine dipeptide, the transition between negative and positive dihedral angles is the slowest process, as depicted in Figure (3.. Since the positive &p state is less probable, we can introduce bias in our training data to achieve nearly equal density in both states. This approach aids in obtaining more precise estimation of the free energy. To ensure a smooth density function, we incorporate weights based the von Mises distribution frM. The weights are computed along the & dihedral angle as
w(pp) = 150 f.M (vllu 1, 10) 1.
(32)
24

---

We then draw training samples based on the weighted distribution.
C.6 Effective samples sizes
The effective sample sizes are computed with Kish's equation 74]. We use 10 samples for the DW4 and LJ13 system, 2 x 10 for alanine dipeptide, and 1 x 10th for the LJ55 system per model.
C.7
Error hars
Error bars in all plots given by standard deviation, averaged over runs, if not otherwise indicated. The same applies for all errors reported in the tables. Exceptions are the Mean batch transport cost plots, where average over 10 batches and the integration error plots, where average over 100 samples for each number of steps.
C.8 Computing infrastructure
All experiments for the DW4 and LJ13 system were conducted GeForce GTX 1080 Ti with 12 GB RAM. The training for alanine dipeptide and the LJ55 system were conducted on Geroree RTX 3090 with 24 GB RAM. Inference performed NVIDIA A100 GPUs with 80GB RAM for alanine dipeptide and the LJ55 system.
25


Reviews:
[{"review_id": "MFYUlQi2Su", "reviewer": "Reviewer_W92U", "summary": "This paper extends on existing works on Flow Matching with minibatch OT solutions to the case of invariant cost functions. In particular, where the invariance is given by an SO (permutation + rotation) group. This is mainly a method of correcting the minibatch bias, since (non-equivariant) minibatch OT will still converge to the correct mapping when the minibatch size goes to infinity. Empirically, it is shown that equivariance OT matching results in lower transport costs, and hence shorter path lengths which may indicate that it is computationally faster to simulate after training than non-equivariant OT (though not shown).", "questions": "  - Eq 14 gives a suboptimal solution to Eq 13. Has the authors looked into how this affects either the transport cost or the final trained model? For instance, one could perhaps solve in an alternating fashion the optimal permutation and rotation. The current method can be seen as a single step approximation of this procedure.\n  \n  - It seems that all of the experiments are on fitting to potential functions. However, the main algorithm assumes data is sampled IID. This mismatch between the problem statement and benchmark problems seems to not be discussed much at all. Is my understanding correct that the training datasets are all simulated from some MCMC procedure?\n\n  - Related to the above, why not experiment on data sets such as QM9, which seems to be used by prior works such as equivariant diffusion models [1]?\n\n  - Table 1 doesn't seem to show large improvement for equivariant OT. Why is that? From what I understand, at small batch sizes, equivariant OT should perform better. But is it perhaps that at regular batch sizes, the difference between OT and equivariant OT becomes much smaller? \n\n  - Also, it would be interesting to see how the regular flow matching performs here, since flow matching with conditional OT paths can approximate OT solutions quite well already [2], and there has been a few papers discussing the relation between standard diffusion models and optimal transport [3].\n\n  - As part of future work section, authors discuss approximations to solving OT problems. It might be worth testing out the approximate algorithms proposed in the published work on minibatch OT + flow matching [2], which the authors should cite. Here they showed heuristic algorithms that have similar performance to minibatch OT but with faster compute cost.\n\n  - Regarding the transport cost, one reason for wanting a low transport cost that this reviewer is familiar with, is the sampling efficiency. However, this aspect isn't shown quantitatively in the experiments. For instance, a plot of ESS (or any performance metric) vs NFE, comparing OT and equivariant OT, would be very useful. From the figures, it is not strongly convincing that equivariant OT is actually performing significantly better than OT yet. \n\n  - Another avenue perhaps is a plot of ESS vs batch size. It would be good to understand at what batch sizes, equivariant OT exhibits better performance than OT.\n\n  - Showing wallclock time of equivariant OT, OT and regular flow matching would be ideal. For instance, Table 9 of [2] shows that using minibatch OT solutions barely increases compute by 4%, but I think having to solve for optimal permutation (especially when the number of particles is high) is a much harder task. It isn't clear whether the increase in compute cost is a good tradeoff for better transport cost (I can believe it is; just that it isn't shown in the paper). For instance, a convergence plot vs wallclock time would be very convincing.\n\n  - As a reviewer who is not familiar with the experimental setups and cannot recognize these systems by name, it would be useful to have a table summarizing each experiment. For instance, list out what is the potential function, and how many particles are in the system? This can help provide a better sense of how challenging each task is.\n\n  - Finally, I am just somewhat perplexed by the motivation of training an equivariant generative model from potential functions, when the training aspect requires MCMC sampling from the desired potential function as a first step. This seems to nullify the simulation-free aspect of flow matching, and it seems that an approach for posterior inference would be much better suited. At the same time, I'm sure the proposed algorithm could work well in settings where a dataset is provided (e.g. QM9 seems to be standard), but these experiments do not appear in this paper.\n\n[1] \"Equivariant Diffusion for Molecule Generation in 3D\" ICML 2022. https://arxiv.org/abs/2203.17003.\n\n[2] \"Multisample Flow Matching: Straightening Flows with Minibatch Couplings\". ICML 2023. https://arxiv.org/abs/2304.14772.\n\n[3] \"Understanding DDPM Latent Codes Through Optimal Transport\". ICLR 2023. https://openreview.net/forum?id=6PIrhAx1j4i.\n\n----\n\nI am satisfied with the authors' answers and have increased my score in light of the new experiments and clarifications. See my reply for some comments regarding the rebuttal.", "limitations": "Overall, I think the benefits of equivariant OT compared to OT (and regular flow matching) can be showcased more. Theoretically, I can believe what the authors are claiming, but empirically I don't quite see these emphasized as part of the empirical results yet. I suggested some ideas for additional plots in the Questions section.", "rating": 7, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "  - Well-written and easy to understand the high-level idea.\n\n  - Straightforward extension of existing works to training equivariant flows.", "weaknesses": "  - Eq 13 and 14 are not easy to solve, and it is unclear how much compute (or wall-clock time) these subproblems require.\n\n  - Lack of comparison to other baselines (equivariant diffusion models, standard flow matching).\n\n  - The empirical differences between equivariant OT and OT seem marginal.\n\n  - While the paper discusses transport cost, it is not shown that these translate to faster sampling algorithms.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their detailed review and questions. We now address their comments individually. To keep the rebuttal within the character limit, we often only cite parts of each question.\n\n> Eq 13 and 14 are not easy to solve ...\n\nEq 13 is indeed intractable, but Eq 14 gives a good tractable approximation, see global rebuttal.\n\n> Lack of comparison to other baselines ...\n\nWe added the suggested standard flow matching baseline in the global rebuttal.\n\n> The empirical differences between equivariant OT and OT seem marginal.\n\nThe difference for equivariant OT compared to normal OT is only marginal for systems with few symmetries, i.e. less than 15 interchangeable particles, which is the case for DW4, LJ13 and ALA2. However, for the much larger LJ55 system (at least 3 times as many identical particles) the difference is significant. Crucially, the straight OT sampling paths allow using a Runge-Kutta integrator instead of the adaptive dopri5, which is significantly faster. With the “normal” OT the particle integration paths change directions (see Figure 2b,c,e), which requires small step sizes where the particles turn. Hence, this shows that it is important to use equivariant OT flow matching when scaling to larger systems to obtain optimal paths. This effect is not that prominent for the smaller systems, where normal OT flow matching already works quite well. See also the global rebuttal.\n\n> While the paper discusses transport cost, it is not shown that these translate to faster sampling algorithms.\n\nWe report in section 6.2 a speed-up of about 10 for the inference for the LJ55 system. Moreover, we performed additional experiments to highlight his further, see global rebuttal.\n\n> Eq 14 gives a suboptimal solution to Eq 13...\n\nFor the discussed systems, our approximation is quite close. We also tested the suggested approximation algorithm, see global rebuttal.\n\n> It seems that all of the experiments are on fitting to potential functions...\n\nThe algorithm does not require IID samples, the samples can be biased. In fact, we show this for the molecule experiments, where the training samples do not stem from the target potential, which is a semi-empirical potential. Instead, we generate the samples with a classical force field, which is about three orders of magnitude cheaper, and then relax these samples with respect to the semi-empirical force field. Hence, the training samples are biased and not IID. Although, the flow then learns a biased potential function, we can reweight the flow samples to the unbiased distribution (see appendix B.3). This is potentially much faster than doing the simulation wrt to the semi-empirical force field. We further validate that we indeed generate samples from the unbiased distribution by comparing to an umbrella simulation (Sec 6.3).\n\n> why not experiment on data sets such as QM9, ...\n\n QM9 is a different task, i.e. conformer generation (single conformations instead of the Boltzmann distribution). Moreover, the QM9 dataset consists of only small molecules and the aim of this paper is to scale to Boltzmann Generators for the first time to significantly larger systems.\n\n> Table 1 doesn't seem to show large improvement for equivariant OT ...\n\nSee above. We investigated different batch sizes in the global rebuttal.\n\n> ... how the regular flow matching performs here,...\n\nWe added this baseline, see global rebuttal.\n\n> ... approximations to solving OT problems ...\n\nWe will cite that in the related work section. See global response.  \n\n> ... plot of ESS  vs NFE... and a plot of ESS vs batch size ...\n\nThese are excellent suggestion. We compare as suggested, highlighting the importance of equivariant flow matching when scaling to larger systems. See general rebuttal.\n\n> ... wallclock time ...\n\nWe compare the wall-clock training time in appendix C.3 Table 5. The training does take longer for the equivariant flow matching. However, a simple way to speed up the equivariant OT training is to generate the batch pairs not during training. This process is highly parallelizable and can be performed on CPUs, which are in practice usually more available than GPUs. See global rebuttal for the suggested plot in Figure 1h.\n\n> ... table summarizing each experiment\n\nWe agree that the information of the dataset should be more visible. We included most of the information only in the appendix C.2 and will move parts to the main part in the final version of the manuscript.\n\n> ... motivation of training an equivariant generative model from potential functions, when the training aspect requires MCMC sampling from the desired potential function as a first step...\n\nWe show for Alanine dipeptide that we do not require samples from the equilibrium target Boltzmann distribution, as discussed above. This is a common setting, as the target potential function is often known up to a constant, but sampling from the equilibrium distribution is difficult. However, generating biased sample, e.g. in different meta-stable states, is often feasible, which can then be used to train a Boltzmann Generator and produce unbiased samples. \nUnfortunately, posterior inference requires backpropagating through the whole integration path, which is infeasible for the larger system, as discussed in appendix A.2. Note that also for posterior inference, we usually require some data from the target distribution for initial training of the model.\nHowever, our work paves the way for transferable Boltzmann generators, as we operate in Cartesian coordinates. These will allow training on (biased) trajectories of a few molecules and be applicable to unseen ones. Hence, requiring simulations only for the training molecules. We leave this exciting avenue for future research."}, {"review_id": "xvoFfJpgW9", "reviewer": "Reviewer_nDvW", "summary": "The authors propose equivariant flow matching, which provides a way of incorporating syymetries into flow matching objectives.\nSpecifically, they propose to replace the squared Euclidian distance cost used in the flow matching objective, with (approximately) the minimum squared Euclidian distance over all possible group actions.\nThey demonstrate in their experiments that this improves training of equivariant flows - obtaining similar or improved performance than vanilla OT flow matching with shorter paths.\n", "questions": "- How was the effective sample size calculated? (equation, number of samples etc)\n- The training time for the Equivariant OT flow matching is significantly longer than the vanilla OT flow matching for LJ55 and alanine dipeptide, even though it is trained for less epochs with a smaller batch size - why is this the case? (e.g. is it due to the extra compute required for the search over rotations/permutations?)\n- Could you please provide further description on the rationale behind the hyper-parameters used for each experiment in the appendix?\n- Could you please specify the compute used for the MD simulation (e.g. number of target energy evaluations, runtime) in the appendix?", "limitations": "The Equivariant flow matching method results in longer training times for LJ55 and alanine dipeptide, but no discussion of why / analysis of this is provided (see Questions). \nOne of the headline results from the abstract (re alanine dipeptide) is obtained using an existing method (without the new proposed method). ", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": " - They identify a clear problem of how symmetries (especially permutation symmetries) cause issues with OT flow matching, resulting a (possibly prohibitively) large batch size required.\n - Their solution is simple and fits the problem well.\n - Their experimental results are consistent with the theory (their CNF has shorter paths).", "weaknesses": " - The performance results are mixed - the flow trained with vanilla OT is sometimes better.\nSpecifically for alanine dipeptide the flow trained with vanilla OT performs better - thus the Equivariant flow matching method does not seem relevant to one of the headline results \"for the first time we obtain a Boltzmann generator with significant sampling efficiency without relying on tailored internal coordinate featurization\". \nOn first read of the abstract the wording makes it seem like this result was obtained due to the authors proposed  method rather than applying the existing flow matching technique to the problem.\n - A CNF trained with a classic score matching loss instead of flow OT would be a good obvious baseline but is not included. ", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their detailed review and questions. We now address their comments individually. \n\n> The performance results are mixed - the flow trained with vanilla OT is sometimes better. \n\nThe difference for equivariant OT compared to normal OT is only marginal for systems with few symmetries, i.e. less than 15 interchangeable particles, which is the case for DW4, LJ13 and ALA2. However, for the much larger LJ55 system (at least 3 times as many identical particles) the difference is significant. Crucially, the straight OT sampling paths allow using a Runge-Kutta integrator instead of the adaptive dopri5, which is significantly faster. With the “normal” OT the particle integration paths change directions (see Figure 2b,c,e), which requires small step sizes where the particles turn. Hence, this shows that it is important to use equivariant OT flow matching when scaling to larger systems to obtain optimal paths. This effect is not that prominent for the smaller systems, where normal OT flow matching already works quite well.\n\n> Specifically for alanine dipeptide the flow trained with vanilla OT performs better - thus the Equivariant flow matching method does not seem relevant to one of the headline results \"for the first time we obtain a Boltzmann generator with significant sampling efficiency without relying on tailored internal coordinate featurization\". On first read of the abstract the wording makes it seem like this result was obtained due to the authors proposed method rather than applying the existing flow matching technique to the problem.\n\nWe obtain the described result with both the OT flow matching and the equivariant OT flow matching. We extend Table 2 with results obtained from equivariant OT flow matching, see global rebuttal.\nEnhancing MD simulations using ML models is a rapidly growing field, as evidenced by the increasing number of AI4Science groups focused on this topic. Prior to our work, it was uncertain whether training a Boltzmann Generator in Cartesian coordinates to generate samples from the equilibrium distribution of real molecules was possible. We are the first to present a model and training algorithm that successfully achieve this. We believe our results are of great interest to the community and will stimulate further research in this area, potentially leading to transferable Boltzmann Generators.\nMoreover, we consider an even more difficult problem as we train the model with biased samples, instead of sample from the equilibrium distribution. We present two training algorithms that achieve this. Both including the known concepts of flow matching and EGNNs. However, we are the first to combine these concepts, resulting in this novel finding. Moreover, we show that when scaling to larger systems with more symmetries, e.g. as shown for LJ55 with 55 identical Lennard-Jones particles, it becomes critical to do equivariant OT flow matching, to preserve the optimal transport sampling paths. Interestingly our work shows that for systems with less symmetries, e.g. fewer identical particles like DW4, LJ13, and alanine dipeptide, the benefits of including equivariant OT flow matching are marginal and sometimes even show worse performance as in the case of alainine dipeptides. We will highlight these findings more in the final version of the manuscript and change the abstract accordingly. \n\n> A CNF trained with a classic score matching loss instead of flow OT would be a good obvious baseline but is not included.\n\nWe included normal flow matching as baseline, see global rebuttal. \n\n> How was the effective sample size calculated? \n\nKish's equation. We used 10000 samples for the DW4 and LJ13 system and 100000 for LJ55 and alanine dipeptide. We will add the equation, reference, and values to the appendix.\n\n> The training time for the Equivariant OT flow matching is significantly longer than the vanilla OT flow matching for LJ55 and alanine dipeptide, even though it is trained for less epochs with a smaller batch size - why is this the case?  (e.g. is it due to the extra compute required for the search over rotations/permutations?)\n\nYes, the reviewer's intuition is correct. The training takes longer for the equivariant flow matching because of the search over rotations and especially permutations. However, the training data preparation can be performed prior to or during the training and is trivially parallelizable on CPUs. Allowing to essentially train the equivariant flow matching models as fast as the OT flow matching models. See the global rebuttal for more details. \n\n> Could you please provide further description on the rationale behind the hyper-parameters used for each experiment in the appendix?\n\nThe used hyperparameters are discussed in appendix C.3. We performed mostly manual hyperparameter searches and used larger models for the more challenging systems.\n\n> Could you please specify the compute used for the MD simulation (e.g. number of target energy evaluations, runtime) in the appendix?\n\nWe will include more details in the appendix. Many simulation details are already included in appendix C.2.\n\n> The Equivariant flow matching method results in longer training times for LJ55 and alanine dipeptide, but no discussion of why / analysis of this is provided (see Questions). \n\nWe add a more detailed discussion in the final version, also exploring alternatives and the parallel way to prepare the training data, as discussed above and in the global rebuttal.\n\n> One of the headline results from the abstract (re alanine dipeptide) is obtained using an existing method (without the new proposed method).\n\nThe novel result of training a Boltzmann generator for molecules in Cartesian coordinates is obtained with both the equivariant OT flow matching and OT flow matching. See the discussion above as well as the global rebuttal. "}, {"review_id": "OneG8D2iGh", "reviewer": "Reviewer_coc5", "summary": "This paper studies equivariant flow matching, an extension to the flow matching paradigm for training continuous normalizing flows by regressing parametric vector fields to conditional vector fields. Building upon prior work, conditional vector fields can be derived from the distribution of the 2-Wasserstein optimal transport map $\\pi(x0, x1)$  between the prior $q(x0)$ and the target $\\mu(x1)$. The central contribution of this paper is to bake symmetries into the optimal transport cost $c(x_0, x_1)$ through a sequential search procedure for the symmetry group $S(N)$ and $SO(N)$. This is done by first applying the Hungarian algorithm followed by finding a rotation matrix that minimizes the cost, for each element in the cost matrix. Finally, the authors parameterize an equivariant vector field using the EGNN architecture of Satorras et. al 2021. Experiments are done on $n$-body particle dynamics and training a Boltzmann Generator for small proteins in Alanine dipeptide. ", "questions": "\n1. One key benefit of equivariant models is that they are more data efficient. This avenue seems to be underexplored. A table similar to Table 1 in $E(N)$-NF would be nice to see. \n\n2. There are many equivariant models for molecular data and simulation but only a small dataset in Alanine dipeptide was used. Can the authors justify not including a large body of equivariant generative models for molecular simulation here?\n\n3. While this paper claims equivariant flow matching as a general procedure the main method is limited to the groups $S(N)$ and $SO(D)$. This is a bit misleading as the current approach for finding the optimal cost matrix is not applicable to other groups. The authors may consider a different title for this work as well as a fairer presentation in the abstract and introduction of what is actually done.", "limitations": "Yes", "rating": 3, "confidence": 5, "soundness": 2, "presentation": 2, "contribution": 2, "strengths": "**Originality**\n\nThe main strength of this paper is that all presented material follows naturally from the desire to bake in equivariance in continuous normalizing flows. As equivariance has already been studied extensively in the generative modeling literature this application to flow matching is a reasonable extension. The originality of this work is limited to solely baking symmetries in the cost matrix as equivariant vector fields using EGNN have already been employed in the literature e.g. $E(N)$-normalizing flow (Satorras et. al 2022). \n\n**Quality**\n\nThe quality of the presented work is a good first attempt at tackling the problem of equivariant flow matching, but unfortunately, it is below the standard that is expected on several fronts which are outlined in the weaknesses section. \n\n**Clarity**\nIn general, the work is fairly clear. The presented ideas are straightforward to grasp but a few technical details are omitted which could improve understanding and readability. For example, how do you align the rotations after performing the Hungarian algorithm? We are minimizing over $SO(D)$, this is a non-trivial manifold. Of course, the appendix and code have this information but given that this is a crucial part of the contribution more detail would improve readability.\n\n**Signifance**\nAs equivariance in generative modeling has largely been studied in the literature the contribution of this work is limited. While equivariance has not been exclusively studied in flow matching, this extension is an early step and has limited novelty. This would be fine if there were large benefits empirically, theoretically, or computationally but this is not the case as far as I can tell and as a result, this paper has limited significance at present.\n", "weaknesses": "This paper has many potential weaknesses some of which are already alluded to in the previous section. \nFirstly, it has very little novelty as many of the concepts regarding equivariance and flows are known in the literature. In fact, using EGNN in normalizing flows for this very symmetry group has already been studied by Satorras et. al 2022. Moreover, one of the main benefits of flow matching is that one sidesteps the training complexity of regular CNFs as we do not need to backprop through an ODE solver. This brings huge computational benefits during training, albeit inference is still costly. In this paper, training is also expensive as solving for the optimal coupling in mini-batch OT is already expensive but even more so because the Hungarian algorithm is employed which scales $O(N^3)$. This is prohibitively expensive for any large-scale machine learning system. I encourage the authors to investigate other means of approximately solving for permutations. Some examples and directions for investigation can include learning the permutations (see Git Rebasin Ainsworth et. al 2022) using the straight-through estimator, or using Gumbel Sinkhorn (Mena 2018).\n\nWith regard to the proposed approach, the authors break down the problem by doing a sequential minimization by first finding the best permutation and then finding the best rotation. But this is not the original problem which may be indeed intractable. This is a ripe avenue to do a bit of theory to justify the proposed approach. For example, why is this sensible and not the first thing one can do? Can we do a little bit of error analysis to bound the error of the optimal cost matrix found in the sequential search to the actual one?\n\nThe experiments section is also rudimentary. While I appreciated the visualization of the shorter and straighter OT-Paths it is difficult to rationalize the gains here knowing that the overall algorithm is likely more expensive. A detailed time-complexity analysis of the overall method is needed here and proper comparison to regular flow matching and $E(N)$ NF is also a good idea. Moreover, the results seem a bit mixed as equivariant models do worse on Alanine dipeptide compared to non-equivariant models. The current explanation in the main text for this result is unsatisfactory. We expect equivariance to help here, why is it not?\n", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their detailed review and questions. To keep the rebuttal within the character limit, we often only cite parts of each question, but always answer the whole one.\n\n> How do you align the rotations after performing the Hungarian algorithm?\n\nWith the Kabsch algorithm, see also global rebuttal. \n\n> While equivariance has not been exclusively studied in flow matching, this extension is an early step and has limited novelty. ... this paper has limited significance at present ...\n\nWe respectfully disagree regarding the limited novelty and significance. Enhancing MD simulations using ML models is a rapidly growing field, as evidenced by the increasing number of AI4Science groups focused on this topic. Prior to our work, it was uncertain whether training a Boltzmann Generator in Cartesian coordinates to generate samples from the equilibrium distribution of real molecules was possible. We are the first to present a model and training algorithm that successfully achieve this. We believe our results are of great interest to the community and will stimulate further research in this area, potentially leading to transferable Boltzmann Generators.\nMoreover, we consider an even more difficult problem as we train the model with biased samples, instead of sample from the equilibrium distribution. We present two training algorithms that achieve this. Both including the known concepts of flow matching and EGNNs. However, we are the first to combine these, resulting in this novel finding. Additionally, we show that when scaling to larger systems with more symmetries, it becomes critical to do equivariant OT flow matching, to preserve the optimal transport sampling paths. Interestingly our work shows that for systems with less symmetries, e.g. fewer identical particles like DW4, LJ13, and alanine dipeptide, the benefits of including equivariant OT flow matching are marginal (see also appendix A.1). We will highlight these findings more in the final version of the manuscript.\n\n> ... training is also expensive as solving for the optimal coupling in mini-batch OT is already expensive ...\n\nWe agree that especially the reordering introduces additional computational cost, as shown in appendix C.3 and discussed in Sec. 8. We share the suggested approximations of the Hungarian algorithm, as detailed in Sec. 8.\nHowever, the training data preparation can be performed prior to or during the training and is trivially parallelizable on CPUs. Allowing to essentially train the equivariant flow matching models as fast as the OT flow matching models. See the global rebuttal for more details. \n\n> the authors break down the problem by doing a sequential minimization by first finding the best permutation and then finding the best rotation...\n\nWe discuss equivariant OT flow matching for more general symmetry groups in appendix B.1 in great detail. We will include the main findings in Sec. 4 in the final version.\nWe agree that this minimization problem is intractable, but we show that our approach is close to the correct solution. See global rebuttal for a detailed analysis.\n\n> The experiments section is also rudimentary...\n\nWe respectfully disagree that the experiments section is rudimentary. As mentioned, we are the first to train a Boltzmann generator for real molecules in Cartesian coordinates. Nevertheless, we added more experiments and additional evaluations, as shown in the global rebuttal. \nWe report runtimes of the training in appendix C.3. However, note that in practice we can do the data preparation in parallel before or during the training.\nWe already compare to $E(N)$-NF flows, also on their biased test set in appendix A.3. Note that we do not compare for LJ55 and alanine dipeptide, as likelihood training requires too much memory, as shown in appendix A.2. \nWe added comparison to normal flow matching as suggested.\n\n> the results seem a bit mixed as equivariant models do worse on Alanine dipeptide compared to non-equivariant models...\n\nThe difference for equivariant OT compared to normal OT is marginal for systems with few symmetries, i.e. less than 15 interchangeable particles, which is the case for DW4, LJ13 and alanine dipeptide. For alanine dipeptide, there are moreover different particle types. However, for the much larger LJ55 system (at least 3 times as many identical particles) the difference is significant. Crucially, the straight OT sampling paths allow using a Runge-Kutta integrator instead of the adaptive dopri5, which is significantly faster. With the “normal” OT the particle integration paths change directions (see Figure 2b,c,e), which requires small step sizes where the particles turn. Hence, this shows that it is important to use equivariant OT flow matching when scaling to larger systems to obtain optimal paths. \n\n> One key benefit of equivariant models is that they are more data efficient. This avenue seems to be underexplored...\n\nWe thank the reviewer for this suggestion, we included this experiment in the global rebuttal. \n\n> There are many equivariant models for molecular data and simulation...\n\nWe investigate standard benchmarking systems in the literature and even propose new ones, which were previously out of reach for Boltzmann generators in Cartesian coordinates. We leave it for future work to scale to even larger systems. However, we showed that in order to get OT integration paths, equivariant OT flow matching is required when scaling to larger systems.\n\n> While this paper claims equivariant flow matching as a general procedure the main method is limited to the groups $S(N)$ and  $SO(D)$...\n\nAs mentioned, the theory for more general symmetry groups is somewhat hidden in appendix B.1. However, the challenge remains of finding a good approximation to eq. 14 for other symmetry groups. Nevertheless, the investigated symmetry groups in our paper are the most important for molecular data. We will change the abstract as suggested."}, {"review_id": "dGdAn8UgzM", "reviewer": "Reviewer_nzPr", "summary": "The manuscript builds on recent progress in simulation-free loss functions for continuous normalizing flows which allow to scale CNFs to significantly larger dimensions and can be thought of as a generalizion of continuous time diffusion models. Specifically, the manuscript extends the recently proposed conditional flow matching approach which proposed to use results of optimal transport to construct particular efficient probability paths.\n\nA notable novelty of the paper is to consider the question of how to incorporate equivariance in the conditional flow matching loss. This is achieved by modifying the cost matrix c(x, x') of the Wasserstein distance to be the minimal cost *over the entire orbit* of x' (or equivalently x). \n\nThe proposal is studied experimental in the context of Boltzmann generators which are normalizing flows trained to act as sampling density for importance weighting (or, alternatively, Markov Chain Monte Carlo) of unnormalized physical target densities of Quantum Chemistry. Normalizing flows are particularly suited for this task as they provide a tractable likelihood as well as fast sampling. Equivariance is of pivotal importance for Boltzmann generators since the studied physical systems often have a high degree of symmetry, e.g. SE(3) and permutation symmetry. \n\nThe paper establishes in detailed numerical experiments that the flow matching approach is beneficial in the context of Boltzmann generators and compares favorably to likelihood-based training. This is shown for standard benchmarks such as Lennard Jones particles and Alanine dipeptide.  \n\nOverall, I think however this is a valuable contribution in a very exciting and rapidly evolving field of research. I therefore tend to recommend acceptance. I would however encourage the authors to add an appendix summarising the conditional OT flow matching procedure of 2302.00482 on which their method builds. Unfortunately, the conditional flow matching paper is not the most readable and it would make the ms for self-contained and accessible to add a brief summary.\n\n", "questions": "- How is the minimum over the SO(D) group orbit in Eq (14) actually performed (the group is continuous)?\n- What is the advantage of cartesian coordinates for alanine dipeptide as opposed to the standard parameterization in terms of dihedrals, bond angles, ...?\n- What is the precise definition of path length in Table 1?  ", "limitations": "Limitations are properly addressed.", "rating": 6, "confidence": 4, "soundness": 4, "presentation": 3, "contribution": 3, "strengths": "- The paper establishes, to the best of my knowledge, for the first time that flow matching is benefical in the context of Botzmann generators. It also presents a detailed comparison to both likelihood-based as well as OT conditional flow matching training objectives on standard benchmark sets.\n- The question of how equivariance is of high relevance as CNFs are widely deployed in physics applications for which symmetry is a fundamental ingredient for successful learning.\n- The presentation is well structured and clear\n", "weaknesses": "- The proposed method is rather specific to the case of SE(3) and permutation invariance. It seems non-trival to me how one would extend the treatment to larger symmetry groups, such as in applications to Lattice Field Theories. \n- A crucial element of Boltzmann generators is that they are often trained using self-sampled energy training (Variational Inference). The flow matching condition does not allow for such type of training. In the original conditional flow matching paper, a reweighting procedure was proposed. I think the authors made a good choice in not discussing this in this contexts as it is to be expected that reweighting will fail for reasonably sized systems such as alanine dipeptide. Nevertheless, the fact that (an efficient version of) energy-based training is not available for flow matching is a major downside of flow-matching.\n- I am a bit unclear on how much the proposed equivariant method actually helps. There is a substantial gain in the LJ55 setup while on Alanine dipeptide the previously proposed OT flow matching leads to higher ESS (at the cost of slightly longer trajectory length). \n\nMinor comments:\n\nP.3 L.80-81: Continuous time diffusion models provide a tractable likelihood. In fact, they are continuous normalizing flows albeit trained with a different objective as can be seen by noticing that the reverse process is equivalent to a deterministic ODE (known as probability flow ODE).\n\nP.3. L.100: distribution -> density\n\nP.3 Eq 3: I find the notation a bit contrived. Why not simply use f_\\theta(t, x)?\n\nP.4 L .111: Jacobian trace -> trace of Jacobian or divergence\n\nP.4 Eq 6: x_1 \\sim \\mu(x_1) -> x_1\\sim \\mu\n\nP.4 L. 137 \\sim q(x_0) -> \\sim q and \\sim \\mu(x_1) -> \\sim \\mu\n\nP.6 Table 1: Please mention in the caption how the errors are determined. Most readers will wonder about this while looking at the table so it would be good to provide this information already here.\n\nP.6 Eq 15-18: Various commas are missing after the equations\n\nP.6 L. 208: mean free Gaussian is not equal to \\mathcal{N}(x_0| 0, 1), as it has different normalizer. \n", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their insightful review, and appreciate their assessment that our paper is a valuable contribution in a very exciting and rapidly evolving field. We now address their comments individually.\n\n> Unfortunately, the conditional flow matching paper is not the most readable and it would make the ms for self-contained and accessible to add a brief summary.\n\nWe agree on the readability of the original flow matching paper. We will include a more detailed discussion both in Sec. 3.4 and in the appendix. We agree that this will make our work more cohere.\n\n> The proposed method is rather specific to the case of SE(3) and permutation invariance. It seems non-trival to me how one would extend the treatment to larger symmetry groups, such as in applications to Lattice Field Theories\n\nWe discuss equivariant OT flow matching for more general symmetry groups in appendix B.1 in great detail. We will include the main findings in Sec. 4 in the main part of the manuscript. However, how Eq. 14 can be efficiently approximated for other symmetry groups remains for future research. The $SO(3)$ and $S(N)$ symmetry groups (as well as the translation) are the most important ones for molecular systems, which our paper aims to tackle.\n\n> A crucial element of Boltzmann generators is that they are often trained using self-sampled energy training (Variational Inference). The flow matching condition does not allow for such type of training. In the original conditional flow matching paper, a reweighting procedure was proposed. I think the authors made a good choice in not discussing this in this contexts as it is to be expected that reweighting will fail for reasonably sized systems such as alanine dipeptide. Nevertheless, the fact that (an efficient version of) energy-based training is not available for flow matching is a major downside of flow-matching.\n\nWe agree that this is limitation of flow matching. However, it is not feasible to train CNF energy based, as shown in appendix A.2 for the larger systems we investigated. As the reviewer mentions, the proposed reweighting procedure in the flow matching paper is not feasible for larger systems. As the importance weights will be zero, when reweighting from a Gaussian to the target distribution. An alternative approach is to first train a CNF with equivariant OT flow matching, with only a few (potentially biased) samples and then generating new samples with the flow. As these can then be reweighted with non-zero weights, they can then be added to the training set. This could be done iteratively and the final training set would resemble samples from the target equilibrium distribution. We will include this idea in the future work section, as it is beyond the scope of this paper.\n\n> I am a bit unclear on how much the proposed equivariant method actually helps. There is a substantial gain in the LJ55 setup while on Alanine dipeptide the previously proposed OT flow matching leads to higher ESS (at the cost of slightly longer trajectory length).\n\nThe difference for equivariant OT compared to normal OT is only marginal for systems with few symmetries, i.e. less than 15 interchangeable particles, which is the case for DW4, LJ13 and ALA2. However, for the much larger LJ55 system (at least 3 times as many identical particles) the difference is significant. Crucially, the straight OT sampling paths allow using a Runge-Kutta integrator instead of the adaptive dopri5, which is significantly faster. With the “normal” OT the particle integration paths change directions (see Figure 2b,c,e), which requires small step sizes where the particles turn. Hence, this shows that it is important to use equivariant OT flow matching when scaling to larger systems to obtain optimal paths. This effect is not that prominent for the smaller systems, where normal OT flow matching already works quite well. We included further experiments in the global rebuttal that highlight the benefits of equivaraint OT flow matching.\n\n> How is the minimum over the SO(D) group orbit in Eq (14) actually performed (the group is continuous)?\n\nIt is performed with the Kabsch algorithm, we will include this in the final version of the manuscript. For more details see also the global rebuttal. \n\n> What is the advantage of cartesian coordinates for alanine dipeptide as opposed to the standard parameterization in terms of dihedrals, bond angles, ...?\n\nThere are multiple advantages of using Cartesian coordinates over internal coordinates. The main two are transferability and the long *robot arm problem*. (i) Transferability: Training a transferable Boltzmann generator will be easier in Cartesian coordinates, as these are not specific to each system. \n(ii) The long robot arm problem: Peptides and proteins often have different metastable states, which can for example be folded or unfolded. In these folded states some non-bonded parts will be close, but the information that these are close needs to be propagated though the whole backbone structure in the form of torsion angles, angles and bond lengths (long robot arm). Hence, a model in internal coordinates, as it does have the information of this distance implicitly, will fail to learn the distribution accurately. Moreover, going to explicit solvent systems, i.e. with water molecules, internal coordinates will not be possible and accounting for the large amount of possible permutation of interchangeable water molecules will be crucial. Highlighting again the importance of using our equivariant OT flow matching when scaling to larger systems, as shown in the LJ55 results. \n\n> What is the precise definition of path length in Table 1?\n\nThe path length as reported in the table and shown in the figures is the usual arc length between two points in $N \\times D$ dimensions. We will include this in the final version of the manuscript.\n\nWe thank the reviewer for their helpful minor comments and agree with all of them. We will change the final version accordingly."}]

Author rebuttal:
We thank the reviewers for their time reviewing our paper and for their insightful questions and suggestions. 
We present results for most of the suggested additional experiments and evaluations below and in the attached pdf. Moreover, we address some common questions and clarifications as well.

**The theory for equivariant flow matching** for more general symmetry groups was somewhat hidden in appendix B.1 We will move the main findings to Sec. 4 and highlight this more in the abstract as suggested. 

**Approximation of Eq 13** The approximation given in Eq. 14 is in practice performed using the Hungarian algorithm for permutations and the Kabsch algorithm for rotations. We compare several different other suggested approximations of Eq. 13 in Fig. 1c,d in the pdf. The baseline reference is computed with an expensive search over the approximation given in Eq. 14. Namely, we evaluate Eq. 14 for 100 random rotations combined with the global reflection, denoted as $O_{200}$, for each sample, i.e.
$\hat{c}(x_0, x_1)=\min_{o\in O_{200}(D)}\tilde{c}(x_0, \rho(o) x_1),$ where $\tilde{c}(x_0,  x_1)$ is given by Eq 14. 
Hence, this is 200 times more expensive than our approach. This baseline should be much closer to the true batch OT solution. 
The presented results for alanine dipeptide show that our approximation is simple, while also being a very good approximation, which we also show in our experiments in the pdf and the main paper. Applying our approximation multiple times reduced the transportation cost slightly. Performing the rotations first, lead to inferior results. We observe the same behavior for the LJ55 system.

**Parallel batch generation** In the current version of the paper, we perform the batch preparation during training.
However, a simple way to speed up the equivariant OT training is to generate the batches beforehand or in parallel to the training process. This process is highly parallelizable and can be performed on CPUs, which are in practice usually more available than GPUs and also in higher numbers. This also allows for larger batch sizes for the equivariant OT model and comes at little additional cost. Hence, scaling equivariant flow matching to even larger systems should not be an issue. We use this procedure for the new experiments and will mention this in the final version.

**Naïve flow matching baseline** We include naïve flow matching with the same equivariant architecture as an additional baseline as requested by multiple reviewers. Naïve flow matching results in even longer integration paths, as shown in Figure 1 and Table 1. The other results are close to the results of OT flow matching. 
Flow matching with a non equivariant architecture, i.e. a dense neural network, failed for all systems but DW4 and is hence not reported.

**Benefits of equivariant OT flow matching**
We show that when scaling to larger system sizes, e.g. the LJ55 system, equivariant OT flow matching is crucial to maintain optimal integration paths, which result in significantly faster sampling (Figure 1b) and allows the usage of fixed step integrators. Moreover, also the training is faster for the equivariant OT models, as they converge faster (Figure 1h). 

**Training set sizes**
We compare different training set sizes for alanine dipeptide and LJ55 in Table 2 in the pdf. Especially the integration paths lengths are the same for the different training set sizes. Note that we used $10^6$ training samples for all the results in the main paper. All other hyperparameters are the same. As observed in prior work, equivariant models are quite data efficient, which is also reproduced by our findings.  

**Batch sizes** 
We compare different batch sizes for the LJ55 system in Table 3 in the pdf. The integration paths lengths are again similar across the different batch sizes. However, smaller batch sizes resulted in better likelihoods for both flow matching models.

**Simulation details**
We prepare the training data before training as described above. 
We will use more runs for the error calculation and more samples to estimate ESS in the final version.
The reported ESS might be misleading as we do not generate many samples and the negative log likelihoods are higher than observed for the larger training set sizes. Hence, some states might be missed by the model for the LJ55 system.  

We will merge Table 1 in the pdf with Table 1 in the paper. Table 2 and 3 in the pdf will be added to the appendix, while the results will be discussed in the main part. Table 4 in the pdf will replace Table 2 in the paper. 

We thank the reviewers for their valuable suggestions and hope for an engaging discussion period.


Meta-review: This paper extends Flow Matching to equivariant flows. The paper is well written. It also uses symmetry to reduce generation path length and find better coupling of noise and data for training. The main issues in the paper are: It has somewhat limited baselines, missing time-complexity (e.g., in the context of solving eqs. 13,14), marginal improvement of symmetric-aware sampling over independent coupling of noise and data, missing motivation for the MCMC sampling as a first step to FM training, missing link of transport cost and faster sampling, and unified contribution claims of equivariant FM and better symmetry-induced coupling. During rebuttal the authors agreed to address these issues, and incorporate relevant changes in the camera ready version.
Decision: Accept (poster)