SE(3) Equivariant Augmented Coupling Flows
Authors: Laurence Illing Midgley
Vincent Stimper
Javier Antoran
Emile Mathieu
Bernhard Schölkopf
José Miguel Hernández-Lobato
Venue: NeurIPS 2023 spotlight (2023)
Keywords: Boltzmann generator
normalizing flow
diffusion
molecular dynamics
TL;DR: We propose SE(3) equivariant augmented coupling flows for modeling physical systems.

Abstract: Coupling normalizing flows allow for fast sampling and density evaluation, making them the tool of choice for probabilistic modeling of physical systems. 
However, the standard coupling architecture precludes endowing flows that operate on the Cartesian coordinates of atoms with the SE(3) and permutation invariances of physical systems. 
This work proposes a coupling flow that preserves SE(3) and permutation equivariance by performing coordinate splits along additional augmented dimensions. At each layer, the flow maps atoms' positions into learned SE(3) invariant bases, where we apply standard flow transformations, such as monotonic rational-quadratic splines, before returning to the original basis.
Crucially, our flow preserves fast sampling and density evaluation, and may be used to produce unbiased estimates of expectations with respect to the target distribution via importance sampling.
When trained on the DW4, LJ13, and QM9-positional datasets, our flow is competitive with equivariant continuous normalizing flows and diffusion models, while allowing sampling more than an order of magnitude faster.
Moreover, to the best of our knowledge, we are the first to learn the full Boltzmann distribution of alanine dipeptide by only modeling the Cartesian positions of its atoms.
Lastly, we demonstrate that our flow can be trained to approximately sample from the Boltzmann distribution of the DW4 and LJ13 particle systems using only their energy functions.

SE(3) Equivariant Augmented Coupling Flows
Laurence I. Midgley" University Cambridge lim240cam. ac ac.uk
Vincent Stimper" Max Planck Institute for Intelligent Systems University of Cambridge
v84880cam.ace ac.uk
Javier Antorán" University of Cambridge j ja6660cam.c ac. uk
Emile Mathieu
Bernhard Schölkopf Max Planck Institute for Intelligent Systems bs@tue. mpg.
University of Cambridge ebm320cam ac. uk
José Miguel Hernández- University of Cambridge jmh2330cam ac
Abstract
Coupling normalizing flows allow for fast sampling and density evaluation, making them the tool of choice for probabilistic modeling of physical systems. However, the standard coupling architecture precludes endowing flows that operate the Cartesian coordinates of atoms with the SE(3) and permutation invariances of physical systems. This work proposes coupling flow that preserves SE(3) and per- mutation equivariance by performing coordinate splits along additional augmented dimensions. At each layer, the flow maps atoms' positions into learned SE(3) invariant bases, where we apply standard flow transformations, such as monotonic rational -uanraate splines, before returning to the original basis. Crucially, our flow preserves fast sampling and density evaluation, and may be used to produce unbi- ased estimates of expectations with respect to the target distribution via importance sampling. When trained on the DW4, IJ13, and QM9 9M9-positional datasets, our flow is competitive with equivariant continuous normalizing flows and diffusion models, while allowing sampling more than order of magnitude faster. Moreover, to the best of our knowledge, we are the first to learn the full Boltzmann distribution alanine dipeptide by only modeling the Cartesian positions of its atoms. Lastly, we demonstrate that our flow can be trained to approximately sample from the Boltzmann distribution of the DW4 and LJ13 particle systems using only their energy functions.
1 Introduction
Modeling the distribution of a molecule's configurations at equilibrium, known as the Boltzmann distribution, promising application of deep generative models [Noé et al., 2019]. While the unnormalized density of the Boltzmann distribution can be obtained via physical modeling, sampling from typically requires molecular dynamics (MD) simulations, which are expensive and produce correlated samples. A promising alternative is to rely on surrogate deep generative models, known as Boltzmann generators. We can draw independent samples from these and debias any expectations
"Equal contribution
Code: https://githin http://tttta..ccoooaaaaaaaaaaata
37th Conference on Neural Information Processing Systems (NcurIPS 2023).

---

estimated with the samples via importance weighting. The Boltzmann distribution typically admits rotation and translation symmetries, known as the Special Euclidean group SE(3), as well as per- mutation symmetry. These constraints are important to incorporate into the model as they improve training efficiency and generalization [Cohen and Welling, 2017, Batzner et al., 2022, Köhler et al., 2020]. Other key desiderata of Boltzmann generators are that they allow for fast sampling and density evaluation. These are necessary for energy- beeegy-based training. that is, training using the Boltzmann distribution's 's unnormalized density [Noé et al., 2019, Stimper et al., 2022, Midgley et al., 2023]. Training by energy is critical, as prevents the model quality from being constrained by the quality and quantity of MD samples.
Existing coupling flows which approximate the Boltzmann distribution of molecules are at least partially defined over internal coordinates, i.c. bond distances, bond angles, and dihedral angles [Wu et al., 2020, Camphell et al., 2021, Köhler et al., 2023a, Midgley et al., 2023], which are SF((3) invariant. However, the definition of these depends on the molecular graph and they are non-unique for most graphs. Furthermore, models parameterized in internal coordinates struggle to capture interactions among nodes far apart in the graph and cannot capture the permutation invariance of some atoms. Thus, they are not suitable for particle systems such as LJ13 [Köhler et al., 2020]. SE(3) equivariance constraints have been applied to continuous normalizing flows (CNFs) operating on Cartesian coordinates [Köhler et al., 2020, Satorras et al., 2021a]. and to the closely related diffusion models [Xu et al., 2022, Hoogeboom et al., 2022, Yim et al., 2023]. These models are built upon SE(3) equivariant graph neural networks (GNNs) [Satorras et al., 2021b, Geiger and Smidt, 2022, Batzner et al., 2022]. These architectures can be applied to any molecular graph [Jing et al., 2022], enabling a single generative model to generalize across many molecules. Alas, sampling and evaluating the density of CNFs and diffusion models typically requires thousands of neural network evaluations [Xiao et al., 2022], preventing them from being trained by energy. As such, presently no Boltzmann generator exists that (i) acts on Euclidean coordinates of atoms (ii) enforces SE(() equivariance, and (iii) allows for fast sampling.
To address this gap, we propose a flexible SE(3) equivariant coupling flow that operates on the Carte- sian coordinates of atoms, allowing for fast sampling and density evaluation. Our contributions are:
We extend coupling layers to be SE(3) equivariant by augmenting their input space with auxiliary variables [Huang et al., 2020] which can be acted upon on by SE(3). We update the atom positions conditioned on the auxiliary variables by first projecting the atoms into an SE(3)- (B33)-invariant space and then applying standard normalizing flow transform before projecting its output back onto the equivariant space.
We demonstrate that, when trained by maximum likelihood, our flow matches the performance of existing SE(3) CNFs and diffusion models as well as coupling flows operating on internal coordinates on molecular generation tasks. Our flow is more than 10 times faster to sample from than SE(3) CNFs and diffusion models. Concurrently with Klein et al. [2023b], we are the first to learn the full Boltzmann distribution of alanine dipeptide solely in Cartesian coordinates.
We demonstrate our flow in the energy-based training setting on DW4 and LJ13, where parameters are learned using only the molecular energy function. Energy-based training of CNFs or diffusion models intractable due to slow sampling and density evaluation. Flows operating on internal coordinates are not able to capture the permutation invariance of these problems. Hence, our flow is the only existing permutation and SE(3) equivariant method that can tractably be applied there.
2 Background: coupling flows and invariant models
2.1 Normalizing flows and coupling transforms
A (discrete-time) normalizing flow is flexible parametric family of densities X' as the push- forward of a base density go along an invertible automorphism fo X + X with parameters E θ [Papamakarios al., 2021]. The density is given by the change of variable formula:
qo (x) = 40 ( det
af, ((x) dz
We can efficiently sample from the flow by sampling from go and mapping these samples through fo in single forward pass. A popular way to construct fo is to use coupling transforms. The D

---

dimensional input a E X' is split into two sets, transforming the first set conditional on the second, while leaving the second set unchanged:
Yidd = Ta
(2)
I:+1:D = B.+1.D
They induce lower triangular Jacobian, such that its determinant becomes (0T( (a0e11/D/00mmp Further, choosing T to have an easy to compute determinant, such elementwise transformation [Dinh et al., 2015, 2017, Durkan al., 2019], allows for fast density evaluation and sampling at low computational cost.
2.2 Equivariance and invariance for coupling flow models of molecular conformations
Throughout, we deal with observations of an n-body system represented by a matrix æ [x ,.',..,"""' E X 3 where the rows index Cartesian coordinates and the columns in- dex individual particles. We seek to construct flows on X' endowed with the symmetries present in molecular energy functions. These are invariant rotations and translations of æ (SE(3)), and to permutation of atoms of the same type (S,). We will formalize them in this section.
Symmetry groups The special Euclidean group SE(3) is the set of orientation preserving rigid transformations in Euclidean space. Its elements L E SE(3) can be decomposed into two components t t= (R, u) where R. SO(3) is 3 x 3 rotation matrix and coordinate V E R2, t = Rv + u denotes the action of t on v. The symmetric group S,, defined over set of atoms consists of all n!! permutations that can be performed with said atoms. Its elements E S, act on an n-body system
R represents a translation; for
x (e.1))
Equivariant maps A map f X + Y' said to be G-equivariant if it commutes with the group action i.e. if for any x E X' and g E G we have f(g x) 9 fz)).. Invariance is special case where for any x E X' and g G we have f(g- 2) f(x). There has been a plethora of recent work on constructing graph neural network functions equivariant to the action of G SE(3) x S, [e.g. Thomas et al., 2018, Satorras et al., 2021b, Geiger and Smidt, 2022], which we will leverage to construct our equivariant coupling flow model.
Invariant density A density p X - R is G-invariant if for any E X' and g G have ==px)) Combining an invariant base density go with an equivariant invertible transform f. as in (I), yields an invariant flow density (Papamakarios et al., 2021, Köhler et al., 2020J. This gives practical way to design invariant densities models.
Challenges in constructing SO(3) x Sn, invariant flow models Unfortunately, no coupling transform can be simultancously equivariant to both permutation of the particles and their rotations; coupling splits must be performed either across particles spatial dimensions which would break either permutation or rotational symmetry [Köhler et al., 2020, Bose et al., 2022]-
Furthermore, there does not exist translation invariant probability measure, as any such measure would be proportional to the Lebesgue measure and therefore not have unit volume. This precludes us from defining invariant base distribution directly on X. Fortunately, Proposition A. and its converse allow us to disintegrate the probability measure into a translational measure proportional to the Lebesgue measure and an SO(3)-invarian probability measure on the subspace of R33n with zero center of mass. We can drop the former and only model the latter.
3 Method: SE(3) X Sn, equivariant augmented coupling flow model
This section describes our main contribution, an SE(3) x S, equivariant coupling flow. We first lay the groundwork for achieving translation invariance by defining our flow density on a lower- dimensional "zero Center of Mass (CoM)" space. To preserve permutation and rotation equivariance, we leverage the augmented flow framework of Huang al. [2020]. Specifically, we use sets of augmented variables as pivot for coupling transforms. Sec. 3.1 introduces novel class of coupling transforms that achieve the aforementioned permutation and rotation equivariance by operating on atoms projected using set of equivariant bases. Sec. 3.2 describes our choice of invariant base distribution and, finally, in See. 3.3, we discuss several schemes to train the augmented flow from either samples or energy functions and how perform efficient density evaluation.
3

---

Shift CoM map
a
â
Equivariant Neural Network
Equivariant
paramettoss
Projection
Flow trans- form тe
Inverse pro- jection "r
x
Yr
Core transformations Fi
x xM
Figure 1: Illustration of the equivariant coupling layer of our augmented normalizing flow, where our variable with zero center of mass (CoM) is transformed with the augmented variable a.
Translation invariance is obtained by modelling the data on the quotient space RR / X, where all n--ody systems that only differ by translation are "glued" together, i.c. where æ x' if x'' + p with p E R Constructing parametric probabilistic model X, automatically endowing it with translation invariance. In practice, we still work with Cartesian coordinates, but center the data so as to zero its CoM: ă - x Φ with æ A E4E [z] Thus, ă lies on X, an (n 10-11. 1) x 3 dimensional hyperplane embedded in X.
Augmented variable pivoted coupling To allow for simultancously permutation and rotation equivariant coupling transforms, introduce augmented variables E A. Our coupling layers update the particle positions x conditioned on a and vice-versa. The augmented variables need to "behave" similarly to in that they can also be acted upon by elements of SE(3) Sn, We achieve this by choosing a to be set of k of observation-sized arrays A = X which we will discuss further in App. B.4. Importantly, we do not restrict a to be zero-CoM.
Invariant flow density on the extended space We parameterize a density q over the extended space X x A w.r.t. the product measure X @ λA where Ax E P()) and Y E P(A) respectively denote the Lebesgue measure X' and A. We use q( a) shorthand for the density of the corresponding zero-CoM projection q a). The density is constructed to be invariant to the action of G SO(3) x S, when simultaneously applied to the observed and augmented variables, that is, q(x, a) = (g g-a) for any g E G. We aim to construct G-equivariant flow f f: X x + X X A on this extended space, and combine it with X' A + R+ a G-invariant base density function, yield the invariant flow density
q( = 4ollf (z. a)) det o ((..a) 0(x,a))
Proposition 3.1 (Invariant marginal), Assume q X x + R4 is G-invariant density over the probability space (X A, Ax & Aa), then qu - SA u. a))a(da) X + R, is G-invariant density w.r.. to the measure Xx-
Proof. For any G. X' and A
(g - x) - / g0g x. x.)))a(da)=
-/,-
9-1 qlg x.--x.9- x, (og.a)aaldg a) - -/ qτx a)da(da) = ga (x))
where we used the G-invariance of g and of the measure AA, as well as 1 = A.
3.1 SE(3) and permutation equivariant coupling transform
We now derive our SE(3) x S, equivariant map f : X x A x A defined on the extended space. We introduce two modules, a Shift-CoM transform, swapping the center of mass between the observed and augmented variables, and an equivariant core transformation, which updates ă conditional on a and vice versa. Composing them yields our equivariant coupling layer illustrated in Fig. 1.
SE(3) X Sn equivariant coupling We dissociate the equivariance constraint from the flexible parameterized flow transformation by (1) projecting the atoms' Cartesian coordinates into a learned,

---

local (per-atom) invariant space, (2) applying the flexible flow to the invariant representation of the atoms' positions, and (3) then projecting back into the original Cartesian space, Specifically, we construct a core coupling transformation that composes (a) an invariant map T E x A' - y, where y is isomorphic to X, and r E parametrizes the map. We denote the parametrized map as Yr. is followed by (b) a standard flexible normalizing flow transform, e.g. neural spline, r T - y, and (c) the inverse map y Denoting the inputs with superscript & and with E + for outputs, our core transformation F: (a'; aa))a²) + (x*+1, aa=1) ) is given as
xr s+1 7 Te(or æ²), with (r.0) = h(a'),
a c+1 = a.
Here, h is a (graph) neural network that returns a set of equivariant reference vectors r, which parametrize the map Yr. and invariant parameters B. y a rotation invariant space. This means that any rotations applied to the inputs will be cancelled by Yr, i.e. "gr ?r- g or equivalently (7g-r) = 9 2F for all E SO(3). We use the inverse projection 7r- to map the invariant features back equivariant features. 'he function TO is a standard invertible tiner-rranffommaionn such an affine or spline based transform, that we apply to the invariant features [Papamakarios et al., 2021].
We now show that the above described transform is rotation and permutation equivariant.
Proposition 3.2 (Equivariant augmented coupling flow). If A + X" x O" is -euii-equvvar for its first output, (iivaaaanra for its second, and S., equivariant for both, and (Tg.r) = =9-7, for all g E 5O(3), the transform F given by (4) is SO(3) x Sn, equivariant.
Proof. For SO(3): We first notice that h(g a) = (g r, rr.). and then since (1g.)) = =9.77" we
-.----00 a) have F(g- g.g-
= (Tar) -1 To( g- x²) = g 7, (1 9 g x²) = 9- [g-))---aa (F((xx))) a).
For Sn We first note that h(o a) = (o. a. 0). Then, using that o and T act on x atom- wise, we
have F(a-zz a.) You -o- (0) (a. ( (o-T))((o (a.z) a F a). D
For the Jacobian of the coupling described above to be well-defined, the variable heing transformed must be non-zero CoM (see App. B.1 for derivation). Thus, although our observations live on X, for now, assume that the inputs the transform assumption in the following paragraphs. This choice also allows us to use standard equivariant GNNs for h [Satorras et al., 2021b, Geiger and Smidt, 2022] which leverage per-node features defined in the ambient space, such as atom type and molecular graph connectivity.
not zero-Col) and will deal with this
Choices of projection Y The equivariant vectors parameterize a local (per-atom) SO(3) equiv- ariant reference frame used in the projection "/. We introduce three different projection strategies. (i) The first strategy, is for r to parameterize frame composed of origin and orthonormal rota- tion matrix into which we project each atom's 's positions. We then take T to be a dimension-wise transformation for each of the projected atoms' coordinates. We dub this method CARTESIAN- (ii) Alternatively, we let parameterize a origin, zenith direction and azimuth direction for spherical coordinates, in Liu et al. [2022]. We then apply elementwise transforms to each atom's 's radius, polar angle and azimuthal angle. We call this SPHERICAL-PRON (iii) Lastly, consider a variant of SPHERICAL-PRO where just the radius is transformed and the polar and azimuth angles are held constant. Here, r parameterizes a single reference point, a per-atom origin. We refer to this last variant as VECTOR-PROJ.
Architectural details For the transformations applied in the invariant projected space we consider affine mappings [Dinh et al., 2017] and monotonic rational-quadratic splines [Durkan et al., 2019]. Additionally, to limit computational cost, we have our GNNs h output M sets of reference vectors r and invariant parameters 0. These parametrize M core coupling transformations with single GNN forward pass. For the CARTESIAN-PROT and SPHERICAL-PROJ flow we include a loss term that discourages certain reference vectors from being collinear, which improves the projection's stability. We provide further details for this and the various projection types in App. B.3.
Center of mass shift The shift-CoM transform allows us to apply the aforementioned SF(3) x S, cquivariant coupling in the ambient space rather than zero-COM subspace. In particular, before transforming our observed vector ã E X with F, we lift it onto X. We achieve this by swapping the center of mass between w and a. For now, assume A = X, i.e. k = 1, with App. B.4 providing details for k > 1. Letting A C A be the subspace where all augmented variables that differ by
5

---

translation occupy the same point, and à E A be defined analogously to E, we apply the map ShiftCoM : X' x A + X' x A which acts on both of its arguments by subtracting from each of them the latter's CoM, that is,
(hiftCoM( a) (a - a, a- a) a) with -
(5)
This operation is invertible, with inverse Shift w), and has unit Jacobian determinant.
Algorithm 1:
Flow block f
Algorithm 2: Joint density evaluation
Inputs: Zero-CoM observation á, augmented
Inputs: (ar, a) p. base density qo. (
variable a, Coupling transforms Fi,
(z a (0) +- (æ a, x)
(x, ā) <- Shift a)
logdet
(æ. â) < FMl ...... F1 (x, a)
for I =1=1.... 1,.... L do
(H a) (0(111)) (1-1))
(a, a) < ShiftCoM a. a)
logdet + logdet + a ((--
(a, a) < F
(a,
Output: ã, a
Output: (o( a (2) logdet
Putting the building blocks together Our flow transform is built sequence of L blocks. Each block, described in Alg. I, consists of two equivariant coupling layers, see Fig. 1. Our observations
á E L are lifted onto X' with ShiftCoM, they are transformed with M core transformations (T(1)) i-1 and Shift.CoM is applied one more time map the observations back to the zero-CoM hyperplane. After this, our augmented variables a are transformed with (F(²)) M
==1'
Joint density evaluation g( a) is performed with Alg. 2. We first subtract the observed variables' CoM from both the observed and augmented variables. We then apply our L flow transform blocks before evaluating the transformed variables' density under our base distribution qu. which is described next. The log determinant of the core flow transform, f. has a contribution from the projection, transform in the invariant space, and inverse projection (see App. B.3 for details).
3.2 S3)3) x Sn. Invariant base distribution
Again, we assume X, i.e. k = 1, with the generalization given in App. B.4. Our invariant choice of base distribution is (x, a) = (xx; 0, 7) (a; x, 1 where E 3R and a E R refer to a and a flattened into vectors, 7 is a hyperparameter and we denote Gaussian distributions on X as N [Satorras et al., 2021a, Yim et al., 2023] with density
(6) We sample from it by first sampling from a standard Gaussian (N I) and then removing the CoM. On the other hand, the distribution for a is supported on A which includes non-zer ourzerr CoM points. centered æ, yielding joint invariance to translations. The isotropic nature of go makes its density invariant to rotations, reflections, and permutations [Satorras et al., 2021a, Yim et al., 2023].
N(: 0, = (2m)-3(nn1))) (xx=) - 11
3.3 Training and likelihood evaluation
In this section, we discuss learning and density evaluation with augmented variables.
Invariant augmented target distribution We assume the density of our observations p SE(3) x S,, invariant. Our target for augmented variables is π (a)ar) = N(a; ær, 7 I), where o/2 matches the variance of the base Gaussian density over a. This satisfies joint invariance p(g x x)){(g alg x)) p(æ)T(a|az for any g E SE(3) x Sn. as shown in App. B.5.
Learning from samples When data samples by maximizing the joint likelihood, which is a lower bound on the marginal log-likelihood over observations up to fixed constant
p are available, we train our flow parameters
E (erplat). [log g( a)] S
+ C.
Learning from energy When samples are not available but we can query the unnormalized energy of a state Uz),, with p(x) x exp(-U(z)), we can minimize the joint reverse KL divergence. By

---

the chain rule of the KL divergence, this upper bounds the KL between marginals
(8) However, the reverse KL encourages mode-seeking [Minka, 2005] which may result in the model failing to characterize the full set of meta-stable molecular states. Therefore, we instead use flow annealed importance sampling bootstrap (FAB) [Midgley al., 2023], which targets the mass covering a-divergence with a = 2. In particular, we minimize the a-divergence over the joint which leads to an upper bound on the divergence of the marginals
DKL. (1(2, a) (()))(()))) 2 DRr. (g(æ)) [()))
D2 (g(2,, a) [| A p(a))m(alaa) da dæ 2 m(a) dæ A (2(a())) (9) g(2.a)
q(=)
To compute unbiased expectations with the augmented flow we rely on the estimator Ex(a)]((=)) agg [w( (z)) where w(x, a) ((z))((lz)//q(a) a). Minimizing the joint a-divergence with a = 2 corresponds to minimizing the variance in the joint importance sampling weights w(x a),
which allows for the aforementioned expectation to be approximated accurately.
Evaluating densities To evaluate the marginal density of observations we use the importance
weighed estimator g(z) (a)a(
<[a]=]] noting that is Gaussian and thus supported every-
g(a.a)
where. The estimator variance vanishes when (a(xz) = ग T (a)zz) as shown in App. B.9.
4 Experiments
4.1 Training with samples: DW4, LJ13 and QM9 positional
First, we consider problems that involve only positional information, with no additional features such atom type or connectivity. Thus, the target densities are fully permutation invariant. The first two of these, namely DW4 and LJ13, are toy problems from Köhler et al. [2020], where samples are obtained by running MCMC on the 4 partiele double well energy function (DW4) and 13 particles Leonard Jones energy function (LJ13) respectively. The third problem, i.c. QM9 positional [Satorras al., 2021a], selects the subset of molecules with 19 atoms from the commonly used QM9 dataset [Ramakrishnan et al., 2014] and discards their node features.
For our model, Equivariant Augmented Coupling Flow (E-ACF), we consider all projection types VECTOR-PROL CARTESIAN PROJ, SPHERICAL-PROI) and compare them to: (1) NON-E-ACF: An augmented flow that is not rotation equivariant but is translation and permutation equivariant, in [Klein et al., 2023a]. This model uses the same structure as the E-ACF but replaces the EGNN with a transformer which acts directly atom positions, without any projection. We train NON-E-ACH with antaettgmenttiion whereby we apply a random rotation to each sample within each training batch. (2) E-CNF ML: The SE(3) equivariant continuous normalizing flow from Satorras et al. [2021a] trained by maximum likelihood. (3) E-CNF FM: An SE(3) equivariant continuous normalizing flow trained via flow matching [Lipman al., 2023, Klein al., 2023b]. (4) E-CNF-DUFF An SE(3) equivariant diffusion model [Hoogeboom et al., 2022] evaluated as a continuous normalizing flow. All equivariant generative models use the SE GNN architecture proposed by Satorras et al. [2021b]. The CARTESIAN- PRROOI exhibited numerical instability on 1M99positional causing runs to crash early. To prevent these crashes was trained at a lower learning rate than the other E-ACF models. App. C provides a detailed description of these experiments.
On Tab. 1, we that DW4 and IJ13,, E-CNF FM performs best, while E-ACF is competitive with E-CNF MI. and E-CNF-DIFF. We note that both DW4 and IJJ13 have biased datasets (see App. C.3.I). On QM9-posiiional, the E-ACF is competitive with E-CNF-DIFF and E-CNF FM, while the under-trained E-CNF ML from Satorras et al. [2021a] performs poorly. We expect that with further tuning E-CNF-DIN and E-CNF FM would match the performance of SPHERICAL-PROJ on QM9 -o9-positional. The NON-E-ACE performs much worse than the E-ACF, despite being trained for more epochs, demonstrating the utility of in- build equivariance. Furthermore, Fig. 2 shows that the distribution of inter-atomic distances of samples from our flow matches training data well. Importantly, sampling and density evaluation of the E-ACF on an A100 GPU takes roughly 0.01 seconds. For the CNF trained by flow matching (E-CNF ML) and score matching (E-CNF-DIFF), sampling takes on average 0.2, and seconds respectively. Thus, the E-ACF is faster for sampling than the CNF by more than an order of magnitude.
7

---

Table I: Negative log-likelihood results for flows trained by maximum likelihood on DW4, LJ13 and positionnl.... E-CNF ML results are from Satorras et al. [2021a]. Best results are emphasized in bold. The results averaged over seeded runs, with the standard error reported as uncertainty.
DW4
LJ13
QM9 positional
E-CNF MI. E-CNF FM E-CNF-DIFE NON-E-ACE
8.15 + N/A
30.56 + N/A
-70.2 N/A
6.31 +
28.81 + 0.34
-143.52 0.26
8.00
31.02
158.30
10.07
33.32 0.33
-76.76
VECTOR-PROI E-ACF
8.603
0..19
-152.23
CARTESIAN- B-ACF
8.00
30.89
-188.62
SPHERICAL- E-ACF
8.61
30.16
-165.71 + 1.35
DW4
LJ13
QM9 Positional
normalized count:
0.06
flow
data
0.02
0.00
2
6
2
3
2
interatomic distance
interatomic distance
interatomic distance
Figure 2: Inter-atomic distances for samples from the train-data and SPHERICAL-PROI E-ACE.
4.2 Training with samples: Alanine dipeptide
Next, approximate the Boltzmann distribution of alanine dipeptide in implicit solvent at temperature T' 800 K. We train the models via maximum likelihood on samples generated by replica exchange MD simulation [Mori and Okamoto, 2010], which serve as ground truth. Our generated dataset consists 10th samples in the training and validation set as well as 10? samples in the test set. Besides our E-ACF using the different projection schemes that introduced in Sec. 3.1, we train the non-SO(3) equivariant flow -NON-E-ACH previous experiments. Moreover, train flow internal coordinates in Midgley et al. [2023]. The joint effective sample size (ESS) is reported for the E-ACF models which a lower bound of the marginal ESS (see Eq. (9)). For the CNFs, the Hutchinson trace estimator is used for the log density [Grathwohl et al., 2018]. This results in a biased estimate of the ESS, which may therefore be spurious. Further details about model architectures, training and evaluation are given App. C.3.2.
with data augmentation similarly to the
Table 2: Alanine dipeptide results. KLD the empirical KLD of the Ramachandran plots (see Fig. 3). Forward ESS is estimated with the test set. Reverse ESS is estimated with 10y model samples. Results are averaged over 3 seeded runs, with the standard error reported as uncertainty.
KLD
NLL
Rev ESS (%)
Fwd ESS (%)
Flow on internal coordinates
(2.01 0.04)
10 190..15
0.02
1.61 0.03
E-CNT FM
(3.83100.18)
-0
190..20 0.09 0..093.11... 10- -2
(2310.8)
To 125
B-CNF-DIFF
(8.86 + 0.49)
10
-188.31 0.01
10
NOX-E-ACF
(1.66 + 0.01)
-00 -184.57
0.35
0.14
(..5
10
VECTOR-PROI E-ACF
(6.15 + 1211)
10 -188.56
0.01
19.4
79.4
CARTESIAN B-ACF
(3.46 + 0.28)
10 -188.59
0.00
52.5
7..7
SPHERICAL-PROI E-ACF
(2.55 + 0.29)
-18..57 0.00
7..
(5.0 4.1) 10
10-1
The results are shown in Fig. and Tab. All variants our E-ACF clearly outperform the NON-E- ACF, as the Kullback Leibler divergence (KLD) of the Ramachandran plots is significantly lower and the NI.L as well as the reverse and forward ESS, see App. 10, are higher. The flow trained on internal coordinates is only marginally better regarding the NLL and KLD than our best models, i.c.

---

Ground trith
Flow on intermal coordinates
E-CNP
E-CNE-
10
0
Non-EACCP with data ong.
Voctor-proj E-ACF
Curtesiun-gnoj E-ACF
Sphorical-poo E-CP
10
0
10-
17
Figure 3: Ramachandran plots, i.e. marginal distribution of the dihedral angles and &b (see
App. C.3.2), obtained with MD (ground truth) and various models.
the CARTESIAN- PRIOO and SPHERICAL-PROJ E-ACF, while we outperform it considering the ESS. Note that the flow internal coordinates explicitly models the angles ø and v. while the E-ACF operates on the underlying Cartesian coordinates. The E-ACFs outperform the E-CNF-DUFT model in all metrics, but the E- E-CNF trained with flow matching has slightly lower NLL while having significantly higher KLD on the Ramachandran plot. This could be due to a better performance on other marginals on some outlier data points in the test set. The forward ESS is very low for all models, which suggests that the models do not cover some regions in the target distribution, and that the reverse ESS is spurious. Alternatively, this may be from numerical instabilities in the models. To the hest of our knowledge, our models are the first to learn the full Boltzmann distribution of molecule purely Cartesian coordinates while being competitive with flow trained internal coordinates. Additionally, we train the best E-ACF model at 7 = 300 K where it also performs well (see App. C.3.2).
4.3 Energy- based training: DW4 and LJ13
Lastly, we demonstrate that our proposed flow can trained on the DW4 and LJ13 problems using only the target's unnormalized density with the FAB algorithm [Midgley al., 2023]. The annealed importance sampling procedure within FAB requires sampling from the flow and evaluating its density multiple times. This is used within the training loop of FAB making it significantly more expensive per parameter update than training by maximum likelihood. Given that sampling and density evaluation with CNFs is very expensive, training them with FAB is intractable. Thus, only report results for our flow, as well as for the -CFFE-ACF, We train the NoN-E-ACF for more iterations than the E-ACF, such that the training times are similar, given that the NON-E-ACF is faster per iteration. App. C.4 provides further details on the experimental setup.
Table 3: Results for training by energy with FAB. Best results are emphasized in bold. Results are averaged over 3 seeded runs, with the standard error reported as uncertainty. Reverse ESS is estimated with 10a samples from the flow. Forward ESS is estimated with the test sets.
DW4
LJ13 Fwd ESS (%)
Rev ESS (%)
Fwd ESS (%)
NLL
Rev ESS (%)
NLL
NOF-F-ACF
2.63
5.45
7.38 0.01
5.38 3.66
4.14
33.22 0.96
VECTOR-PROI B-ACF
84.29 0.41
$3.39.10.79
7.11
19.60
65.20
30.33
CARTESIAN- E-ACF
82.44
0.50
80.08 0.664
7.13
60.68 10.41
65.54
30.34
SPHERICAL-PROI E-ACF
80.44 0.88
81.46 0.995
7.14
0.776
66.13 0.11
30.21 0.02

---

Tab. 3 shows that the E-ACF trained with FAB successfully approximates the target Boltzmann distributions, with reasonably high joint ESS, and NLL comparable to the flows trained by maximum likelihood. Additionally, the ESS may be improved further by combining the trained flow with AIS, this is shown in App. C.4. In both problems the NoN-E- E-ACF performs worse, both terms of ESS and NLL. All models trained by maximum likelihood have much lower ESS (see App. C.3.1). 1). This is expected, as unlike the a = 2-divergence loss used in FAB, the maximum likelihood objective does not explicitly encourage minimizing importance weight variance. Furthermore, the flows trained by maximum likelihood use a relatively small, biased training set, which therefore limits their quality.
5 Discussion and Related Work
Augmented flows have been used for improving expressiveness of Boltzmann generators [Köhler et al., 2023b.c]; however, these models were not equivariant. Klein et al. [2023a] proposed an augmented normalizing flow architecture to provide conditional proposal distributions for MD simulations and use coupling scheme similar ours. However, this model only achieves translation and permutation equivariance and the authors make their flow approximately rotation invariant through data augmentation. In our experiments, we found data augmentation to perform significantly worse than intrinsic invariance. In principle, Klein et al. [2023a]'s 's model could be made fully invariant by substituting our flow's projection-basse based coupling transform.
An alternative to our equivariant flow and equivariant CNFs are the equivariant residual flows proposed in Bose et al. [2022]. Alas, residual flows require fixed-point iteration for training and inversion. This is expensive and may interact poorly with energy-based training methods such as FAB [Midgley et al., 2023] which require fast exact sampling and densities. Furthermore, Bose et al. [2022] found that the spectral normalization required for residual flows did not interact well with the cquivariant CNN architecture in their experiments.
There has been recent progress in improving the sampling speed of CNFs/diffusion models [Tong et al., 2023, Song et al., 2023] and on using these models for sampling from unnormalized densities [Vargas et al., 2023, Zhang and Chen, 2022, Zhang et al.. 2023]. Thus, in the future CNF/diffusion models trained by energy may prove be competitive with discrete-time flow based methods.
The strategy of projection into a local reference frame to enforce equivariance has been successfully employed in existing literature, specifically for protein backbone generation [Jumper et al., 2021, Yim et al., 2023]. Here we have focused on modelling the full set of Cartesian coordinates of a molecule, but an interesting avenue for future work is to extend our framework to other domains, such as modelling rigid bodies, which has applications to protein backbone generation [Jumper et al., 2021, Yim et al., 2023] and many-body systems [Köhler et al., 2023c].
Limitations Although our flow is significantly faster than alternatives such as CNFs, the expensive EGNN forward pass required each layer of the flow makes it more computationally expensive than flows on internal coordinates. Additionally, we found our flow to be less numerically stable than flows on internal coordinates, which we mitigate via adjustments to the loss, optimizer and neural network (see App. B.3, App. C. App. C.2). Our implementation uses the B(3) equivariant EGNN proposed by Satorras et al. [2021a]. However, recently there have been large efforts towards developing more expressive, efficient and stable EGNNs architectures [Fuchs et al., 2020, Batatia et al., 2022, Musaelian et al., 2023, Liao and Smidt, 2023]. Incorporating these into our flow may improve performance, efficiency and stability. This would be especially useful for energy- based training, where the efficiency of the flow is a critical factor.
6 Conclusion
We have proposed an SE(3) equivariant augmented coupling flow that achieves similar performance to CNFs when trained by maximum likelihood, while allowing for faster sampling and density evaluation by more than an order of magnitude. Furthermore, we showed that our flow can be trained as a Boltzmann generator using only the target's unnormalized density, on problems where internal coordinates are inadequate due to permutation symmetries, and doing so with a CNF is computationally intractable. It is possible to extend our model to learn the Boltzmann distribution of diverse molecules, by conditioning their molecular graph, which we hope to explore in the future.
10

---

Acknowledgments and Disclosure of Funding
We thank Gábor Csányi and his group for the helpful discussions. Laurence Midgley acknowledges support from Google's s TPU Research Cloud (TRC) program and from the EPSRC through the Syntech PhD program. Vincent Stimper acknowledges the Max Planck Computing and Data Facilities for providing extensive computing resources and support. Javier Antorán acknowledges support from Microsoft Research, through its PhD Scholarship Programme, and from the EPSRC. José Miguel Hernández-Lobate acknowledges support from a Turing AI Fellowship under grant EP/V023756/1 José Miguel Hrrmánddz-Loobrtte and Emile Mathicu are supported by an EPSRC Prosperity Partnership EP/T005386/1 between Microsoft Research and the University of Cambridge. This work has been performed using resources provided by the Cambridge Tier-2 system operated by the University of Cambridge Research Computing Service (http://www.hps www.hhc.ca.a..... capital grant. was also supported by the German Federal Ministry of Education and Research (BMBF): Tübingen AI Center, FKZ: 01IS18039B; and by the Machine Learning Cluster of Excellence, EXC number 2064/1 Project number 390727645.
funded by EPSRC Tier-2
References
llyes Batatia, David Peter Kovacs, Gregor N. C. Simm, Christoph Ortner, and Gabor Csanyi. MACE: Higher order equivariant message passing neural networks for fast and accurate force fields. In Alice H. Oh, Alekh Agarwal, Danielle Belgrave, and Kyunghyun Cho. editors, Advances in Neural Information Processing Systems, 2022. URL YPpSngE-24 ZU
Simon Batzner, Albert Musaelian, Lixin Sun, Mario Geiger, Jonathan P Mailoa, Mordechai Korbluuhh Nicola Molinari, Tess E Smidt, and Boris Kozinsky. E(3)-equivarian graph neural networks for data-efficient and accurate interatomic potentials. Nature communications, 13(1):2453, 2022.
Avishek Joey Bose, Marcus Brubaker, and Ivan Kobyzev. Equivariant finite normalizing flows, 2022.
Andrew Campbell, Wenlong Chen, Vincent Stimper, Jose Miguel Hemnandez-Lobato, and Yichuan Zhang. A gradient based strategy for Hamiltonian Monte Carlo hyperparameter optimization. In Proceedings of the 38th International Conference on Machine Learning, pages 2021.
PMLR,
Taco S. Cohen and Max Welling. Steerable cnns. International Conference on Learning Representa- fions, 2017.
Laurent Dinh, David Krueger, and Yoshua Bengio. NICE: non-linear independent components estimation. In Yoshua Bengio and Yann LeCun, editors, 3rd International Conference Learning Representations, ICLR 2015, San Diego CA, USA, May 7-9, 2015, Workshop Truck Proceedings, 2015. URL http:://trr::
htggabs//444 8516.
Laurent Dinh, Jascha Sohl-Dickstein, and Samy Bengio. Density estimation using real NVP. In 5th International Conference on Learning Representations, ICLR 2017, Toulon, France, April 24-26, 2017, Conference Track Proceedings. OpenReview.net, 2017. URL https::/ net/
Herum?iddHHppbHl1..
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.
Conor Durkan, Artur Bekasov, lain Murray, and George Papamakarios. Neural spline flows. In Hanna M. Wallach, Hugo Larochelle, Alina Beygelzimer, Florence d' Alché-Bus Buc, Emily B. Fox, and Roman Garnett, editors, Advances in Neural Information Processing Systems 32: Annual Conference on Neural Information Processing Systems 2019, NeurIPS 2019, December 8-14, 2019, Vancouver, BC, Canada, pages 7509-75222 7520, 2019. URL https: //tpr//eeeeeeeeeeeeeeeeeeee wtw.oeeediinnnnnniiii..... neurips cc/ paper
htt.://72224aaac22--bbrtere...e Abst .ttll.
Fabian B. Fuchs, Daniel E. Worrall, Volker Fischer, and Max Welling. transfrranssormmrs:: 3d roto- translation equivariant attention networks. Advances in Neural Information Processing Systems 33, 2020.
11

---

Mario Geiger and Tess E. Smidt. e3nn: Euclidean neural networks. CoRR. doi: 10.48850 .9453. URL https: htLtttt/:///// / 488.0..//////.i 2207 09453.
2022.
Will Grathwohl, Ricky T. Q. Chen, Jesse Bettencourt, Ilya Sutskever, and David Duvenaud. FFJORD: free-form continuous dynamics for scalable reversible generative models. CoRR, abs// 810.01367, 2018. URL / orgg www.asg///aa.oeeeo.e abs//810..11....... 1810. 01367.
Jonathan Ho, Ajay Jain, and Picter Abbeel. Denoising diffusion probabilistic models, In Advances in Neural Information Processing Systems, 2020.
Emiel Hoogeboom, Victor Garcia Satorras, Clément Vignac, and Max Welling. Equivariant dif- fusion for molecule generation in 3d. In Kamalika Chaudhuri, Stefanie Jegelka, Le Song, Csaba Szepesvári, Gang Niu, and Sivan Sabato, editors, International Conference on Ma- chine Learning, ICML 2022, 17-23 July 2022, Baltimore, Maryland, USA. volume 162 of URL https:
Proceedings of Machine Learning Research, pages 8867-8887, PMLR, 2022. //proceedings.
mlr .ml.r.g.mmm......mmm.mm.....mmmm
htm... html.
Chin-Wei Huang, Laurent Dinh, and Aaron Courville. Augmented normalizing flows: Bridging the gap between generative flows and latent variable models. arXiv preprint arXiv:2002.07101,, 2020.
Bowen Jing, Gabriele Corso, Jeffrey Chang, Regina Barzilay, and Tommi S. Jaakkola. Torsional diffusion for molecular conformer generation. CoRR, abs/2206.01 729, 2022. doi: 10.48550/arXiv. 2206.01729. URL htw.s::tt//eetc /toi. org/10 4005485500/ 48550/arXivv arriv. 2206. 017299.
John M. Jumper, Richard Evans, Alexander Pritzel, Tim Green, Michael Figurnov, Olaf Ron- neberger, Kathryn uunyasuvunakool, Russ Bates, Augustin Zidek, Anna Potapenko, Alex Bridg- land, Clemens Meyer, Simon A A Kohl, Andy Ballard, Andrew Cowie, Bernardino Romera- Paredes, Stanislav Nikolov, Rishub Jain, Jonas Adler, Trevor Back, Stig Petersen, David A. Reiman, Ellen Clancy, Michal Zielinski, Martin Steinegger, Michalina Pacholska, Tamas Bergham- mer, Sebastian Bodenstein, David Silver, Oriol Vinyals, Andrew W. Senior, Koray Kavukcuoglu, Pushmeet Kohli, and Demis Hassabis. Highly accurate protein structure prediction with AlphaFold. Nature, 596(7873):583 589, 2021.
Diederik Kingma, Tim Salimans, Ben Poole, and Jonathan Ho. Variational diffusion models. In Advances in Neural Information Processing Systems, volume 34, 2021.
Diederik P Kingma and Jimmy Ba. Adam: method for stochastic optimization. arXiv preprint DrXiv:1412.69800 2014.
Leon Klein, Andrew YK Foong, Tor Erlend Fjelde, Bruno Mlodozeniec, Mare Brockschmidt, Scbastian Nowozin, Frank Noć, and Ryota Tomioka. Timewarp: Transferable acceleration of molecular dynamics by learning time-coarsened dynamics. arXiv preprint arXiv:2302.011 01000111700 170, 2023a.
Leon Klein, Andreas Krämer, and Frank Noé. Equivariant flow matching. arXiv preprint arXiv:2306.15000, 2023b.
Jonas Köhler, Leon Klein, and Frank Noé. Equivariant flows: Exact likelihood generative learning for symmetric densities. In Proceedings of the 37th International Conference on Machine Learning, ICML 2020, 13-18 July 2020, Virtual Event, volume 119 of Proceedings of Machine Learning Research, pages 5361 PMLR, 2020. URL http://proceediie mlr. press/v119/ kohler20a.htee html.
Jonas Köhler, Yaoyi Chen, Andreas Krämer, Cecilia Clementi, and Frank Noé. Flow-matching: Efficient coarse-graining molecular dynamics without forces. Journal of Chemical Theory and Computation, 19(3):942 952, 2023a. doi: 10. 1o.102/acs.cttc.c0000
Jonas Köhler, Yaoyi Chen, Andreas Krämer, Cecilia Clementi, and Frank Noé. Flow-matching: Efficient coarse-graining molecular dynamics without forces. Journal of Chemical Theory and Computation, 1933)942-9522 2023b. doi: 10. 20.1021/0
URL https://doi. org/
10. 1al10022sss.... http://csscccc....oeeeeeeoo 3c00016. PMID: 36668906.
Jonas Köhler, Michele Invernizzi, Pim de Haan, and Frank Noé. Rigid body flows for sampling molecular crystal structures. In International Conference Machine Learning, 2023c.
12

---

Yi-Lun Liao and Tess Smidt. Equiformer: Equivariant graph attention transformer for 3d atomistic graphs. In The Eleventh International Conference on Learning Representations, 2023. URL
https://openrevieet http:::///oemeeeeeeeeeeettrrrrt Lerfforrmmid--KmPPPP ARg.TTDD.
Yaron Lipman, Ricky T. Q. Chen, Heli Ben-Hamu, Maximilian Nickel, and Matthew Le. Flow matching for generative modeling. The Eleventh International Conference on Learning Repre- sentations, 2023. URL https: http//pppereviieee net/ fost/forum?i www.mnnddpu/uuDiaa.. PqVMRDDJTT..
Yi Liu, Limei Wang, Meng Liu, Yuchao Lin, Xuan Zhang, Bora Oztekin, and Shuiwang Ji. Spherical message passing for 3d molecular graphs. In The Tenth International Conference Learning Representations, ICLR 2022, Virtual Event, April 25-29, 2022. OpenReview.net, 2022. URL https://opennreview
Laurence Illing Midgley, Vincent Stimper, Gregor N. C. Simm, Bernhard Schölkopf, and José Miguel Blemández-Lobald Flow annealed importance sampling bootstrap. International Conference on Learning Representations, 2023.
Tom Minka. Divergence measures and message passing. Technical report, Microsoft Research, 2005.
Yoshiharu Mori and Yuko Okamoto. Replica-exchange molecular dynamics simulations for various constant temperature algorithms. Journal of the Physical Society of Japan, 79(7):074001, 2010.
Albert Musaelian, Simon Batzner, Anders Johansson, Lixin Sun, Cameron Owen, Mordechai Korbluth, and Boris Kozinsky. Learning local equivariant representations for large-scale atomistic dynamics. Nature Communications, 14(1) ):579, 2023.
Frank Noć, Simon Olsson, Jonas Köhler, and Hao Wu. Boltzmann generators: Sampling equilibrium states many-body systems with deep learning. Science, 365(6455) 2019. doi: 1147. URL https://vvin http::::/wwwssiiin....... science. www.or/g/////////////ereeeeeae// htt//ddi//bbba/ att////ss/10000/ //cinnc.
aaw1147.
Art B. Owen. Monte Carlo theory, methods and examples. 2013.
George Papamakarios, Eric Nalisnick, Danilo Jimenez Rezende, Shakir Mohamed, and Balaji Lakshminarayanan. Normalizing flows for probabilistic modeling and inference. Journal of Machine Learning Research, 22(57): 2021.
Peter Petersen. Linear Algebra. Springer, 2012.
David Pollard. A User's Guide to Measure Theoretic Probability. Cambridge University Press, 2002.
Raghunathan Ramakrishnan, Pavlo O Dral, Matthias Rupp, and O Anatole Von Lilienfeld. Quantum chemistry structures and properties of 134 kilo molecules. Scientific data, 1(1):-7,, 2014.
Danilo Jimenez Rezende, George Papamakarios, Sébastien Racanière, Michael S. Albergo, Gurtej Kanwar, Phiala E. Shanahan, and Kyle Cranmer. Normalizing flows on tori and spheres. In Proceedings of the 37th International Conference Machine Learning, volume 119, pages 8083-8092. PMLR, 2020.
Victor Garcia Satorras, Emiel Hoogeboom, Fabian Fuchs, Ingmar Posner, and Max Welling. E(n) equivariant normalizing flows. In Marc' Aurelio Ranzato, Alina Beygelzimer, Yann N. Dauphin, Percy Liang, and Jennifer Wortman Vaughan, editors, Advances in Neural Information Processing Systems 34: Annual Conference Neural Information Processing Systems 2021, NeurIPS 2021, neurips..... December 14, 2021, virtual, pages 4181 4192, 2021a. URL www.aaafaaacatttattttaaaaaaaaaa. html. hc//aaper2021/// http:/aaaa........0.0.0000...10.
Victor Garcia Satorras, Emiel Hoogeboom, and Max Welling. E(n) equivariant graph neural networks. In Marina Meila and Tong Zhang, editors, Proceedings of the 38th International Conference on Machine Learning, ICML 2021, 18-24 July 2021, Virtual Event, volume 139 of Proceedings of Machine Learning Research, pages 9323 9332. PMLR, 2021b. URL http:://t// httt:://pproceedngg. mlr. wwwpp//sssveee
html.
Yang Song, Prafulla Dhariwal, Mark Chen, and Ilya Sutskever. Consistency models. International Conference on Machine Learning, ICML. 2023, 23-29 July 2023, Honolulu, Hawaii, USA, 202: 32211 12211-32252, 2023.
13

---

Vincent Stimper, Bernhard Schölkopf, and José Miguel Herándezz-Lobato Resampling Base Distributions of Normalizing Flows. In Proceedings of The 25th International Conference on Artificial Intelligence and Statistics, volume 151, pages 4936, 2022.
Nathaniel Thomas, Tess Smidt, Steven Kearnes, Lusann Yang. Li Li, Kai Kohlhoff, and Patrick Riley, Tensor field networks: Rotation- and translation- eqquiiirinn neural networks for 3d point clouds. arXiv preprint arXiv: arXhv:1802.080 0821 2018.
Alexander Tong, Nikolay Malkin, Guillaume Huguet, Yanlei Zhang, Jarrid Rector- Brotor, Kilian FATRAS, Guy Wolf, and Yoshua Bengio. Improving and generalizing flow-based generative models with minibatch optimal transport. In ICML Workshop on New Frontiers in Learning, Control, and Dynamical Systems, 2023. URL
nee.i/eweeeee/eeeeeeme.eeeeeeeee net. net/forum? http//ormmmniii/iiiiiiteeeee id=HgDwiZrpVq
Francisco Vargas, Will Sussman Grathwohl, and Arnaud Doucet. Denoising diffusion samplers. In The Eleventh International Conference on Learning Representations, 2023. URL https: // //openreview
Hao Wu, Jonas Köhler, and Frank Noć. Stochastic normalizing flows. In Hugo Larochelle, Marc' Aurelio Ranzato, Raia Hadsell, Maria-Florina Balcan, and Hsuan-Tien Lin, editors, Advances in Neural Information Processing Systems 33: Annual Con- ference on Neural Information Processing Systems 2020, NeurIPS 2020, December 6- 12. 2020. virtual, 2020. URL https://proceeding neurips. pttppappp///2222// 116808813777119808822228888111.1
A-stracttt.....a html.
Zhisheng Xiao, Karsten Kreis, and Arash Vahdat. Tackling the generative learning trilemma with denoising diffusion gans. In The Tenth International Conference Learning Representations, ICLR 2022, Virtual Event, April 25-29, 2022. OpenReview.net, 2022. URL www.htttp//// httbt::://ppeerevieee http:/oornnid-jjoopp--00000
Minkai Xu, Lantao Yu, Yang Song, Chence Shi, Stefano Ermon, and Jian Tang. GeoDiff: Diffusion Model for Molecular Conformation Generation. In International Conference on Learning Representations, 2022.
Geometric
Jason Yim, Brian L. Trippe. Valentin De Bortoli, Emile Mathicu, Arnaud Doucet, Regina Barzilay, and Tommi S. Jaakkola. SE(3) diffusion model with application to protein backbone generation. 202:40001 -4000990144009, 2023.
Dinghuai Zhang, Ricky Tian Qi Chen, Cheng-Hao Liu, Aaron Courville, and Yoshua Bengio. Diffu- sion generative flow samplers: Improving learning signals through partial trajectory optimization. arXiv preprint arXiv:2310.02679 2023.
Qinsheng Zhang and Yongxin Chen. Path integral sampler: A stochastic control approach for sampling. In International Conference on Learning Representations, 2022. URL https:// openreview. net/ met//orum? or _uCb2ynRu7Y.
14

---

A Background
A.1 Tensors and representation
We assume group (G, algebraic structure that consists of a set G with group operator that follows the closure, inverse and associativity axioms.
A representation is an invertible linear transformation p(g) V + V, indexed by group elements g E G. and defines how group elements g acts on elements of the vector space V. It satisfies the group structure p(gh) = p())
A tensor T is geometrical object that is acted on by group elements 9 G in a particular way, characterized by its representation ft g = p(g)T. Typical examples include scalar with Periv and vectors living in V with pla(g) = g. Higher order tensors which live in V x can be constructed as 'T' = (g) @ ... @ Pld (g)) vec(T". A collection of tensors ; T.) with representations (P1...., Pn) can be stacked with fiber representation the direct sum p
V
P1 O ... + Pr.
A.2 Translational invariant densities
Proposition A.1 (Disintegration of measures for translation invariance [Pollard. 2002, Yim et al., 20231). Under mild assumptions, for every y an S0(3)-invariant probability measure on [R E I; = 0} and p proportional to the Lebesgue measure on R& s.t. dp([I1,..., = ( 1/4 Ei= x;) x cn1( DE Is...., ,
measure μ R33*" there exists
1 > z:)).
The converse also true, which motivates constructing an (0(3)-invariant probability measure on the subspace of R3xn with zero center of mass.
B
Method
B.1
Transforming variables on subspace
This section shows that a flow transform F X X applied to observations embedded in the zero-CoM subspace X does not have well-defined Jacobian. We will show this by constructing such transform, which leaves zero-CoM input unchanged, but for which the log-determinant is arbitrary.
Without loss of generality, augmented variables are one dimensional, i.e., X' = A = R Consider the orthonormal basis V E Raxx constructed by concatenating orthonormal basis vectors v; E R² as [Vi,.... 7T We choose U1 to be equal to 1,.......]" //r/ such that inner products with i2 retrieve the CoM of a system particles up to multiplication by a scalar: of x /næ and of a 0. Let F(x)- = s
will assume in this subsection and the next that our observed and
V SV x be mapping which rotates onto the basis defined by V. multiplies by the diagonal matrix = diag([s, I,..., 1)) and then rotates back to original space with VT, This transformation only acts on the CoM. Applying this transformation to zero-CoM variable à we have
F()) V SVÆ [= ā, 2y ã ..... - V uf ..... e + u
leaving the zero CoM input unchanged. However, the log deterremminn
0F(x)
log
log VT SV = log | + log |S] + log [V| = log s,
dæ
since for the unitary matrices | |
1. Because is arbitrary, is the density returned by any flow applying such a layer to zero-CoM variables. In the maximum likelihood setting, the aforementioned flow would lead degenerate solution of 0 such that the log determinant would approach negative infinity and the log likelihood would approach infinity.
While our flow does use transforms which operate on non-CoM variables, we lift our CoM ob- servations to the non-zero- space before applying said transformations. We do this using the ShiftCoM operation, described in Sec. 3 and below.
15

---

B.2 Jacobian of ShiftCoM
We now show our ShiftCoM: X' x'xAA> X Á transformation have unit Jacobian log determinant. We first, re-state the function definition
ShiftCoM(E, a) _ (0 a. a a) with a 2
71 Σ aa
We now re-write it using the orthonormal basis V E _ defined in the previous subsection. For this, we use I2 @ V E 222nx2n to refer to a 2 x block diagonal matrix with V both blocks, and we also stack the inputs and outputs of Shift CoM, yielding ShiftCoM([ al) (I2 @ V) P(I2 @ [) a), where a permutation matrix that exchanges the first and n + 1th elements of the vector acts upon. It also flips this element's sign. To see this note
ShiftCoM(( al) = (12 @ P(2 @ V) ã a
(42 @ v) u ã a.... u k. u a a...... o
(12 @ v)"l-v a, 2 x ................................ of 0, 02
0 aaa
= [ - a, a - a].
For the determinant we use that [/] = 1 and [V = 1, and thus the determinant of their Kronecker
product will be equal one (12 @ v)1 1. Combining this with the fact that permutation matrices and matrices that flip signs have unit determinant, we arrive at
[ShiftCoM((2 al) = log
log
[I2 &
a|ž,
log (I2 & V)² log Pl + log 122 & =1
=1.
It is worth noting that ShiftCoMM preserves the inherent dimensionality of its input; both its inputs and outputs consist of one zero-CoM vector and one non-zero-CoM vector. Because ShiftCoM does apply any transformations to the CoM, it does not suffer from the issue of ill-definednes discussed in the previous subsection.
B.3 Projection variants and their respective transformations
In this section we provide detailed description of the core transformation of our equivariant coupling flow layer, see Fig. 1. It consists of a projection into an invariant space via an equivariant reference, followed by flow transformation in this space and a back-projection.
We denote the inputs with superscript and use l + 1 for outputs to define our transformation block F: (x"; at) + (x++1, 11) Without loss of generality, we assume that we are transforming xf. Then, JF takes the following form.
x f+1
To(Tr x²), with (r. = h(a').
"rr
(10)
E+1 = a.
a
Here, h is (graph) neural network that returns set of equivariant reference vectors [ri,..., r"]. and invariant parameters 0. The equivariant vectors r parametrize ?r. a projection
operator onto rotation invariant feature space. This means that any rotations applied the inputs
will be cancelled by Tr. i.c. Ygr 7/- g or equivalently (7/-r)=
7r for all g E SO(3).
In the next paragraphs, we discuss three projection operations and their respective flow transforma- tions.
Cartesian projection Here, we predict a 3D Cartesian coordinate system as illustrated in Fig. 4a. The coordinate system has an origin و and an orthonormal basis Q, both being equivariant. Q is constructed from two equivariant vectors v1 = , - y and 12 p3 1d in the following way. To obtain the first basis vector, we simply normalize u1, i.e. b, = | We get the second basis vector via a
mi
[Petersen, 2012] iteration, i.e. we compute the component of 02 orthogonal to v1, which is U2 I 1/2 bJ V2, and then normalize it to obtain b2 I Note that this is ill-posed if W1 and Uy are collinear. To avoid this case, we introduce an auxiliary loss
aux-loss = log (c (abcss o2/la///////la/////l/ll/l
(11)
16

---

by
b1
(xc)
g
1
V2,1
b2
o
b3
02
80
22
3
(x*
01
b2
(a) Cartesian projection.
(b) Spherical projection.
Figure 4: Illustration of the Cartesian and spherical projection.
which uses logarithmic barrier function the angle between the vectors prevent them from being collinear. e a small constant that prevents the aux-loss from being equal to infinity, Finally, the third basis vector is by = 62 x 62, and we can assemble
[b1, b2, b3].
The projection into this coordinate system and the respective inverse operation are given by 7/ (a) = Q (x o)),
(12)
(ap) = Qm, + o,
(13)
where xp is the transformed projected variable.
In the projected space, we can apply any standard flow transform, as long as acts on each particle individually. We choose To affine transformation in the Real NVP architecture [Dinh et al., 2017].
Note that since by is constructed via cross product, it is not parity equivariant, which is why the entire transformation block not parity equivariant either. However, since sign(b3 r2) is a pseudo-scalar, we can replace by with b/g sign(b, (63 y²) 62,,, which is vector, to retain parity equivariance.
The determinant of the Jacobian of F when using the Cartesian projection is just the determinant of the Jacobian of 7o, which show below.
(F(#;a)
0F(x; aa
OF(æ; a)la da
a)
Zx
105 (z)
ÔTa (z)
@x dx
1
dz
Oz
T= (m))
(x=)
[TT(2)
Q
Q-1
dz
x=))(()
iro (TT(2))
Q
1Q-1
Oz
(x++(())
Oz
(.))
Spherical projection For the spherical projection, which illustrated Fig. 4b, we compute an origin and three Cartesian basis vectors b,. by, and by as in the Cartesian projection. Instead of
17

---

computing the Cartesian coordinates of each x', we compute the spherical coordinates (p. B, Hence, for the i-th atom the projection operation and its inverse is given by
(*x(2)))* || arccos
atan2 (bg x2 bJ x4))
(14)
( (cos( yb + sin( (cos(sp") yb2 + sin sin())))b)))
(15)
Here, we used "f1 = b1 x², sr x² rfj- and x/ = (p', 3r, sot) for brevity.
The spherical coordinates are transformed with monotonic rational -uaaoraai splines [Durkan et al., 2019]. The radial coordinate p is transformed through a spline with the interval starting at 0 to ensure that the transformed coordinate is also non-negative to ensure invertibility. The two angles are transformed via circular splines [Rezende et al., 2020]. Similarly to the Cartesian projection, we use auxiliary loss (Eq. (11)) to prevent the vectors output from the GNN from being co-linear.
When using the spherical projection, we compute the log-determinant of the Jacobian of F by using the chain rule in the same way as we do for standard normalizing flows. The Jacobian determinant of the projection operation is given by the well-known term
a (x)()))) = (ot)² sin 3".
(16)
Ox:
For the inverse, we can just raise this term to the power of 1.
Vector projection This is special case of the spherical projection, where we only transform the radial coordinate but leave the angles the same. Since practice we only need to compute p. we call this case vector projection. Note that just predicting the origin neural network is sufficient here. Unlike the Cartesian and spherical projections, the vector projection does not rely the cross product, and hence is parity equivariant by default.
reference with the equivariant
B.4 Multiple augmented variables
Consider augmented variables consisting of multiple sets of observation-sizee arrays, i.c. A x with k > 1. For the base distribution, and also the target, we deal with multiple augmented variables by letting them follow conditional normal distributions centered at æ. We thus have (aax) = I N(a; x, (2) and (z, a) Naw; 0, IT
x, o²1).
For flow block, (0(x-1)/2+1)... az), and update one set conditional on the other, and vice-versa, while always ensuring the arrays being transformed have non-zero CoM.
group our variables into two sets (ã, ai..... a(k- and
Initially, a is non-zero-CoM with each augmented variable a, having CoM located at a different point in space. Before the first coupling transform F(I). the ShiftCoM function applies the mapping
(17) taking (kk-1)/2+1 to the zero-COM hyper-plane, while a and every other augmented variable remain non-zero-CoM. Before the second coupling transform F(2). we apply the ShiftCoM maps in reverse, taking æ back to the zero-COM hyper-plane.
ShiftCoMMi, a) (Ø ā d a x+7+ with دي4باة 1 1 |
The coupling transforms (()) M and F work as previously, except now they transform (k + 1)/2 sets of variables (holding the other set constant), where for each of the variables, each of the nodes get their own local reference frame and inner flow transform T0-
B.5 Invariant base and target distributions
We now show the SE(3) x S, invariance of our extended base distribution choice, which trivial extends to our choice of augmented target distribution. Consider the choice mallz) N(a; 7 for 2 scalar hyperparameter. With this choice, the constraint interest satisfied since the Isotropic Gaussian is rationally invariant and centered at To see this, let E SE(3) be decomposed
18

---

into two components = (R, u) where R E SO(3) is a 3 x 3 rotation matrix and u E R represents a translation,
1
-1
M(Ra + Ra u, 1)
oxp
[Ra
Rx u) (n I + Ra + u) (18)
Z
1
-1
(a x) T
TR [m= xx)
exp
Z
2
(a - x) R R(n =1) - x)
exp
Z
2
<aa
11²1)
(19)
An analogous argument holds for permutations, since any Sn, can be represented unitary n x n matrix which is simultaneously applied to a and æ.
B.6 Maximum likelihood objective
We train our models from samples by maximising the joint log-density (g a la)) [log g( a)]. This can be shown be lower bound the expected marginal log-density over observations, which the quantity targeted by models without augmented variables as
((x a) < x:)
Ep(w) [[x)]] = Ep(a) log En(aaa)
(p(e))(aaa) glog a) log ा (aax)) (20)
where we drop the latter term from our training objective since it constant w.r.t. flow parameters. The bound becomes tight when ((a2z) = T (a)
B.7 Reverse KL objective
We now show the reverse KL objective involving the joint distribution of observations and augmented variables presented in the main text upper bounds the reverse KL over observations. Using the chain rule of the KL divergence we have
DKL (q(2, a) (1 DKL (q(z) || p(x)) Bg( (g(a|z) || (ala))
2 DRL (q(x) [[)))
The looseness in the bound is (e)) DKL (4(ala:) [[ which again vanishes when g(a/æ)
T (a(
B.8 Alpha divergence objective
We now show that our joint objective represents analogous upper bound for the divergence when
a = 2, i.e., the objective targeted by FAB [Midgley et al., 2023]. Let U(x) be the unnormalized
energy function of interest evaluated state æ, with p(æ)) exp(-U((n) (x)).. We have
= p(x)²
D2 (((x))[((()))
dx
((x)
/
) π T(alx)2
(aaxx) da dæ
<(x)
(aaz))
da
(())g(a(())
D2 (g(0, a) (<)))((()))
which also becomes tight when q(a(æ)) = क (a|x). The inequality is true because To see this, take any densities π and q, and apply Jensen's 's inequality as
s(al=)) q(n|m) da
T())2 da Enja) ((a)
T(a)
((a)
-1
g(a) -1 Taa) = =/
Ex(a)
> En(a)
g(a)
(aa)
19

---

B.9 Importance weighted estimator
We define the following estimator a(m.a) then q(x, a) = O, we have that this estimator trivially unbiased:
S[a|=| with x(-[x). Assuming that whenever <(aa))
((x,, a)
Biw]
=
π <ala))
π (aal))dp(a
/
g( a)dp(a)
gxx)
and that variance is given by
V[w]
B[(w - x(x))
= E E[w²] - E2 [[2wg(x)] + E[g(z))
/ g(0r, a)²
<ala=)? (a (d)())))(a))))))))))))))))) 2q(x))[(vo] + 2(x))
f ((x))(((a)))
<<aa)) q( a)dx(()))))))))))))))))a)) - q(x))
thus if ((a(=)) x (a)x)) V[w] = 0.
The log marginal likelihood may therefore be estimated with
1 M
q am)
log ((xr) RJ logg M
Σ
and Om ~ x(112)),
(21)
n (a (z))
which becomes accurate M becomes large.
B.10 Effective sample size
For evaluation we use both the reverse and forward effective sample size (ESS), where the former relies on samples from the model, and the latter from samples from the target. For the augmented flow models we report the joint ESS, which lower bounds the marginal ESS, Below, we provide the formulas for the reverse and forward ESS. For simplicity, we define these in terms of the marginal ESS. The joint ESS estimators may be obtained using the same formulas, replacing æ with (x, a).
The reverse effective sample size is given by
(E wxx;)))2 and 2; ~ g(xx)) Σn wxx))²2
Theav
(22)
where w(x;)) p (q(x))((xx and is the unnormalized version of p [Owen, 2013].
The forward ESS may be estimated with
n
Re.fwd -- E wxx))
and ~ p(x,),
(23)
where the normalizing constant 7 p/p may be estimated using
9(x))
z = Ep
(24)
[p(z))
(xx))
and x; ~ p(x)))
(25)
((x,)
B.11 Conversion of internal coordinate flow density density on ZeroCoM Cartesian space with global rotation represented
The density of the flow on internal coordinates is not directly comparable the flows on Cartesian coordinates, the variables live in different space (bond distances, angles and torsions instead
20

---

of Cartesian coordinates). The internal coordinates, and density may be converted into Cartesian coordinates via passing them through the appropriate bijector and computing the log determinant of the bijector. This is commonly done in existing literature as the target energy typically takes Cartesian coordinates as input [Midgley et al., 2023]. However, this representation invariant to global rotations of the molecule, making it lower dimensional than the space with which the flows which introduced in this paper operate. Furthermore, this representation has the first atom centered on the origin, where the flows in this paper instead have zero-CoM. To compare the density of the flow on internal coordinates with the flows that operate directly on Cartesian coordinates we have to account for both of these aforementioned factors. We describe how to do this below.
Conversion to density that includes degrees of freedom representing global rotation The representation of internal coordinates in Cartesian space places the first atom at the origin, the second atom a distance of 61 along the X-axis (i.e. position {61, O, 0}), and the third atom on the XY plane, at a distance of by from the origin at an angle of a2 with respect to the x-axis (i.e. at position {cos(ag) b2, sin(az) b2. 0}). We define this representation of Cartesian coordinates as 2 E Z, where we drop all coordinates of the first atom (atom 0), the y and & coordinates of atom 1, and the z coordinate of atom 2, all of which are zero, so that Z contains as many elements as there are degrees of freedom.
To convert this into space that has degrees of freedom representing global rotations, which we refer to as w E W, we simply sample rotation uniformly over SO3 and apply this to w.
The rotation has three degrees of freedom, which represent follows; First, apply rotation about the X-axis of 0₁ sampled uniformly over [-и, æ/. Secondly, we rotate about the Z-Axis by sampling 02 uniformly between [-1, and rotating by sin (02). Finally, we rotate about the Y-axis by Og sampled uniformly between [-T,, rr. Together this ensures that the resultant orientation of the molecule in Cartesian space has a uniform density over SO3. We denote the concatenation of these random variables that parameterize the rotation
O where O [[- T], [-1, [- I,, x|}.
The density on W then given by
((w) = (())((0) det Ōg(z,
(26)
O(z,0]
where g: x O > W is bijector which applies the rotation parameterized by o to 2.
can be shown that the log jacobian determinant of is 2lo((b1) + log(b>) + log(sin(az)) Further- more, log (()) is constant value of log 8T.
Conversion density on Zero-CoM space After applying the above transform we have w that includes rotations, but is centered on atom O. Thus, the last step needed is to convert into the Zero-CoM space. As previoulsy, to ensure that our representations have the same number of elements as there are degrees of freedom, we drop the first atom's 's Cartesian coordinates, these contain redundant information. We note that for the representation w E W the first atom's coordinates are at the origin, and for the Zero-CoM space æ E X they are at Em= Ky. The bijector h : W + X is given by
h(w) = w
Σ
w
(27)
where N is the number of nodes. The log jacobian determinant of h is simply equal to -3 log n.
Overall conversion
Thus we can convert the density in Z into an equivalent density in X using,
log ((x) = log q(z) + log q(0) + log det 0g(z, o) 0(z,0)
+ log det ah(w)
Sw
(28)
log q(2) - log ST - 2 log(bi) - log(b2) - (og(sin(a2)) + 3 log
C Experiments: Details and Further Results
We provide the details of the hyper-parameters and compute used for each experiment. Addition- ally, we provide the code https: //tp:////ttttt//e htt.g/ib...ooooolaata///////ee/e aerr--aummeneddeee -ugmente--ccuullggg------ couplnggfloow
21

---

for replicating our experiments. It also contains the configuration files with the exact hyperparameter for each model and training run.
C.1 Custom Optimizer
For all experiments we use a custom defined optimizer, which improves training stability by skipping updates with very large gradient norms, and clipping updates with moderately large gradient norms. Our optimizer keeps track of a window the the last 100 gradient norms, and then for a given step: (1) If the gradient norm is more than 20 times greater than the median gradient norm within the window, then skip the gradient step without updating the parameters. (2) Otherwise, if the gradient norm more than five times the median within the window then clip the gradient norm to five times the median. (3) Use the (possibly clipped) gradients to perform a parameter update using the adam optimizer [Kingma and Ba, 2014]. Lastly, we also include an optional setting which allows for masking of gradient contribution of samples within each batch with very large negative log probability during training, which further improves training stability.
C.2 EGNN implementation
We use our own custom implementation of the EGNN proposed by Satorras et al. [2021b] for our experiments. We found that applying softmax the invariant features output from the EGNN helped improve stability. Additionally, we extend Satorras et al. [2021b]'s EGNN to allow for multiple channels of positions/vectors each node, with interaction across channels. This is needed as the local reference frame typically is parameterized by multiple equivariant vectors, and for the case where the multiple sets of augmented variables are used such that there are multiple sets of positions input to the EGNN. We refer to the code for the details of our EGNN implementation.
C.3 Training with samples
C.3.1 DW4, LJ13 and QM9 Positional
Experimental details Below provide description of the syper-parameteer used within each
experiment. Additionally Tab. 4 provides the run-times for each experiment.
Datasets: For DW4 and IJ13 we use training and test set of 1000 samples following Satorras et al. [2021a]. We note that this data is biased with a clear difference between the distribution of train and test data (see Fig. 5. For
use training, validation and test set of 13831, 2501 and 1,813 samples respectively also following Satorras et al. [2021a].
DW4
Bomaalized count
train
normalized count
0.06
0.04
0.03
0.01
0.02
0.01
0.00
0.00
1.0
1.5
2.0
2..6
3.0
3.5
4.0
inseratomic distance
inseratomic distance
Figure 5: Inter-atomic distances for samples from the train and test data DW4 and LJ13. Both datasets are biased with visible differences between the train and test sets.
Flow experiment details All flow models using 12 blocks (see Alg. 1 for the definition of block). For the noise scale of the augmented variables sampled from N(a; I),,, we use = 0.1 for E-ACF and
both the base of the flow, and augmented target distribution. For the CARTESIAN-PRE NON-E-ACF we use RealNVP [Dinh et al., 2017] for the flow transform, while for CARTESIAN-PROI and VECTOR-PROI we use a rational quadratic spline transform [Durkan et al., 2019]. For the spline use 8 bins for all flows. We use set of augmented variables for DW4 and LJ13, and 3 sets
22

---

of augmented variables for DM9-positional. For all equivariant GNN's we implemented the GNN proposed by Satorras et al. [2021b] with 3 blocks. All MLP's have 2 layers, with a width of 64 units. For the NON-E-ACE model we use a transformer with 3 layers of self-attention with 8 heads. We use a cosine learning rate schedule, that is initialised to value of 0.00002, peaks at 0.0002 after 30 epochs, and then decays back to a final learning rate of 0.00002 for DW4 and LJ13. For QM9 the learning rate is initialised to a value of 0.00006, peaks at 0.0002 after 30 epochs, and then decays back to final learning rate of 0.00006. With the above learning rate schedule for QM9 the CARTESIAN-PRE E-ACF was unstable during training which resulted in crashed runs (NaN loss). Hence, it was instead trained with constant learning rate of 0,00006 at which training was stable. We use a batch size of 32 for all experiments, and 100, 400 and 800 epochs of the datasets for DW4, LJ13 and QM9 posiiionli respectively. We train the all problems, which results in roughly the same run time as the SPHERICAL- PROJ E- A-F. For the CARTESIAN-PROJ and SPHERICAL-PRE B-ACF, we use coefficient of 10. for the auxiliary loss (Eq. (11)), for weighting it relative to the maximum likelihood loss. Marginal log likelihood in Tab. 1 is calculated using Eq. (21) with 20 samples from the augmented variables per test-set sample of
for 4 times many epochs for
Continuous normalizing flow difftusion experiment details For the E-CNF ML we report the results from Satorras et al. [2021a]. For the E-CNE FM we implement flow matching, with the model's vector field parameterized by the E(n) GNN of Satorras et al. [2021b]. For DW4 and LJ13 an EGNN with 3 message passing layers, with 3-hidden layer MLP with 128 units per layer. For QM0--MM-P--- AM9-positional we use an EGNN with 5 message passing layers, with an 4-hidden layer MLP with 256 units per layer. For all problems decay from le-4 to O. Additionally, for QM9, we use EMA with , = 0.999 to compute the final training weights which the model is evaluated with. We train for 200, 400, and 1600 epochs on DW4, LJ13 and QM9-positional respectively. For sampling and evaluating log-densities we use the dopri5 ODE solver [Dormand and Prince, 1980] with an rtol and atol of le-5. For DW4 and LJ13 we evaluate densities using the exact divergence of the vector field, while for QM9 we approximate the divergence using the Hutchinson trace estimator [Grathwohl al., 2018] as the exact divergence is computationally expensive.
use Adam optimizer with cosine learning rate
Our E-CNF-DIFF implementation roughly follows Hoogeboom et al. [2022]. but we use the noise schedule as Ho et al. [2020]. We use the same GNN architecture as for our flows, that is the E(n) GNN of Satorras et al. [2021b], to predict the diffusion noise. The network is composed of message passing layers, each of which uses a 3 hidden layer MLP with 256 units per layer. GNN expands its inputs to 256 scalars and 196 vectors per graph node and returns its predictions as linear combination of the latter 196 vectors. This model has around 90M learnable parameters. We use the AdamW optimiser with weight decay of 10-4 and run it for 100, 200, and 5000 epochs for the DW4, LJ13 and QM9 dataset, respectively. The initial learning rate is set to 10-5, It is linearly increased to 5 1 for the first 15% of training before being decreased schedule. Following, Kingma et al. [2021]. we sample noise timesteps during such that these are uniformly spaced within each batch, reducing gradient variance. We set the smallest accessible time value to 10 For sampling and evaluating log-densities we use the dopri5 ODE solver with a maximum number of evaluations of 1000.
10-6 with a cosine
Table 4: Run time for training by maximum likelihood on the DW4, LJ13 and QM9 1M9--ositional datasets, in hours. All runs used GeForee RTX 2080Ti GPU Run times include evaluation that was performed intermittently throughout training. The results are averaged over 3 seeded runs, with the standard error reported as uncertainty.
DW4
LJ13
QM9 positional
E-CNF ML E-CNF-DIE NON-E-ACE VECTOR-PROI E-ACF
N/A N/A N/A
N/A
336 N/A
N/ALN/A
3.80
0.10
0.000
20.68
0.16
0.61
19.0
CARTESIAN-PRO E-ACF
0.02
0.00
17.13+0.33
SPHERICAL- E-ACF
0.05
0.00
0.117
Further results Tab. 5 provides the effective sample size of the flows trained by maximum likelihood, with respect to the joint target distribution p(x) (a)a))
23

---

Table 5: Joint effective sample size (%) evaluated using 10000 samples. The results averaged over seeded runs, with the standard error reported as uncertainty.
DW4
LJ13
Rev ESS (%)
Fwd ESS (%)
Rev ESS (%)
Fwd ESS (%)
E-CNF FM NON-E-ACE VECTOR- E-ACF
5.02
0.01 + 0.00
5.37
0.10
0.68 +
0.06 0.02
0.00
0.00
1.87 0.68
0.16
1.10
0.06
CARTESIAN-PROI E-ACF
3.08 0.05
0.09
0.37
0.00
SPHERICAL-PROI E-ACF
3.06
0.17 0.05
1.30
0.00
Varying the number training samples for DW4 and IJJ13 Tab. 6 shows the effect of varying the training dataset size on the performance of the flows. We tune the number of epochs via the validation set; we use 1000, 100, 50 and 35 epochs for DW4 and 1500, 1500, 400 and 100 LJ13 for each dataset size respectively. Additionally we train NON-E-ACF (with data augmentation) for 4 times as many epochs for each experiment, as it is computationally cheaper per epoch and takes larger number of epochs before it overfits. We that the performance gap between the non-equivariant and equivariant flow models larger in the lower data regime.
Table 6: Negative log likelihood results test set for flows trained by maximum likelihood for the DW4 and LJ13 datasets varying the number of training samples. E-CNF results are from Satorras al. [2021a]. Best results are emphasized in bold.
DW4
LJ13
samples
10a
10a
10a
10a
10
10a
10a
10a
E-CNF MON-E-ACF VECTOR- E-ACF
8.31
8.15
7.69
7.48
33.12
30.99
30.56
30.41
10.24
10.07
8.23
7.39
40.21
34.83
33.32
31.94
8.65
8.69
7.95
7.47
33.64
30,61
30.19
30,09
CARTESIAN-PROJ E- A-ACF
8.79
8.82
7.93
7.52
36.66
31.72
30.89
30.48
SPHERICAL-PROI B-ACE
8.66
8.61
7.97
7.38
34.79
31.26
30.33
30.16
C.3.2 Alanine dipeptide
4
Figure 6: Visualization of the 22-atom molecule alanine dipeptide with its two dihedral angles & and wb.
The molecule alanine dipeptide and its dihedral angles ó and th which are used in the Ramachandran plots in Fig. 3, are depicted in Fig. 6. We model it in an implicit solvent at a temperature of 800 K. The data used for maximum likelihood training and testing was generated with replica exchange MD simulation [Mori and Okamoto, 2010] with systems having temperatures from 300 until 1300 K with 50 K temperature difference between each of them.
Each ACF has layers with core transformations per layer for the VECTOR-PROS flow and for the CARTESIAN- PROJ and the SPHERICAL- the same architecture as the models in Midgley et [2023], i.e. it is coupling neural spline flow with 12 layers. The models were trained for 50 epochs.
flow. The flow trained on internal coordinates has

---

10
10
10
10
(a) Ground truth
(b) Raw model samples
(c) Reweighted model samples
Figure 7: Ramachandran plot of SPHERICAI-PPOI B-ACF model trained alanine dipeptide
800 K. The samples in Fig. 7c are reweighted with the importance weights. Thereby, the KLD drops from 1.88 10-3 to 1.39 10-3
The E-CNF and the E-CNF-DIFF model have E(n) GNN with 4 message passing layers, each using a 3-layered MI. with 128 hidden units each, as a vector field and score function, respectively. Both trained for 20 epochs.
Since the version of the VECTOR PROJ flow as well as the E-CNF and E-CNF-DIFF models are parity equivariant, the model learns to generate both chiral forms of alanine dipeptide [Midgley et al., 2023]. For evaluating the model, we filtered the samples from the model to only include those that correspond to the L-form, which is the almost exclusively found in nature.
To plot and compute the KLD of the Ramachandran plots, we drew 10" samples from each model. Since computing the likelihood of the E-CNF and the E-CNF-DIFF model is so expensive, we evaluated these models only on 10% of the test set, i.e. 10 datapoints.
Limiting the number of samples To test how well E-ACFs generalizes thanks to their equivariance, we trained our three variants of E-ACFs and an NON- NoN-E-ACFor subsets the training dataset with various sizes. The results are given in Table 7. For all training dataset sizes the equivariant architectures clearly outperform the non-equivariant counterpart. Notably, the NON-E-AS IEACF trained on the full dataset performs worse than the VECTOR-PROT R-ACH trained on 10% of the data and the CARTESIAN-PROI and SPHERICAL-PRUS E-ACF trained on just 1% of the data.
Table 7: NLL of the models trained on subsets of the training dataset with different sizes. The best performing model, i.e. the one with the lowerst NLL, marked in bold for each training dataset size.
Training dataset size
10a
10a
10a
10a
10
NON-E-ACE VECTOR-PROI E-ACF
-28.8
-76.8
120.4
-164.3
-177.8
-31.0
-111.4
-169.5
-183.6
-18855
CARTESIAN-PROI E-ACF
-65.1
-155.7
-185.9
-188.1
-188.6
SPHERICAL-PROI B-ACF
-53.1
-119.3
-183.6
-188.3
-188.6
Reweighting Given the high effective sample sizes we obtained for our models (see Tab. 2), we also attempted to reweight the samples. As already observed by Midgley et al. [2023], for a large number of samples, i.e. 10", there were few high outlier importance weights, which spoil the importance weight distribution. Hence, we clipped the highest 0.2% of the importance weights as in [Midgley al., 2023).
The results for our SPHERICAL-PROI E-ACF model are shown in Fig. 7. The reweighted Ramachan- dran plot is significantly closer to the ground truth low density regions. The KLD reduced from 1.88 10 to 1.39 10- 10
Alanine dipeptide T = 300K Although our experiments were done on alanine dipeptide at T = 800K, we also tried T" = 300K. We trained our SPHERICAL-PROJ B-ACF model the same
25

---

101
10
1013
10
10
10
(a) Ground truth
(b) Raw model samples
(c) Reweighted model samples
Figure 8: Ramachandran plot of SPHERICAL-PROI E-ACF model trained alanine dipeptide at 300 K. The raw and reweighted results are shown. Through reweighting, the KLD drops from 2.79 10T to 1.97 10 10-3
dataset used in [Midgley et al., 2023]. As shown in Fig. 8, we got good performance. Even without reweighting our model has KL.D 2.79-10-3 10- When doing reweighting, decreases to 1.97 10 For comparison, the flows trained internal coordinates presented in [Midgley et al., 2023] had KLD (7.57 + 3.80) 10 and rewcighting worsened the results given the low ESS. However, these models had about 6 times fewer parameters than our model.
C.4 Energy-based training
For both the DW4 and LJ13 experiments we use an identical flow architecture and optimizer setup to the flows trained by maximum likelihood with samples App. C.3.1. We use batch size of 128 and 1024 for DW4 and I. IJ13 respectively. We train for 20000 and 14000 epochs for DW4 and IJJ13 for the E-ACF. For the NoN-E-ACF we train for 60000 and 56000 epochs for DW4 and LJ13 respectively, which corresponds to roughly the same training time as the SPHERICAL- NoN-E-ACF is faster per iteration. Reverse ESS is estimated with 10,000 samples from the flow.
E-ACF because the
Forward ESS estimated with the test sets.
Hyper-parameters for FAB: We set the minimum replay buffer size to 64 batches and maximum size 128 batches. For the importance weight adjustment for re-weighting the buffer samples to account for the latest flow parameters, we clip to maximum of 10. For AIS we use HMC with leapfrog steps, and and intermediate distributions for DW4 and IJ13 respectively. We tune the step size of HMC to target an acceptance probability of 65%. We perform 8 buffer sampling and parameter update steps per AIS forward pass.
Table Run time for energy-based training with FAB on the DW4 and LJ13 problems. DW4 was using a GeForce RTX 2080Ti GPU and LJ13 with a V-4 TPU. Run times include evaluation that was performed intermittently throughout training. The results are averaged over seeded runs, with the standard error reported as uncertainty.
DW4
LJ13
NON-E-ACF VECTOR- PROJ E-ACF
8.9
3..
7.5
45.0
CARTESIAN- E-ACF
6.3
42.7
SPHERICAL-PROI E-ACF
7.310.0
44.711.0
26


Reviews:
[{"review_id": "fraAiN4VoH", "reviewer": "Reviewer_goAi", "summary": "This paper constructs equivariant normalizing flows that are discrete to sidestep speed issues with equivariant continuous normalizing flows. The method does this by augmenting the input with some an additional variable to avoid representation issues.", "questions": "Nothing outside of the listed weaknesses.", "limitations": "Yes", "rating": 5, "confidence": 3, "soundness": 3, "presentation": 1, "contribution": 2, "strengths": "* The method is a reasonable extension of standard equivariant normalizing flows literature that seeks to address the fundamental problem of balancing expressivity and speed.\n* The method produces good results (ie likelihoods and other metrics are trained to a very good value while not using much compute).", "weaknesses": "* I still don't understand the usage of $\\gamma_r$. Line 160-161 says that $\\gamma_r$ is a projection to an invariant space, but the definitions in Line 162 don't seem to accurately describe this (which should be $\\gamma_r \\cdot gx = \\gamma_r \\cdot x$ if I'm correct. Also, what is the definition of $\\gamma_r^{-1}$, since I assume that necessarily $\\gamma_r$ must got a lower-dimensional subspace. It would also help tremendously if the \"Choice of projection $\\gamma$\" section was written to be more mathematically precise.\n* This \"invariant feature projection\" that $\\gamma_r$ is known to be a rather weak form of constructing invariant functions. Some expressivity proof would be helpful to allay these types of concerns.\n* If $\\gamma_r$ projects to something like S^2 angles, which I think might be happening for sphere-proj, then wouldn't one need a manifold-based flow model to accurately build a flow on this? See e.g. [1, 2].\n* To evaluate densities, the paper computes using only 20 samples for a. This seems like a low amount.\n* Is there a reason why results aren't reported for DW4 or LJ13 for diffusion in table 1? I assume this is trained from samples, which should allow for diffusion.\n* In the appendix, diffusion seems to be trained for a very short amount of time (less than even the for the proposed fast ACF methods. From the literature, diffusion tends to take a while to train but this tends to improve log-likelihood.\n* For table 2, why can't the negative log-likelihood be compared? Also I think its KL divergence is lower than the proposed method (at least from the table).\n\n[1] https://arxiv.org/abs/2002.02428\n[2] https://arxiv.org/abs/2006.10254", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your feedback! We have addressed each of the weaknesses you highlighted below. We hope that this helps clarify the paper. \n## Usage of $\\gamma$:\nYou are correct; there is a typo in the text - thanks for pointing this out. The property of interest is $\\gamma_{g \\cdot r} = \\gamma_{ r} \\cdot g^{-1}$. Furthermore, we have updated text to make the definition of $\\gamma$ more clear. We have highlighted these changes in our [general response to reviewers](https://openreview.net/forum?id=KKxO6wwx8p&noteId=Narle0scqP). \n\nIn the case of the Cartesian and spherical projections the invariant space is the same dimension as the equivariant space. The vector projection invariant space is lower dimensional, but this can simply be seen as a case of the spherical projection where the two angles are held constant. The projections are standard - hence we relegate their details to the Appendix B3. For example, in the case of the Cartesian projection the projection into the invariant space is given by $x_{invariant} = R^{-1} (x_{equivariant} - o)$ where the projection is composed by an orthonormal rotation matrix R and the origin o, where both R and o are equivariant.  The inverse of this projection is then simply $x_{equivariant} = R x_{invariant} + o$. We hope this example makes $\\gamma$ more clear.\n\n## Expressivity limitations of invariant projection\nInvariant projections have been quite successful for expressive transformations of equivariant spaces. Specifically the Invariant Point Attention (IPA) used in AlphaFold 2 uses a similar idea of projection into a local invariant frame associated with each protein residue. The IPA network was also used in frame diffusion (https://arxiv.org/abs/2302.02277). This similarity has now been noted in the Discussion. \n\nAs the invariant features are of the same dimensionality as the equivariant features, information is not lost in the projection. One similar case which may be causing some confusion is GNNs that use lower dimensional representations (such as distances between points instead of vectors), which is very harmful for expressiveness. \nEmpirically, we are able to capture complex functions, like QM9 and alanine density, which are challenging even for flexible diffusion models. \n\n## Manifold flow for angles\nYes, for the spherical projection we have one angle bounded between $[0, \\pi]$ and another between $[-\\pi, \\pi]$ and we indeed use the flows from [1] that you listed (which we also cite). This is described in Appendix B.3. \n## Number of samples of $a$ for marginal density estimation\nEmpirically we find that the marginal log likelihood plateaus around this value. For example, on DW4 we tested increasing from 20 to 50 samples of $a$ per $x$ changes the marginal log likelihood is 0.01 (tested on a single trained model) and for LJ13 it made a 0.03 difference. Setting the number of samples to 20 allows us to obtain good estimates of the marginal log likelihood, while preventing the evaluation costs from becoming unnecessarily large. \n\nAs we discuss in line 244, “evaluating densities”, the variance of our density estimator vanishes when our flow is able to match the augmented target distribution exactly. This distribution is a simple Gaussian, which is often modelled well by out flow. This result explains why we can get away with a relatively small number of samples.\n## Why no diffusion for DW4 and LJ13?\nOn these datasets we report the CNF results of (Satorras et. al.). This model is equivalent to modern diffusion models, although the authors train directly by likelihood instead of score matching. \nAlthough we did train diffusion models on these datasets our results were poor for these problems (i.e. worse than the CNF results of (Satorras et. al.)).\nThe models for these problems therefore require further fine-tuning which we were not able to complete before the deadline.\nSpecifically, the noise schedule requires further tuning as we found it to dramatically effect performance.\nWe will add results with the CNF-Diffusion baseline for the camera-ready version of the paper. We expect it to get similar results to the E-CNF for these problems. \nOn QM9 positional, we were able to substantially improve upon the results of Satorras et. al. with our diffusion models. This is likely because  Satorras et. al. were unable to train to convergence. Thus we report the results for both CNFs, with CNF-Diff representing the strong baseline.\n## Diffusion trained for short\nWe found that the biggest difference for training diffusion was optimizing the networks and noise schedule, while further increases in the training time did not make a large difference. \nHowever, for the camera ready version of the paper we look into the effect of training the diffusion model for a much larger training time. \n\nThe speed of our flow is advantageuous compared to CNFs for training that requires density evaluation or sampling. This applies to training by maximum likelihood and energy-based training. Hence, in the maximum likelihood regime our flow trains for much shorter than that of Satorras et. al. \nHowever, the score-matching objective used in diffusion does not require density evaluation and is highly parellelizable. \nThus, in the training-by-samples regime we do not expect our flow to be faster for training than diffusion models. \n\n## Comparing log likelihood on Aldp, KLD of histograms\nThe spaces on which densities are defined are different on internal coordinates and on 3D coordinates making their NLL not comparable.\nYou are correct that the flow on internal coordinates does better on the KL metric than our method. This is consistent with Figure 3, the histogram from which this empirical KLD is calculated, where we see that the internal coordinates match the data the closest. The internal coordinates represent a very favourable but non-transferable representation of a molecule’s conformations. "}, {"review_id": "qJbcKWc8yz", "reviewer": "Reviewer_mxfv", "summary": "This paper tackles the problem of sampling molecule configurations using deep generative models. The learned distribution should be permutation invariant (in terms of permutations of the particles in $\\mathbb{R}^3$) and SE(3) invariant (i.e., invariant to rotations and translations applied to the particles' cartestian coordinates). At the same time, the deep generative models should be fast to sample from and allow quick density evaluation.\n\nContinuous normalizing flows (and relatedly, diffusion models) can be built with these symmetry constraints, but are slow to sample from. Coupling flows are fast at sampling/density evaluation but existing approaches rely on internal coordinates rather than explicit cartesian coordinates. Generally, the goal of this paper is to endow coupling flows with the SE(3) and permutation symmetries of existing CNFs, while preserving the speed benefits of coupling flows.\n\nCoupling transforms pose a problem for SE(3)xS_n symmetry because splitting the coordinates will break either rotation or permutation symmetry, depending on which dimension one splits on. The paper gets around this problem by introducing augmented dimensions, which are instantiated as additional observations $a$. The original observations are transformed conditioned on the augmented variables and vice versa, with equivariance enforced by having a particular structure to the transformation of x given a (Eq 4). In experiments, the paper shows that the proposed method performs similarly to CNF based methods while preserving the computation benefits of coupling flows.", "questions": "N/A", "limitations": "The paper is not particularly clear on limitations of the proposed approach.", "rating": 7, "confidence": 2, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "This paper presents a solid and principled design for making coupling layers (in a normalizing flow) SE(3) and permutation equivariant, a property of practical interest in situations such as modeling molecular configurations. The experiments show that their approach preserves the computation benefits of coupling flows and the performance benefits of symmetry priors.", "weaknesses": "Although the preliminaries are set up well, the presentation of the method starting from 3.1 is somewhat confusing and could benefit from more intuitive description (eg, of Prop 3.2) and concrete examples. ", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your review! We have addressed the issues you highlighted below. \n\n## Clarity of methods section\nWe have updated the description of our method to make it more clear. Additionally [Reviewer goAi](https://openreview.net/forum?id=KKxO6wwx8p&noteId=fraAiN4VoH) found a typo in one of the equations in the text for this section which we have now fixed. The updated description of the method can be seen in the [general response to reviewers](https://openreview.net/forum?id=KKxO6wwx8p&noteId=Narle0scqP).  We now briefly note in the text an example $\\gamma\\_r$ defined by a projection into a frame composed of a rotation matrix and origin. Due to space constraints we are currently unable to fit a more detailed description illustrating how each of the components of the flow work.  However, Figure 1 provides a (very high level) overview of the flow steps, while Appendix B.3 goes through the various projection types in detail. \n\nWe hope to add further a more detailed example of $\\gamma$ to the main text if we can make more space. We provide this example below - we hope it makes the projection-into-an-invariant-space component of the flow more clear. \n\n>  The Cartesian projection into the invariant space is given by $x_{invariant} = R^{-1} (x_{equivariant} - o)$ where the projection is composed by an orthonormal rotation matrix R and the origin o, where both R and o are equivariant.  The inverse of this projection is then simply $x_{equivariant} = R x_{invariant} + o$. \n\n\n## Unclear limitations\nWe agree that the limitations were not made sufficiently clear, hence we have added an explicit limitations subsection to our discussion. This new section is provided in [general response to reviewers](https://openreview.net/forum?id=KKxO6wwx8p&noteId=Narle0scqP). "}, {"review_id": "doAMk1umPQ", "reviewer": "Reviewer_JYCH", "summary": "This paper proposes a novel method for developing an SE(3) equivariant coupling flow that works on the Cartesian coordinates of atoms, crucially enabling quick sampling and density evaluation. The standard coupling architecture is unable to operate on the Cartesian coordinates with the Special Euclidean group SE(3) and permutation invariances of physical systems, hence the researchers introduce an advanced coupling flow that preserves these by using coordinate splits along added augmented dimensions. The method includes mapping atoms’ positions into learned SE(3) invariant bases at each layer, applying standard flow transformations such as monotonic rational-quadratic splines, and then reverting to the original basis. This proposed model demonstrated competitive performance when trained on the DW4, LJ13, and QM9-positional datasets, offering much faster sampling speeds. The researchers also became the first to learn the full Boltzmann distribution of alanine dipeptide by only modeling the Cartesian positions of its atoms.\n\nFurthermore, the study introduces a new approach to making coupling layers SE(3) equivariant by augmenting the input space with auxiliary variables that can be acted upon by SE(3). This allows the atom positions to be updated based on these auxiliary variables by first projecting the atoms into an SE(3)-invariant space, applying a standard normalizing flow transform and then projecting its output back onto the equivariant space. The researchers showcased that when trained by maximum likelihood, their flow matches the performance of both existing SE(3) CNFs and coupling flows operating on internal coordinates on molecular generation tasks. This flow is over 100 times faster for sampling than SE(3) CNFs. The study also demonstrates the flow's ability to be used in the energy-based training setting on the DW4 and LJ13 problems, showcasing its unique advantages over existing methods.", "questions": "1. Can you provide more details on the graph neural network $h$ in the main paper. It is not clear to me how this network outputs invariant parameters $\\theta$. Is this the sum of all the node embeddings?\n\n2. It is not clear to me how exactly the GNN $h$ outputs equivariant features. Are you using an EGNN and taking the velocity vectors? More detail on this in the main paper would be appreciated.\n\n3. In Fig 1. The $\\oplus$ operation is actually a concatenation rather than an addition right? This could be made more clear.\n\n4. In line 210 it is stated \"We first subtract the observations’ CoM from both the observed and augmented variable\", but previously in line  134 it is stated, \"Importantly, we do not restrict a to be zero-CoM.\" This appears to be a contradiction?", "limitations": "N/A", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "This paper has a few strengths that I think are worth mentioning:\n1. The overall idea of using an augmented normalizing flow and baking in symmetries in the product space is a novel direction and sensible idea.\n\n2. In general the presentation of the paper is of high quality (except a few notational quirks). The writing is easy to follow and is supplanted by a great architecture diagram in Fig 1.\n\n3. The experimental results show a lot of promise and validate the design of the flow.", "weaknesses": "\n\n\n**Technical Details**\n\nI had a difficult time following the proof of equivariance it uses slightly non-standard notation for defining the group action and how it works. Specifically, if $\\theta$ is invariant and the standard flow is invariant I'm not sure how this could lead to an equivariant flow. That is we often require an equivariant pushforward to map an invariant base density to an invariant target density. There are two ways you could convince me: 1.) Randomly sample rotations from $g \\in SO(3)$ and numerically prove equivariance of the overall flow is satisfied. 2.) Or you could provide a lot more detail and use more standard notation. For example, I think you mean $\\tau_{\\theta}$ is an invariant density but is actually an equivariant pushforward map. Otherwise, I'm not sure how this goes through.\n\nThe idea of using augmented variables is an interesting algorithmic choice, but using a standard coupling flow for $\\tau_{\\theta}$ is only 1 option. In fact, I believe you could use any flow architecture for $\\tau_{\\theta}$ as such it would be great to see if other mainstay flow architectures could be employed here (e.g. ResFlow, MintNet, BNAF).\n\n**Experiments**\nThe biggest question I have from the experiments section is from Tab 1. It seems for DW4 training with samples is not competitive with E-CNF, but with training with energies, we find that the proposed approach clearly beats E-CNF. Can the authors provide more intuition and justification here?\n\nOne of the main benefits of using symmetries as an inductive bias is that it leads to increased sample efficiency. This is not explored in the current manuscript. One could consider an additional experiment with samples where an equivariant model and a non-equivariant model are trained with various different dataset sizes. One would expect the equivariant model to be more performant in the low data regime. Can the authors try this experiment?\n\n**Minor**\n- In the text you use CoM-Shift but in equations, this is called ShiftCoM. Better consistency in naming conventions should improve readability.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thanks for your review! Below we have responded to your feedback and questions.\n\n## Proof of equivarance\nIndeed in our code we do test equivariance of our flow by testing that  $g \\cdot F(x, a) = F(g \\cdot x, g \\cdot a)$ (and that the log det is invariant). \n\nWe do not think that our notation for group action is non-standard. For instance https://arxiv.org/abs/2302.02277 uses very similar notation. In particular, we use $g$ to refer to a member of a group and $g \\cdot x$ to refer to it being applied to a particular object $x$. We do not denote the representations of $g$ which act on different spaces separately, as these can always be distinguished by looking at the object upon which $g$ acts. Could the reviewer please clarify which specific notation is non-standard? \n\nWe have updated the methods section to make it more clear and precise, specifically with respect to the definition of $\\gamma$. This update is described in our [general response to reviewers](https://openreview.net/forum?id=KKxO6wwx8p&noteId=Narle0scqP). \n\nLet us try to further clarify how our coupling layer is constructed. \n$\\tau$ is any bijective transformation, it is neither required to be invariant nor equivariant. \nWe apply $\\tau$ to features that live in an “invariant space”, that is, the outputs of $\\gamma_r$.\n\nThe map $\\gamma: \\psi \\times X \\rightarrow Y$ is invariant, meaning that $\\gamma_{g \\cdot r} \\cdot g \\cdot  x = \\gamma_{r} \\cdot  x^{\\ell}$. Any rotations applied to $x$ (and $r$) will be undone by $\\gamma$.\n\nAfter the bijection $\\tau$ is applied, we apply the inverse of $\\gamma_{r}$. This re-introduces any rotations present in $x$ (and $r$) that were undone by  $\\gamma_{r}$ in the first place, resulting in an equivariant map. \n\nWe do note that the wording is a bit tricky - specifically the density of the flow and base distribution is **invariant** but the flow pushforward is an **equivariant** map. We have tried to make this distinction clear in the text. \n\n\n## Using other flows besides coupling flows\nCoupling flows are fast for both sampling and density evaluation (i.e. fast for both forwards transform and inversion). This is critical for our application of interest, where a flow that has fast density evaluation and sampling can be used for energy based training. \nWe mention residual flows in the discussion noting that they are limited for our purposes by slow inversion. Similarly, the MintNet and Autoregressive Flows either have a slow forwards or inverse transform, hence are not well suited for application of interest. \n\n\n## Why better than CNF when trained with energy but not when with samples?\nWe are not certain that we are understanding the Reviewer's question here, we interpret is as asking why the E-CNF trained by maximum likelihood is worse than the E-ACF trained by energy, even though it is better than the E-ACF trained by samples. Could the reviewer please clarify the question if this is not the case. \n\nGenerally, training by energy which uses ground truth density is better than training on a small (possibly biased) dataset from MD simulation, as training by energy results in the flow fitting the ground truth target function very well, where training by maximum likelihood is limited by the size and quality of the training dataset. \nHence, the E-ACF trained by energy on DW4 obtains far better performance than when trained by maximum likelihood. The comparison between the E-CNF trained with samples, and the E-ACF trained by energy is not apples-to-apples, as both the training method and model type are being changes. However, overall this difference is mainly explained by the aforementioned benifits of training-by-energy, rather than differences in the models expresiveness. \n\n\n## Testing sample efficiency with differerent dataset sizes\nWe agree this is an interesting question. Hence, we tested this for DW4, LJ13 and Alanine Dipeptide. We summarized these results in the [general response to reviewers](https://openreview.net/forum?id=KKxO6wwx8p&noteId=Narle0scqP). Generally, we found that the equivariant flows are significantly more data efficient than the non-equivariant flows. \n\n\n## COMShift vs ShiftCOM inconsistency\nThanks for highlighting this! We have fixed it in the text. \n\n\n## Details of EGNN\nThe EGNN from sattoras et al, which has become a standard architecture in the field, outputs equivariant vectors and invariant scalars per node, see equations 3-7 in https://arxiv.org/abs/2102.09844. We use this design essentially verbatim (with a few minor tweaks that we have added to Appendix C.2.), where the input is the cartesian coordinates and node embeddings of the atoms. The invariant parameters are per node (not global) hence there is no pooling over nodes. \n\n## Figure 1 concatenation instead of addition?\nThis should be a subtraction (see Equation 5). We have fixed this in the image.\n\n## Line 210 seems to imply restricting both $x$ and $a$ to zero-CoM?\nGiven observations $x$ and augmented variables $a$, we first perform $(x - \\bar{x}, a - \\bar{x})$. \nThis restricts $x$ to the Zero-Com subspace but crucially does not restrict $a$ to be ZeroCom. We agree that this was ambiguous in the text and have edited it to make this more clear. "}, {"review_id": "H0Yv8lLx0i", "reviewer": "Reviewer_2cbB", "summary": "The paper proposes a method of producing a special Euclidean group - SE(3) - equivariant method of training coupling flows by augmenting the input space then conducting standard normalizing flow transformations in a rotationally invariant space by means of an invertible transformation. This is applicable particularly to Boltzmann distribution problems.\n\nUpdate: I have read the rebuttal, and the additional details seem convincing. I have raised my score accordingly.", "questions": "In the paper, the primary baselines and comparisons are against CNFs, primarily due to the fact that existing literature on SE(3) equivariant flows have mainly through the CNF regime (though the authors do also mention equivariant residual flows as well). How does enforcing equivariance vs not enforcing affect subsequent performance on the benchmarks in the case of coupling flows?", "limitations": "Yes.", "rating": 7, "confidence": 3, "soundness": 3, "presentation": 4, "contribution": 2, "strengths": "The proposed SE(3) Equivariant Flows offer a specialized framework for dealing with atomical problems which utilize SE(3) structure. For the task of approximating the Boltzmann distribution, the model works quite well compared to some of the baselines established, and features higher speedup. In addition, it admits a FAB-version of training the model without maximum likelihood.", "weaknesses": "Seeing the disparity between the NON E-ACF and \tthe other models on the QM9 dataset, I would have liked to see some comparisons against the E-CNF DIFF baseline on the Alanine dipeptide baseline as well. \n\nIn addition, the baseline speedup comparison is only against the diffusion model for which they claim a 500x speedup. Though this speedup might be due to the advantages proposed by the framework, a large portion of it might be due to the denoising process intrinsically taking up a lot of time. I would like to see the authors either talk about speedups compared to other standard CNFs, or along the lines of a performance/speed tradeoff that can be made with respect to the results shown in the tables.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thanks for your feedback! Below we have responded to your comments.\n\n## Diffusion on Alanine Dipeptide\nWe agree that this would be a good additional baseline. Hence, we will add it to the camera ready version of the paper. \n\n## Speedup of CNF\nThe reviewer is correct that the large relative speedup is due to \"to the (diffusion) denoising process intrinsically taking up a lot of time\". Diffusion and discrete normalising flows tradeoff flexibilty for speed; diffusion is more flexible but discrete flows are faster. We will clarify this in the main text as it is one of the main motivations for our work. \n\nWe evaluate the model trained with diffusion as a CNF (i.e. ODE not SDE) and use arguably the most standard CNF ode solver (dopri5) for evaluating the speed of the CNF sample generation. We agree that this could in principle be made faster (e.g. via Optimal transport flow matching or via trading performance for speed). However, we consider this to be outside of the scope of the paper. \n\n## Effect of enforcing vs not enforcing equivariance\nWe study the effect of enforcing vs not enforcing equivariance with our non-$SO3$ equivariant coupling flow baseline (NON-E-ACF in tables 1, 2 and 3). We fine that it consistently performs significantly worse than the equivariant models. Additionally, we have now added experiments comparing the equivariant models and non-equivariant models for different dataset sizes. We find that the equivariant models are substantially more data-efficient. These new results are summarised in our [general response to reviewers](https://openreview.net/forum?id=KKxO6wwx8p&noteId=Narle0scqP). "}, {"review_id": "i9uglQsY1G", "reviewer": "Reviewer_Xvtr", "summary": "The manuscript proposes a normalizing flow architecture that is i.) equivariant wrt to the euclidian group SE(3) and permutation group S_n and ii.) is based on coupling blocks. Existing work on this specific equivariance has relied on continuous normalizing flows which rely on a Neural ODE instead of a coupling architecture to construct a bijective map.\n\nCoupling-based architectures have the advantage that they typically faster - both in training and sampling phase. \n\nThe proposed architecture combines various previously explored ideas in an elegant manner. Specifically, it relies on augmented flows to realise the equivariance, similar to Klein et al and first proposed in by Huang et al., and achieves SE(3) symmetry by relying on zero centre of mass subspaces. \n\nThe experimental part gives a detailed comparison to existing baselines and demonstrates the advantage of the proposed architecture over CNF-based flows as well as the proposal by Klein et al. \n\nOverall, I found this paper an insightful and solid contribution as well as a nice read. The proposed architecture is definitely worth exploring in my own research. I would encourage the authors to publish their code to maximize the impact of their paper.\n", "questions": "- How is the KLD normalizer in the tables estimated? \n- Did you investigate if any mode-collapse is present for the reverse KL training? In particular, did you estimate the forward ESS, i.e. 1/E_p w where p is the target/data density and w the importance weight?\n- Cartesian projection seems to lead to numerical instabilities. Do you have an understanding of why they arise?", "limitations": "Yes, I think the limitations are sufficiently summarised although I would have appreciated a dedicated section.", "rating": 6, "confidence": 3, "soundness": 4, "presentation": 4, "contribution": 3, "strengths": "- Very clearly written manuscript\n- A thorough and extensive experimental analysis\n- Proposed architecture is a elegant combination of well-known mechanisms achieving full equivariance\n- Both likelihood and energy-based training are considered", "weaknesses": "- The ms mainly combines previously proposed ideas (albeit in an innovative and elegant fashion). In particular, its proposal is relatively similar to the one of Klein et al.\n\n", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your review - we appreciated that you find our work relevant to your research! Below we have responded to your feedback.\n\n## Similarity to Klein et al\nWe assume by Klein et al you are referring to https://arxiv.org/pdf/2302.01170.pdf.\n\nWe are aware of this work and indeed it was an inspiration for our work, but we do not think it solves the problems that we set out to tackle in our research. In particular:\n\n - Klein et. al. learn their models from samples only, while we pursue energy-based learning. Our contributions are all targeted towards making energy-based learning tractable.\n\n - Klein et. al. use a non-rotation-equivariant flow architecture. They use data augmentation to learn a rotation invariant density. Our key contribution is the introduction of a novel SE(3) equivariant flow architecture. \n\n - Klein et. al. learn conditional metropolis proposals to speed up existing MD schemes. We target the full, unconditional, Boltzmann distribution directly. \n\nIn terms of similarities, we see the largest one as both papers using coupling flows with augmented variables to obtain permutation equivariance. Could the reviewer clarify to which further similarities in both papers’ proposals they are referring to?\n\n\n## Code release\nOur code is available in the supplementary material. We will also be releasing the code repository upon publication!\n\n## KLD normalizer estimate\nFor alanine dipeptide the KLD is an empirical KLD between the histograms/Ramachandran plots in Figure 3. Because this is simply the KLD between two categorical distributions, it already normalized. This would not be the KLD between the flow and (unnormalized) target densities themselves, which which we do not report.  \n\n## Assessment of mode collapse from training by reverse KL. \nWe do not perform training via reverse KL. Specifically for training-by-energy we use the Flow Annealed Importance Sampling (FAB) training technique, which encourages mass covering rather than mode seeking. We can see that the models do not suffer from mode collapse by their test set likelihood (which becomes very low if mode collapse occurs). The models trained with FAB have a similar (or better) test set likelihood than the models trained by maximum likelihood - hence we can see that these models are fitting the target well. Forward ESS would be a good additional metric, however this requires the normalized target density where we only have the unnormalized density. \n\n## Numerical instability of Cartesian Projection\nIf vectors output from EGNN become collinear then the orthonormalization becomes numerically unstable. We use an additional loss that encourages non-collinear output vectors to mitigate this. It is described in Appendix B3. \n\n## Dedicated limitations section\nWe agree - hence we have added a limitations subsection to the discussion. This is given in our [general response to reviewers](https://openreview.net/forum?id=KKxO6wwx8p&noteId=Narle0scqP). \n"}]

Author rebuttal:
# General response
Thank you for the reviews - these have been very helpful and we have incorporated the feedback into the paper. Below we highlight key changes to the paper: **(1)** Added new results studying the data-efficiency of the equivariant vs non-equivariant variants of the flow. **(2)** Improved the clarity of the methods section. **(3)** Added explicit limitations section.

## Testing sample efficiency with differerent dataset sizes
We have run additional experiments comparing the data efficiency of the equivariant vs non-equivariant models for DW4, LJ13 and Alanine dipeptide. These are summarized below. 

**DW4 and LJ13:**
We vary the dataset sizes following Satorras et al. [2021a]. We see that the equivariant flow is more data efficient, with larger performance gaps between equivariant and non-equivariant flow variants in the lower data regime. 


| Problem        | DW4         | DW4          | DW4          | DW4 | LJ13   | LJ13          | LJ13            | LJ13 |
|-------------------|-------|-------|------|------|-------|--------|-------|-------|
| **Training dataset size**          | $10^2$           | $10^3$            | $10^4$            | $10^5$  | $10$    | $10^2$           | $10^3$            | $10^4$ |
| **E-CNF**              | **8.31**  | **8.15**  | **7.69** | 7.48 | **33.12** | 30.99  | 30.56 | 30.41 |
| **Non-E-ACF**           | 10.24 | 10.07 | 8.23 | 7.39 | 40.21 | 34.83  | 33.32 | 31.94 |
| **Vector-proj E-ACF**  | 8.65  | 8.69  | 7.95 | 7.47 | 33.64 | **30.61** | **30.19** | **30.09** |
|  **Cartesian-proj E-ACF**| 8.79  | 8.82  | 7.93 | 7.52 | 36.66 | 31.72  | 30.89 | 30.48 |
| **Spherical-proj E-ACF**   | 8.66  | 8.61  | 7.97 | **7.38** | 34.79 | 31.26  | 30.33 | 30.16 |


**Alanine dipeptide:**
We varied the training dataset size from $10^2$ to $10^6$ by subsampling the original training dataset, which has $10^6$ datapoints. We report the NLL for each model, with the best, i.e. lowest, value marked in **bold**. For all dataset sizes the equivariant models clearly outperform the non-equivariant model. Notably, if we only use 10% (vector-proj) or even just 1% (Cartesian-proj and spherical-proj) of the original dataset, our equivariant models are still better then the non-equivariant one trained on the full dataset.


| Training dataset size    | $10^2$           | $10^3$            | $10^4$            | $10^5$            | $10^6$            |
| ------------------------ | ---------------- | ----------------- | ----------------- | ----------------- | ----------------- |
| **Non-E-ACF**            | $-28.8$          | $-76.8$           | $-120.4$          | $-164.3$          | $-177.8$          |
| **Vector-proj E-ACF**    | $-31.0$          | $-111.4$          | $-169.5$          | $-183.6$          | $-188.5$          |
| **Cartesian-proj E-ACF** | $\mathbf{-65.1}$ | $\mathbf{-155.7}$ | $\mathbf{-185.9}$ | $-188.1$          | $\mathbf{-188.6}$ |
| **Spherical-proj E-ACF** | $-53.1$          | $-119.3$          | $-183.6$          | $\mathbf{-188.3}$ | $\mathbf{-188.6}$ |



## Improving methods clarity.
We have edited the Methods section to make the definition of $\gamma$ more precise, additionally we fixed a typo in one of the equations on Line 160-161. Here is the language we now use:

> $SE3 \times S_n$ equivariant coupling:
We dissociate the equivariance constraint from the tractability of the Jacobian determinant by constructing a core coupling  transformation that composes (a) an invariant map $\gamma: \psi \times X \rightarrow Y$, where $Y$ is isomorphic to $X$, and $r \in \psi$ parametrises the map. We denote the parametrised map as $\gamma_r$, e.g., a frame composed of an origin and a rotation matrix (b) a standard normalizing flow transform $\tau: Y \rightarrow  Y$, \(c\) the inverse map $\gamma^{-1}$.
Denoting the inputs with superscript $\ell$ and with $\ell+1$ for outputs, our core transformation $F: (x^{\ell}; a^{\ell}) \mapsto (x^{\ell+1}, a^{\ell+1})$ is given as
$$ x^{\ell+1} = \gamma^{-1}_{r} \cdot \tau\_{\theta} ( \gamma\_{r} \cdot x^{\ell}  ), \quad \text{with} \quad (r, \theta) = h(a^{\ell}) \\ a^{\ell+1} = a^{\ell}. $$
Here, $h$ is a (graph) neural network that returns a set of equivariant reference vectors $r$ which parametrise the map $\gamma\_{r}$, and invariant parameters $\theta$. $Y$ is a rotation invariant space. This means that any rotations applied to the inputs will be cancelled by $\gamma\_{r}$, i.e. $\gamma\_{g \cdot r} = \gamma\_{ r} \cdot g^{-1}$ or equivalently $(\gamma\_{g \cdot r})^{-1} =  g \cdot \gamma^{-1}\_{r}$ for all $g \in SO3$.
We use the inverse projection $\gamma\_{r}^{-1}$ to map the invariant features back to equivariant features.
The function $\tau\_{\theta}$ is a standard flow transformation, that we apply to the invariant features.

## Limitations
We have added a limitations subsection to our discussion to make these more clear. 
> **Limitations**: Although our flow is significantly faster than alternatives such as CNFs, the expensive EGNN forward pass required in each layer of the flow makes it more computationally expensive than flows on internal coordinates. Additionally, we found our flow to be less numerically stable than flows on internal coordinates, which we mitigate via adjustments to the loss, optimizer and EGNN network (see App. B.3, App. C.1, App. C.2). Our implementation uses the E(3) equivariant EGNN proposed by Satorras et al. [2021a]. However, recently there have been large efforts towards developing more expressive, efficient and stable EGNNs architectures [Fuchs et al., 2020, Batatia et al., 2022, Musaelian et al., 2023, Liao and Smidt, 2023]. Incorporating these into our flow may improve performance, efficiency and stability. This would be especially useful for energy-based training, where the efficiency of the flow is a critical factor. 

Meta-review: Thank you for your submission and active engagement throughout the review period. The reviewers and I are all in agreement that the paper significantly advances the use of equivariant CNFs for modeling molecular configurations.
Decision: Accept (spotlight)