Title: SE(3) Equivariant Augmented Coupling Flows
Venue: NeurIPS 2023
Accepted: yes
Keywords: Boltzmann generator

normalizing flow

diffusion

molecular dynamics
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

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.

1 Introduction

Modeling the distribution of a molecule's configurations at equilibrium, known as the Boltzmann distribution, is a 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 it 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

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 permutation 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-based training, that is, training using the Boltzmann distribution's unnormalized density [Noé et al., 2019, Stimper et al., 2022, Midgley et al., 2023]. Training by energy is critical, as it 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.e. bond distances, bond angles, and dihedral angles [Wu et al., 2020, Campbell et al., 2021, Köhler et al., 2023a, Midgley et al., 2023], which are SE(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(3) equivariance, and (iii) allows for fast sampling.

To address this gap, we propose a flexible SE(3) equivariant coupling flow that operates on the Cartesian 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) -invariant space and then applying a 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 is 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 a flexible parametric family of densities on X as the pushforward of a base density q 0 along an invertible automorphism f θ : X → X with parameters θ ∈ Θ [Papamakarios et al., 2021]. The density is given by the change of variable formula:

We can efficiently sample from the flow by sampling from q 0 and mapping these samples through f θ in a single forward pass. A popular way to construct f θ is to use coupling transforms. The D

dimensional input x ∈ X is split into two sets, transforming the first set conditional on the second, while leaving the second set unchanged:

They induce a lower triangular Jacobian, such that its determinant becomes | ∂ T ( x 1: d ; x d +1: D ) /∂ x 1: d | . Further, choosing T to have an easy to compute determinant, such as an elementwise transformation [Dinh et al., 2015, 2017, Durkan et 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 = [ x 1 , . . . , x n ] ∈ X = R 3 × n , where the rows index Cartesian coordinates and the columns index individual particles. We seek to construct flows on X endowed with the symmetries present in molecular energy functions. These are invariant to rotations and translations of x ( SE(3) ), and to permutation of atoms of the same type ( S n ). 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 t ∈ SE(3) can be decomposed into two components t = ( R,u ) where R ∈ SO(3) is a 3 × 3 rotation matrix and u ∈ R 3 represents a translation; for a coordinate v ∈ R 3 , t · v = Rv + u denotes the action of t on v . The symmetric group S n defined over a set of n atoms consists of all n ! permutations that can be performed with said atoms. Its elements σ ∈ S n act on an n -body system as σ · x = [ x σ (1) , . . . , x σ ( n ) ] .

Equivariant maps A map f : X → Y is said to be G -equivariant if it commutes with the group action · , i.e. if for any x ∈ X and g ∈ G we have f ( g · x ) = g · f ( x ) . Invariance is a special case where for any x ∈ X and g ∈ G we have f ( g · x ) = f ( x ) . There has been a plethora of recent work on constructing graph neural network functions equivariant to the action of G = SE(3) × S n [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 x ∈ X and g ∈ G we have p ( g · x )= p ( x ) . Combining an invariant base density q 0 with an equivariant invertible transform f , as in (1), yields an invariant flow density [Papamakarios et al., 2021, Köhler et al., 2020]. This gives a practical way to design invariant densities models.

Challenges in constructing SO(3) × S n invariant flow models Unfortunately, no coupling transform can be simultaneously equivariant to both permutation of the particles and their rotations; coupling splits must be performed either across particles or spatial dimensions which would break either permutation or rotational symmetry [Köhler et al., 2020, Bose et al., 2022].

Furthermore, there does not exist a 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 an invariant base distribution directly on X . Fortunately, Proposition A.1 and its converse allow us to disintegrate the probability measure into a translational measure proportional to the Lebesgue measure and an SO(3) -invariant probability measure on the subspace of R 3 × n with zero center of mass. We can drop the former and only model the latter.

3 Method: SE(3) × S n equivariant augmented coupling flow model

This section describes our main contribution, an SE(3) × S n equivariant coupling flow. We first lay the groundwork for achieving translation invariance by defining our flow density on a lowerdimensional 'zero Center of Mass (CoM)' space. To preserve permutation and rotation equivariance, we leverage the augmented flow framework of Huang et al. [2020]. Specifically, we use sets of augmented variables as a pivot for coupling transforms. Sec. 3.1 introduces a novel class of coupling transforms that achieve the aforementioned permutation and rotation equivariance by operating on atoms projected using a set of equivariant bases. Sec. 3.2 describes our choice of invariant base distribution and, finally, in Sec. 3.3, we discuss several schemes to train the augmented flow from either samples or energy functions and how to perform efficient density evaluation.

˜

x

a

x

′

˜

a

ShiftCoM

map

˜

a

≜

a

-

¯

a

¯

a

-

Core transformations

F

i

×

M

Projection

γ

r

Flow trans-

form

τ

θ

Inverse pro-

jection

γ

-

1

r

Equivariant Neural Network

Equivariant

reference

r

Invariant

parameters

θ

Translation invariance is obtained by modelling the data on the quotient space R 3 × n / R 3 ≜ ˜ X ⊆ X , where all n -body systems that only differ by a translation are 'glued' together, i.e. where x ∼ x ′ if x = x ′ + p with p ∈ R 3 . Constructing a parametric probabilistic model over ˜ 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 ≜ x -¯ x with ¯ x ≜ 1 n ∑ n i =1 [ x ] i . Thus, ˜ x lies on ˜ X , an ( n -1) × 3 dimensional hyperplane embedded in X .

Augmented variable pivoted coupling To allow for simultaneously permutation and rotation equivariant coupling transforms, we introduce augmented variables a ∈ A . Our coupling layers update the particle positions x conditioned on a and vice-versa. The augmented variables need to 'behave' similarly to x , in that they can also be acted upon by elements of SE(3) × S n . We achieve this by choosing a to be a set of k of observation-sized arrays A = X k , 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 × A w.r.t. the product measure λ ˜ X ⊗ λ A , where λ ˜ X ∈ P ( ˜ X ) and λ A ∈ P ( A ) respectively denote the Lebesgue measure on ˜ X and A . We use q ( x , a ) as a shorthand for the density of the corresponding zero-CoM projection q (˜ x , a ) . The density q is constructed to be invariant to the action of G ≜ SO(3) × S n when simultaneously applied to the observed and augmented variables, that is, q ( x , a ) = q ( g · x , g · a ) for any g ∈ G . Weaim to construct a G -equivariant flow f : ˜ X×A→ ˜ X×A on this extended space, and combine it with q 0 : ˜ X × A → R + , a G -invariant base density function, to yield the invariant flow density

Proposition 3.1 (Invariant marginal) . Assume q : X × A → R + is a G-invariant density over the probability space ( X × A , λ X ⊗ λ A ) , then q x ≜ ∫ A q ( · , a ) λ A (d a ) : X → R + is a G-invariant density w.r.t. to the measure λ X .

Proof. For any g ∈ G , x ∈ X and a ∈ A

where we used the G -invariance of q and of the measure λ A , as well as g -1 A = A .

3.1 SE(3) and permutation equivariant coupling transform

We now derive our SE(3) × S n equivariant map f : ˜ 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 ˜ x conditional on a and vice versa. Composing them yields our equivariant coupling layer illustrated in Fig. 1.

SE(3) × S n 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 γ : Ξ ×X → Y , where Y is isomorphic to X , and r ∈ Ξ parametrizes the map. We denote the parametrized map as γ r . It is followed by (b) a standard flexible normalizing flow transform, e.g. a neural spline, τ : Y → Y , and (c) the inverse map γ -1 . Denoting the inputs with superscript ℓ and with ℓ +1 for outputs, our core transformation F : ( x ℓ ; a ℓ ) ↦→ ( x ℓ +1 , a ℓ +1 ) is given as

Here, h is a (graph) neural network that returns a set of equivariant reference vectors r , which parametrize the map γ r , and invariant parameters θ . Y is a rotation invariant space. This means that any rotations applied to the inputs will be cancelled by γ r , i.e. γ g · r = γ r · g -1 or equivalently ( γ g · r ) -1 = g · γ -1 r for all g ∈ SO(3) . We use the inverse projection γ -1 r to map the invariant features back to equivariant features. The function τ θ is a standard invertible inner-transformation such as 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 h : A → X n × Θ n is SO(3) -equivariant for its first output, SO(3) -invariant for its second, and S n equivariant for both, and ( γ g · r ) -1 = g · γ -1 r for all g ∈ SO(3) , the transform F given by (4) is SO(3) × S n equivariant.

Proof. For SO(3) : We first notice that h ( g · a ) = ( g · r , θ ) , and then since ( γ g · r ) -1 = g · γ -1 r we have F ( g · x , g · a ) = ( γ g · r ) -1 · τ θ ( γ g · r · g · x ℓ ) = g · γ -1 r · τ θ ( γ r · g -1 · g · x ℓ ) = g · F ( x , a ) .

For the Jacobian of the coupling described above to be well-defined, the variable being transformed must be non-zero CoM (see App. B.1 for a derivation). Thus, although our observations live on ˜ X , for now, assume that the inputs to the transform are not zero-CoM and we will deal with this 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.

Choices of projection γ The equivariant vectors r parameterize a local (per-atom) SO(3) equivariant reference frame used in the projection γ r . We introduce three different projection strategies. (i) The first strategy, is for r to parameterize frame composed of an origin and orthonormal rotation matrix into which we project each atom's positions. We then take τ to be a dimension-wise transformation for each of the projected atoms' coordinates. We dub this method CARTESIAN-PROJ. (ii) Alternatively, we let r parameterize a origin, zenith direction and azimuth direction for spherical coordinates, as in Liu et al. [2022]. We then apply elementwise transforms to each atom's radius, polar angle and azimuthal angle. We call this SPHERICAL-PROJ. (iii) Lastly, consider a variant of SPHERICAL-PROJ 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 θ . These parametrize M core coupling transformations with a single GNN forward pass. For the CARTESIAN-PROJ 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 SE(3) × S n equivariant coupling in the ambient space rather than zero-COM subspace. In particular, before transforming our observed vector ˜ x ∈ ˜ X with F , we lift it onto X . We achieve this by swapping the center of mass between ˜ x and a . For now, assume A = X , i.e. k = 1 , with App. B.4 providing details for k > 1 . Letting ˜ A ⊆ A be the subspace where all augmented variables that differ by

a translation occupy the same point, and ˜ a ∈ ˜ A be defined analogously to ˜ x , we apply the map ShiftCoM : ˜ X × A ↦→ X × ˜ A which acts on both of its arguments by subtracting from each of them the latter's CoM, that is,

This operation is invertible, with inverse ShiftCoM(˜ a , x ) , and has unit Jacobian determinant.

Putting the building blocks together Our flow transform is built as a sequence of L blocks. Each block, described in Alg. 1, consists of two equivariant coupling layers, see Fig. 1. Our observations ˜ x ∈ ˜ X are lifted onto X with ShiftCoM , they are transformed with M core transformations ( F (1) i ) M i =1 , and ShiftCoM is applied one more time to map the observations back to the zero-CoM hyperplane. After this, our augmented variables a are transformed with ( F (2) i ) M i =1 .

Joint density evaluation q ( x , 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 q 0 , 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 SE(3) × S n Invariant base distribution

Again, we assume A = X , i.e. k = 1 , with the generalization given in App. B.4. Our invariant choice of base distribution is q 0 ( x , a ) = ˜ N ( x ; 0 , I ) N ( a ; x , η 2 I ) where x ∈ R 3 n and a ∈ R 3 n refer to x and a flattened into vectors, η 2 is a hyperparameter and we denote Gaussian distributions on ˜ X as ˜ N [Satorras et al., 2021a, Yim et al., 2023] with density

We sample from it by first sampling from a standard Gaussian N (0 , I ) and then removing the CoM. On the other hand, the distribution for a is supported on A which includes non-zero CoM points. It is centered on x , yielding joint invariance to translations. The isotropic nature of q 0 makes its density invariant to rotations, reflections, and permutations [Satorras et al., 2021a, Yim et al., 2023].

3.3 Training and likelihood evaluation

In this section, we discuss learning and density evaluation with augmented variables.

Invariant augmented target distribution Weassume the density of our observations p is SE(3) × S n invariant. Our target for augmented variables is π ( a | x ) = N ( a ; x , η 2 I ) , where η 2 matches the variance of the base Gaussian density over a . This satisfies joint invariance p ( g · x ) π ( g · a | g · x ) = p ( x ) π ( a | x ) for any g ∈ SE(3) × S n , as shown in App. B.5.

Learning from samples When data samples x ∼ p are available, we train our flow parameters by maximizing the joint likelihood, which is a lower bound on the marginal log-likelihood over observations up to a fixed constant

Learning from energy When samples are not available but we can query the unnormalized energy of a state U ( x ) , with p ( x ) ∝ exp( -U ( x )) , we can minimize the joint reverse KL divergence. By

the chain rule of the KL divergence, this upper bounds the KL between marginals

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 et al., 2023], which targets the mass covering α -divergence with α = 2 . In particular, we minimize the α -divergence over the joint which leads to an upper bound on the divergence of the marginals

To compute unbiased expectations with the augmented flow we rely on the estimator E p ( x ) [ f ( x )] = E q ( x , a ) [ w ( x , a ) f ( x )] where w ( x , a ) = p ( x ) π ( a | x ) /q ( x , a ) . Minimizing the joint α -divergence with α = 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 q ( x ) = E a ∼ π ( ·| x ) [ q ( x , a ) π ( a | x ) ] , noting that π is Gaussian and thus supported everywhere. The estimator variance vanishes when q ( a | x ) = π ( a | x ) , as shown in App. B.9.

4 Experiments

4.1 Training with samples: DW4, LJ13 and QM9 positional

First, we consider 3 problems that involve only positional information, with no additional features such as 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 particle double well energy function (DW4) and 13 particles Leonard Jones energy function (LJ13) respectively. The third problem, i.e. QM9 positional [Satorras et 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-PROJ, CARTESIAN-PROJ, SPHERICAL-PROJ) and compare them to: (1) NON-E-ACF: An augmented flow that is not rotation equivariant but is translation and permutation equivariant, as 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 on atom positions, without any projection. We train NON-E-ACF with data-augmentation 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 et al., 2023, Klein et al., 2023b]. (4) E-CNF-DIFF: An SE(3) equivariant diffusion model [Hoogeboom et al., 2022] evaluated as a continuous normalizing flow. All equivariant generative models use the SE(3) GNN architecture proposed by Satorras et al. [2021b]. The CARTESIAN-PROJ exhibited numerical instability on QM9-positional causing runs to crash early. To prevent these crashes it was trained at a lower learning rate than the other E-ACF models. App. C.3.1 provides a detailed description of these experiments.

On Tab. 1, we see that on DW4 and LJ13, E-CNF FM performs best, while E-ACF is competitive with E-CNF ML and E-CNF-DIFF. We note that both DW4 and LJ13 have biased datasets (see App. C.3.1). On QM9-positional, 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-DIFF and E-CNF FM would match the performance of SPHERICAL-PROJ on QM9-positional. The NON-E-ACF 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 5 seconds respectively. Thus, the E-ACF is faster for sampling than the CNF by more than an order of magnitude.

2

3

4

5

6

interatomic distance

0

.

00

0

.

02

0

.

04

0

.

06

normalized count

DW4

flow

data

1

2

3

interatomic distance

LJ13

2

4

6

interatomic distance

QM9 Positional

4.2 Training with samples: Alanine dipeptide

Next, we approximate the Boltzmann distribution of alanine dipeptide in an implicit solvent at temperature T = 800K . We train the models via maximum likelihood on samples generated by a replica exchange MD simulation [Mori and Okamoto, 2010], which serve as a ground truth. Our generated dataset consists of 10 6 samples in the training and validation set as well as 10 7 samples in the test set. Besides our E-ACF using the different projection schemes that we introduced in Sec. 3.1, we train the non-SO(3) equivariant flow (NON-E-ACF) with data augmentation similarly to the previous experiments. Moreover, we train a flow on internal coordinates as in Midgley et al. [2023]. The joint effective sample size (ESS) is reported for the E-ACF models which is 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 in App. C.3.2.

The results are shown in Fig. 3 and Tab. 2. All variants of our E-ACF clearly outperform the NON-EACF, as the Kullback Leibler divergence (KLD) of the Ramachandran plots is significantly lower and the NLL as well as the reverse and forward ESS, see App. B.10, are higher. The flow trained on internal coordinates is only marginally better regarding the NLL and KLD than our best models, i.e.

Ground truth

Flow on internal coordinates

E-CNF FM

E-CNF-Diff

π/2

1-01

0

10-2

一π

Non-E-ACF with data aug.

Vector-proj E-ACF

Cartesian-proj E-ACF

Spherical-proj E-ACF

π

10-3

π/2

0

10-4

一π

一π

π|2

0

π2

Ⅱ—Ⅱ

0

π|2

Ⅱ—Ⅱ

π|2

0

π|2

π一π

π|2

0

元2

T

the CARTESIAN-PROJ and SPHERICAL-PROJ E-ACF, while we outperform it considering the ESS. Note that the flow on internal coordinates explicitly models the angles ϕ and ψ , while the E-ACF operates on the underlying Cartesian coordinates. The E-ACFs outperform the E-CNF-DIFF model in all metrics, but the E-CNF trained with flow matching has a slightly lower NLL while having a significantly higher KLD on the Ramachandran plot. This could be due to a better performance on other marginals or 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 best of our knowledge, our models are the first to learn the full Boltzmann distribution of a molecule purely on Cartesian coordinates while being competitive with a flow trained on internal coordinates. Additionally, we train the best E-ACF model at T = 300K 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 be trained on the DW4 and LJ13 problems using only the target's unnormalized density with the FAB algorithm [Midgley et 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, we only report results for our flow, as well as for the NON-E-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.

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-ACF performs worse, both in terms of ESS and NLL. All models trained by maximum likelihood have a much lower ESS (see App. C.3.1). This is expected, as unlike the α = 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 a coupling scheme similar to 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 model could be made fully invariant by substituting in our flow's projection-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 equivariant 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 to 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 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 neural 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.

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 on their molecular graph, which we hope to explore in the future.

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 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-Lobato acknowledges support from a Turing AI Fellowship under grant EP/V023756/1. José Miguel Hernández-Lobato and Emile Mathieu 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.hpc.cam.ac.uk) funded by an EPSRC Tier-2 capital grant. It 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.

References

Ilyes 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 https://openreview.net/forum?id= YPpSngE-ZU .

Simon Batzner, Albert Musaelian, Lixin Sun, Mario Geiger, Jonathan P Mailoa, Mordechai Kornbluth, Nicola Molinari, Tess E Smidt, and Boris Kozinsky. E(3)-equivariant 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 Hernandez-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 1238-1248. PMLR, 2021.

Taco S. Cohen and Max Welling. Steerable cnns. International Conference on Learning Representations , 2017.

Laurent Dinh, David Krueger, and Yoshua Bengio. NICE: non-linear independent components estimation. In Yoshua Bengio and Yann LeCun, editors, 3rd International Conference on Learning Representations, ICLR 2015, San Diego, CA, USA, May 7-9, 2015, Workshop Track Proceedings , 2015. URL http://arxiv.org/abs/1410.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://openreview.net/ forum?id=HkpbnH9lx .

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, Iain Murray, and George Papamakarios. Neural spline flows. In Hanna M. Wallach, Hugo Larochelle, Alina Beygelzimer, Florence d'Alché-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-7520, 2019. URL https://proceedings.neurips.cc/ paper/2019/hash/7ac71d433f282034e088473244df8c02-Abstract.html .

Fabian B. Fuchs, Daniel E. Worrall, Volker Fischer, and Max Welling. Se(3)-transformers: 3d rototranslation equivariant attention networks. Advances in Neural Information Processing Systems 33 , 2020.

Mario Geiger and Tess E. Smidt. e3nn: Euclidean neural networks. CoRR , abs/2207.09453, 2022. doi: 10.48550/arXiv.2207.09453. URL https://doi.org/10.48550/arXiv.2207.09453 .

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/1810.01367, 2018. URL http://arxiv.org/abs/1810.01367 .

Jonathan Ho, Ajay Jain, and Pieter 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 diffusion for molecule generation in 3d. In Kamalika Chaudhuri, Stefanie Jegelka, Le Song, Csaba Szepesvári, Gang Niu, and Sivan Sabato, editors, International Conference on Machine Learning, ICML 2022, 17-23 July 2022, Baltimore, Maryland, USA , volume 162 of Proceedings of Machine Learning Research , pages 8867-8887. PMLR, 2022. URL https: //proceedings.mlr.press/v162/hoogeboom22a.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.01729, 2022. doi: 10.48550/arXiv. 2206.01729. URL https://doi.org/10.48550/arXiv.2206.01729 .

John M. Jumper, Richard Evans, Alexander Pritzel, Tim Green, Michael Figurnov, Olaf Ronneberger, Kathryn Tunyasuvunakool, Russ Bates, Augustin Zídek, Anna Potapenko, Alex Bridgland, Clemens Meyer, Simon A A Kohl, Andy Ballard, Andrew Cowie, Bernardino RomeraParedes, Stanislav Nikolov, Rishub Jain, Jonas Adler, Trevor Back, Stig Petersen, David A. Reiman, Ellen Clancy, Michal Zielinski, Martin Steinegger, Michalina Pacholska, Tamas Berghammer, 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: A method for stochastic optimization. arXiv preprint arXiv:1412.6980 , 2014.

Leon Klein, Andrew YK Foong, Tor Erlend Fjelde, Bruno Mlodozeniec, Marc Brockschmidt, Sebastian Nowozin, Frank Noé, and Ryota Tomioka. Timewarp: Transferable acceleration of molecular dynamics by learning time-coarsened dynamics. arXiv preprint arXiv:2302.01170 , 2023a.

Leon Klein, Andreas Krämer, and Frank Noé. Equivariant flow matching. arXiv preprint arXiv:2306.15030 , 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-5370. PMLR, 2020. URL http://proceedings.mlr.press/v119/ kohler20a.html .

Jonas Köhler, Yaoyi Chen, Andreas Krämer, Cecilia Clementi, and Frank Noé. Flow-matching: Efficient coarse-graining of molecular dynamics without forces. Journal of Chemical Theory and Computation , 19(3):942-952, 2023a. doi: 10.1021/acs.jctc.3c00016.

Jonas Köhler, Yaoyi Chen, Andreas Krämer, Cecilia Clementi, and Frank Noé. Flow-matching: Efficient coarse-graining of molecular dynamics without forces. Journal of Chemical Theory and Computation , 19(3):942-952, 2023b. doi: 10.1021/acs.jctc.3c00016. URL https://doi.org/ 10.1021/acs.jctc.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 on Machine Learning , 2023c.

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://openreview.net/forum?id=KwmPfARgOTD .

Yaron Lipman, Ricky T. Q. Chen, Heli Ben-Hamu, Maximilian Nickel, and Matthew Le. Flow matching for generative modeling. In The Eleventh International Conference on Learning Representations , 2023. URL https://openreview.net/forum?id=PqvMRDCJT9t .

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 on Learning Representations, ICLR 2022, Virtual Event, April 25-29, 2022 . OpenReview.net, 2022. URL https://openreview.net/forum?id=givsRXsOt9r .

Laurence Illing Midgley, Vincent Stimper, Gregor N. C. Simm, Bernhard Schölkopf, and José Miguel Hernández-Lobato. 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 J Owen, Mordechai Kornbluth, 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 of many-body systems with deep learning. Science , 365(6457):eaaw1147, 2019. doi: 10.1126/science.aaw1147. URL https://www.science.org/doi/abs/10.1126/science. 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):1-64, 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):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 on 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 on Neural Information Processing Systems 2021, NeurIPS 2021, December 6-14, 2021, virtual , pages 4181-4192, 2021a. URL https://proceedings.neurips. cc/paper/2021/hash/21b5680d80f75a616096f2e791affac6-Abstract.html .

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://proceedings. mlr.press/v139/satorras21a.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-32252, 2023.

Vincent Stimper, Bernhard Schölkopf, and José Miguel Hernández-Lobato. Resampling Base Distributions of Normalizing Flows. In Proceedings of The 25th International Conference on Artificial Intelligence and Statistics , volume 151, pages 4915-4936, 2022.

Nathaniel Thomas, Tess Smidt, Steven Kearnes, Lusann Yang, Li Li, Kai Kohlhoff, and Patrick Riley. Tensor field networks: Rotation-and translation-equivariant neural networks for 3d point clouds. arXiv preprint arXiv:1802.08219 , 2018.

Alexander Tong, Nikolay Malkin, Guillaume Huguet, Yanlei Zhang, Jarrid Rector-Brooks, 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 https://openreview.net/forum?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.net/forum?id=8pvnfTAbu1f .

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 Conference on Neural Information Processing Systems 2020, NeurIPS 2020, December 612, 2020, virtual , 2020. URL https://proceedings.neurips.cc/paper/2020/hash/ 41d80bfc327ef980528426fc810a6d7a-Abstract.html .

Zhisheng Xiao, Karsten Kreis, and Arash Vahdat. Tackling the generative learning trilemma with denoising diffusion gans. In The Tenth International Conference on Learning Representations, ICLR 2022, Virtual Event, April 25-29, 2022 . OpenReview.net, 2022. URL https://openreview. net/forum?id=JprM0p-q0Co .

Minkai Xu, Lantao Yu, Yang Song, Chence Shi, Stefano Ermon, and Jian Tang. GeoDiff: A Geometric Diffusion Model for Molecular Conformation Generation. In International Conference on Learning Representations , 2022.

Jason Yim, Brian L. Trippe, Valentin De Bortoli, Emile Mathieu, Arnaud Doucet, Regina Barzilay, and Tommi S. Jaakkola. SE(3) diffusion model with application to protein backbone generation. 202:40001-40039, 2023.

Dinghuai Zhang, Ricky Tian Qi Chen, Cheng-Hao Liu, Aaron Courville, and Yoshua Bengio. Diffusion 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/forum?id=_uCb2ynRu7Y .

A Background

A.1 Tensors and representation

We assume a group ( G, · ) , an algebraic structure that consists of a set G with a group operator · that follows the closure , inverse and associativity axioms.

A representation is an invertible linear transformation ρ ( g ) : V → V , indexed by group elements g ∈ G , and defines how group elements g acts on elements of the vector space V . It satisfies the group structure ρ ( gh ) = ρ ( g ) ρ ( h ) .

A tensor T is a geometrical object that is acted on by group elements g ∈ G in a particular way, characterized by its representation ρ : g · T = ρ ( g ) T . Typical examples include scalar with ρ triv ( g ) = 1 and vectors living in V with ρ Id ( g ) = g . Higher order tensors which live in V ×··· × V can be constructed as g · T = ( ρ Id ( g ) ⊗··· ⊗ ρ Id ( g )) vec ( T ) . A collection of tensors ( T 1 , . . . , T n ) with representations ( ρ 1 , . . . , ρ n ) can be stacked with fiber representation as the direct sum ρ = ρ 1 ⊕··· ⊕ ρ n .

A.2 Translational invariant densities

Proposition A.1 (Disintegration of measures for translation invariance [Pollard, 2002, Yim et al., 2023]) . Under mild assumptions, for every SE(3) -invariant measure µ on R 3 × n , there exists η an SO(3) -invariant probability measure on { R 3 × n | ∑ n i =1 x i = 0 } and ¯ µ proportional to the Lebesgue measure on R 3 s.t. d µ ([ x 1 , . . . , x n ]) = d¯ µ ( 1 n ∑ n i =1 x i ) × d η ([ x 1 -1 n ∑ n i =1 x i , . . . , x n -1 n ∑ n i =1 x i ]) .

The converse is also true, which motivates constructing an SO(3) -invariant probability measure on the subspace of R 3 × n with zero center of mass.

B Method

B.1 Transforming variables on a subspace

This section shows that a flow transform F : X → X applied to observations ˜ x embedded in the zero-CoM subspace ˜ X does not have a well-defined Jacobian. We will show this by constructing such a transform, which leaves zero-CoM input unchanged, but for which the log-determinant is arbitrary.

Without loss of generality, we will assume in this subsection and the next that our observed and augmented variables are one dimensional, i.e., X = A = R n . Consider the orthonormal basis V ∈ R n × n constructed by concatenating orthonormal basis vectors v i ∈ R n as [ v 1 , . . . , v n ] T . We choose v 1 to be equal to [1 , 1 , . . . , 1] T / √ n such that inner products with v 1 retrieve the CoM of a system of particles up to multiplication by a scalar: v T 1 x = √ n ¯ x and v T 1 ˜ x = 0 . Let F ( x ) = V T SV x be a mapping which rotates onto the basis defined by V , multiplies by the diagonal matrix S = diag ([ s, 1 , . . . , 1]) and then rotates back to the original space with V T . This transformation only acts on the CoM. Applying this transformation to a zero-CoM variable ˜ x we have

leaving the zero CoM input unchanged. However, the log-determinant is

since for the unitary matrices | V | = | V T | = 1 . Because s is arbitrary, so 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 to a degenerate solution of s → 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 observations to the non-zero-CoM space before applying said transformations. We do this using the ShiftCoM operation, described in Sec. 3 and below.

B.2 Jacobian of ShiftCoM

We now show our ShiftCoM : ˜ X ×A ↦→ X × ˜ A transformation to have unit Jacobian log determinant. We first, re-state the function definition

Wenowre-write it using the orthonormal basis V ∈ R n × n defined in the previous subsection. For this, we use I 2 ⊗ V ∈ R 2 n × 2 n to refer to a 2 × 2 block diagonal matrix with V in both blocks, and we also stack the inputs and outputs of ShiftCoM , yielding ShiftCoM([˜ x , a ]) = ( I 2 ⊗ V ) T P ( I 2 ⊗ V )[ ˜ x , a ] , where P is a permutation matrix that exchanges the first and n +1 th elements of the vector it acts upon. It also flips this element's sign. To see this note

For the determinant we use that | I 2 | = 1 and | V | = 1 , and thus the determinant of their Kronecker product will be equal to one | ( I 2 ⊗ V ) | = 1 . Combining this with the fact that permutation matrices and matrices that flip signs have unit determinant, we arrive at

It is worth noting that ShiftCoM 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 not 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 a 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 a flow transformation in this space and a back-projection.

We denote the inputs with superscript ℓ and use ℓ +1 for outputs to define our transformation block F : ( x ℓ ; a ℓ ) ↦→ ( x ℓ +1 , a ℓ +1 ) . Without loss of generality, we assume that we are transforming x ℓ . Then, F takes the following form.

Here, h is a (graph) neural network that returns a set of u equivariant reference vectors r = [ r 1 , . . . , r u ] , and invariant parameters θ . The equivariant vectors r parametrize γ r , a projection operator onto a rotation invariant feature space. This means that any rotations applied to the inputs will be cancelled by γ r , i.e. γ g · r = γ -1 r · g -1 or equivalently ( γ g · r ) -1 = g · γ -1 r for all g ∈ SO(3) .

In the next paragraphs, we discuss three projection operations and their respective flow transformations.

Cartesian projection Here, we predict a 3D Cartesian coordinate system as illustrated in Fig. 4a. The coordinate system has an origin o = r 1 and an orthonormal basis Q , both being equivariant. Q is constructed from two equivariant vectors v 1 = r 2 -r 1 and v 2 = r 3 -r 1 in the following way. To obtain the first basis vector, we simply normalize v 1 , i.e. b 1 = v 1 ∥ v 1 ∥ . We get the second basis vector via a Gram-Schmidt [Petersen, 2012] iteration, i.e. we compute the component of v 2 orthogonal to v 1 , which is v 2 , ⊥ = v 2 -b ⊤ 1 v 2 , and then normalize it to obtain b 2 = v 2 , ⊥ ∥ v 2 , ⊥ ∥ . Note that this is ill-posed if v 1 and v 2 are collinear. To avoid this case, we introduce an auxiliary loss,

b

1

b

2

b

3

r

1

=

o

r

2

v

1

r

3

v

2

v

2

,

⊥

x

i

(

x

i

)

1

(

x

i

)

2

(

x

i

)

3

(a) Cartesian projection.

b

2

b

3

b

1

o

x

i

ρ

i

x

i

∥

x

i

⊥

φ

i

β

i

(b) Spherical projection.

which uses a logarithmic barrier function on the angle between the vectors to prevent them from being collinear. ϵ is a small constant that prevents the aux-loss from being equal to infinity. Finally, the third basis vector is b 3 = b 1 × b 2 , and we can assemble Q = [ b 1 , b 2 , b 3 ] .

The projection into this coordinate system and the respective inverse operation are given by

where x p is the transformed projected variable.

In the projected space, we can apply any standard flow transform, as long as it acts on each particle individually. We choose τ θ to an affine transformation as in the Real NVP architecture [Dinh et al., 2017].

Note that since b 3 is constructed via a cross product, it is not parity equivariant, which is why the entire transformation block is not parity equivariant either. However, since sign( b ⊤ 3 r 2 ) is a pseudo-scalar, we can replace b 3 with b ′ 3 = sign( b ⊤ 3 r 2 ) b 3 , which is a 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 τ θ , which we show below.

Spherical projection For the spherical projection, which is illustrated in Fig. 4b, we compute an origin o and three Cartesian basis vectors b 1 , b 2 , and b 3 as in the Cartesian projection. Instead of

computing the Cartesian coordinates of each x i , we compute the spherical coordinates ( ρ, β, φ ) . Hence, for the i -th atom the projection operation and its inverse is given by

Here, we used x i ∥ = b ⊤ 1 x i , x i ⊥ = x i -x i ∥ , and x i p = ( ρ i , β i , φ i ) for brevity.

The spherical coordinates are transformed with monotonic rational-quadratic splines [Durkan et al., 2019]. The radial coordinate ρ 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 the 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

For the inverse, we can just raise this term to the power of -1 .

Vector projection This is a special case of the spherical projection, where we only transform the radial coordinate but leave the angles the same. Since in practice we only need to compute ρ , we call this case vector projection. Note that just predicting the origin o as a reference with the equivariant neural network is sufficient here. Unlike the Cartesian and spherical projections, the vector projection does not rely on the cross product, and hence is parity equivariant by default.

B.4 Multiple augmented variables

Consider augmented variables consisting of multiple sets of observation-sized arrays, i.e. A = X k 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 x . We thus have π ( a | x ) = ∏ k i =1 N ( a i ; x , σ 2 I ) and q 0 ( x , a ) = ˜ N ( x ; 0 , I ) ∏ k i =1 N ( a i ; x , σ 2 I ) .

For our flow block, we group our variables into two sets (˜ x , a 1 , . . . , a ( k -1) / 2 ) and ( a ( k -1) / 2+1 , . . . , a k ) , and update one set conditional on the other, and vice-versa, while always ensuring the arrays being transformed have non-zero CoM.

Initially, a is non-zero-CoM with each augmented variable a i having a CoM located at a different point in space. Before the first coupling transform F (1) M , the ShiftCoM function applies the mapping

taking a ( k -1) / 2+1 to the zero-COM hyper-plane, while x and every other augmented variable remain non-zero-CoM. Before the second coupling transform F (2) M , we apply the ShiftCoM maps in reverse, taking x back to the zero-COM hyper-plane.

The coupling transforms F (1) M and F (2) M work as previously, except now they transform ( k +1) / 2 sets of variables (holding the other set constant), where for each of the k variables, each of the n nodes get their own local reference frame and inner flow transform τ θ .

B.5 Invariant base and target distributions

We now show the SE(3) × S n invariance of our extended base distribution choice, which trivial extends to our choice of augmented target distribution. Consider the choice π ( a | x ) = N ( a ; x , η 2 I ) for η 2 a scalar hyperparameter. With this choice, the constraint of interest is satisfied since the Isotropic Gaussian is rationally invariant and centered at x . To see this, let t ∈ SE(3) be decomposed

into two components t = ( R,u ) where R ∈ SO(3) is a 3 × 3 rotation matrix and u ∈ R 3 represents a translation,

An analogous argument holds for permutations, since any σ ∈ S n can be represented as a unitary n × n matrix which is simultaneously applied to a and x .

B.6 Maximum likelihood objective

Wetrain our models from samples by maximising the joint log-density E x ∼ p ( x ) , a ∼ π ( a | x ) [log q ( x , a )] . This can be shown to be a lower bound on the expected marginal log-density over observations, which is the quantity targeted by models without augmented variables as

where we drop the latter term from our training objective since it is constant w.r.t. flow parameters. The bound becomes tight when q ( a | x ) = π ( a | x ) .

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

The looseness in the bound is E q ( x ) D KL ( q ( a | x ) || π ( a | x )) which again vanishes when q ( a | x ) = π ( a | x ) .

B.8 Alpha divergence objective

We now show that our joint objective represents an analogous upper bound for the α divergence when α = 2 , i.e., the objective targeted by FAB [Midgley et al., 2023]. Let U ( x ) be the unnormalized energy function of interest evaluated at state x , with p ( x ) ∝ exp( -U ( x )) . We have

which also becomes tight when q ( a | x ) = π ( a | x ) . The inequality is true because ∫ π ( a | x ) 2 q ( a | x ) d a > 0 . To see this, take any densities π and q , and apply Jensen's inequality as

B.9 Importance weighted estimator

We define the following estimator w := q ( x , a ) π ( a | x ) with a ∼ π ( ·| x ) . Assuming that whenever π ( a | x ) = 0 then q ( x , a ) = 0 , we have that this estimator is trivially unbiased:

and that is variance is given by

thus if q ( a | x ) = π ( a | x ) , V [ w ] = 0 .

The log marginal likelihood may therefore be estimated with

which becomes accurate as 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 x with ( x , a ) .

The reverse effective sample size is given by

where w ( x i ) = ˜ p ( x i ) /q ( x i ) and ˜ p is the unnormalized version of p [Owen, 2013].

The forward ESS may be estimated with

where the normalizing constant Z = ˜ p/p may be estimated using

B.11 Conversion of internal coordinate flow density to density on ZeroCoM Cartesian space with global rotation represented

The density of the flow on internal coordinates is not directly comparable to the flows on Cartesian coordinates, as the variables live in a different space (bond distances, angles and torsions instead

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 is 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 rotation-invariant representation of internal coordinates in Cartesian space places the first atom at the origin, the second atom a distance of b 1 along the X-axis (i.e. at position { b 1 , 0 , 0 } ), and the third atom on the XY plane, at a distance of b 2 from the origin at an angle of a 2 with respect to the x-axis (i.e. at position { cos( a 2 ) · b 2 , sin( a 2 ) · b 2 , 0 } ). We define this representation of Cartesian coordinates as z ∈ Z , where we drop all coordinates of the first atom (atom 0), the y and z 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 a space that has degrees of freedom representing global rotations, which we refer to as w ∈ W , we simply sample a rotation uniformly over SO 3 and apply this to w .

The rotation has three degrees of freedom, which we represent as follows; First, we apply rotation about the X-axis of o 1 sampled uniformly over [ -π, π ] . Secondly, we rotate about the Z-Axis by sampling o 2 uniformly between [ -1 , 1] and rotating by sin -1 ( o 2 ) . Finally, we rotate about the Y-axis by o 3 sampled uniformly between [ -π, π ] . Together this ensures that the resultant orientation of the molecule in Cartesian space has a uniform density over SO 3 . We denote the concatenation of these random variables that parameterize the rotation as o ∈ O where O = { [ -π, π ] , [ -1 , 1] , [ -π, π ] } .

The density on W is then given by

where g : Z × O → W is a bijector which applies the rotation parameterized by o to z .

It can be shown that the log jacobian determinant of g is 2 log( b 1 ) + log( b 2 ) + log(sin( a 2 )) . Furthermore, log q ( o ) is a constant value of -log 8 π .

Conversion density on Zero-CoM space After applying the above transform we have w that includes rotations, but is centered on atom 0. 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 Cartesian coordinates, as these contain redundant information. We note that for the representation w ∈ W the first atom's coordinates are at the origin, and for the Zero-CoM space x ∈ ˜ X they are at -∑ n -1 i =1 x i . The bijector h : W → ˜ X is given by

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,

C Experiments: Details and Further Results

We provide the details of the hyper-parameters and compute used for each experiment. Additionally, we provide the code https://github.com/lollcat/se3-augmented-coupling-flows

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 a 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 is 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 a softmax to 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 at each node, with interaction across channels. This is needed as the local reference frame r 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 we provide a description of the hyper-parameters used within each experiment. Additionally Tab. 4 provides the run-times for each experiment.

Datasets : For DW4 and LJ13 we use a 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 QM9-positional we use a training, validation and test set of 13831, 2501 and 1,813 samples respectively also following Satorras et al. [2021a].

DW4

LJ4

0.05

train

normalized count

count

0.06

0.04

test

0.03

lized

0.04

0.02

normal

0.02

0.01

0.00

0.00

2

3

4

5

6

1.0

1.5

2.0

2.5

3.0

3.5

4.0

interatomicdistance

interatomicdistance

Flow experiment details All flow models using 12 blocks (see Alg. 1 for the definition of a block). For the noise scale of the augmented variables sampled from N ( a ; x , η 2 I ) , we use η = 0 . 1 for both the base of the flow, and augmented target distribution. For the CARTESIAN-PROJ E-ACF and NON-E-ACF we use RealNVP [Dinh et al., 2017] for the flow transform, while for CARTESIAN-PROJ and VECTOR-PROJ we use a rational quadratic spline transform [Durkan et al., 2019]. For the spline we use 8 bins for all flows. We use 1 set of augmented variables for DW4 and LJ13, and 3 sets

of augmented variables for QM9-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-ACF 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 a 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 a final learning rate of 0.00006. With the above learning rate schedule for QM9 the CARTESIAN-PROJ E-ACF was unstable during training which resulted in crashed runs (NaN loss). Hence, it was instead trained with a 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-positional respectively. We train the NON-E-ACF for 4 times as many epochs for all problems, which results in roughly the same run time as the SPHERICAL-PROJ E-ACF. For the CARTESIAN-PROJ and SPHERICAL-PROJ E-ACF, we use a 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 a per test-set sample of x .

Continuous normalizing flow / diffusion experiment details For the E-CNF ML we report the results from Satorras et al. [2021a]. For the E-CNF 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 we use an EGNN with 3 message passing layers, with a 3-hidden layer MLP with 128 units per layer. For QM9-positional we use an EGNN with 5 message passing layers, with an 4-hidden layer MLP with 256 units per layer. For all problems we use Adam optimizer with a cosine learning rate decay from 1e-4 to 0. 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 1e-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 et al., 2018] as the exact divergence is computationally expensive.

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 3 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 a linear combination of the latter 196 vectors. This model has around 90 Mlearnable parameters. We use the AdamW optimiser with a 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 · 10 -4 for the first 15% of training before being decreased to 10 -6 with a cosine 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 -3 . For sampling and evaluating log-densities we use the dopri5 ODE solver with a maximum number of evaluations of 1000.

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 | x ) .

Varying the number of training samples for DW4 and LJ13 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 a larger number of epochs before it overfits. We see that the performance gap between the non-equivariant and equivariant flow models is larger in the lower data regime.

C.3.2 Alanine dipeptide

中

Y

t

The molecule alanine dipeptide and its dihedral angles ϕ and ψ , 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 T = 800K . The data used for maximum likelihood training and testing was generated with a replica exchange MDsimulation [Mori and Okamoto, 2010] with 21 systems having temperatures from 300 K until 1300 K with 50 K temperature difference between each of them.

Each ACF has 20 layers with 3 core transformations per layer for the VECTOR-PROJ flow and 1 for the CARTESIAN-PROJ and the SPHERICAL-PROJ flow. The flow trained on internal coordinates has the same architecture as the models in Midgley et al. [2023], i.e. it is a coupling neural spline flow with 12 layers. The models were trained for 50 epochs.

(a) Ground truth

(b) Raw model samples

(c) Reweighted model samples

10-1

10-1

10-1

EIN

EIN

EIN

10-2

10-2

10-2

0

0

o

10-3

10-3

10-3

10-4

10-4

10-4

ⅡI

The E-CNF and the E-CNF-DIFF model have E(n) GNN with 4 message passing layers, each using a 3-layered MLP with 128 hidden units each, as a vector field and score function, respectively. Both were 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 one almost exclusively found in nature.

To plot and compute the KLD of the Ramachandran plots, we drew 10 7 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 6 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-E-ACFon subsets of 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-ACF trained on the full dataset performs worse than the VECTOR-PROJ E-ACF trained on 10% of the data and the CARTESIAN-PROJ and SPHERICAL-PROJ E-ACF trained on just 1% of the data.

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 7 , there were a 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 et al., 2023].

The results for our SPHERICAL-PROJ E-ACF model are shown in Fig. 7. The reweighted Ramachandran plot is significantly closer to the ground truth in low density regions. The KLD is reduced from 1 . 88 · 10 -3 to 1 . 39 · 10 -3 .

Alanine dipeptide at T = 300 K Although our experiments were done on alanine dipeptide at T = 800 K, we also tried T = 300 K. We trained our SPHERICAL-PROJ E-ACF model on the same

(a) Ground truth

(b) Raw model samples

(c) Reweighted model samples

10-1

10-1

10-1

EIN

EN

EN

10-2

10-2

10-2

0

4

0

10-3

10-3

10-3

10-4

10-4

10-4

dataset used in [Midgley et al., 2023]. As shown in Fig. 8, we got good performance. Even without reweighting our model has a KLD of 2 . 79 · 10 -3 . When doing reweighting, it decreases to 1 . 97 · 10 -3 . For comparison, the flows trained on internal coordinates presented in [Midgley et al., 2023] had a KLD of (7 . 57 ± 3 . 80) · 10 -3 , and reweighting 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 a batch size of 128 and 1024 for DW4 and LJ13 respectively. We train for 20000 and 14000 epochs for DW4 and LJ13 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-PROJ E-ACF because the NON-E-ACF is faster per iteration. Reverse ESS is estimated with 10,000 samples from the flow. Forward ESS is estimated with the test sets.

Hyper-parameters for FAB : We set the minimum replay buffer size to 64 batches and maximum size to 128 batches. For the importance weight adjustment for re-weighting the buffer samples to account for the latest flow parameters, we clip to a maximum of 10. For AIS we use HMC with 5 leapfrog steps, and 2 and 8 intermediate distributions for DW4 and LJ13 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.

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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\nSTRENGTHS:\n- 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\n\nWEAKNESSES:\n- 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\nQUESTIONS:\n- 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?\n\nLIMITATIONS:\nYes, I think the limitations are sufficiently summarised although I would have appreciated a dedicated section.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 4, "DATE": "2023-06-22"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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.\n\nSTRENGTHS:\nThe 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.\n\nWEAKNESSES:\nSeeing 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.\n\nQUESTIONS:\nIn 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?\n\nLIMITATIONS:\nYes.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-02"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS:\nThis 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.\n\nWEAKNESSES:\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.\n\nQUESTIONS:\n1. 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?\n\nLIMITATIONS:\nN/A", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-04"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS:\nThis 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.\n\nWEAKNESSES:\nAlthough 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.\n\nQUESTIONS:\nN/A\n\nLIMITATIONS:\nThe paper is not particularly clear on limitations of the proposed approach.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 2, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-07"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS:\n* 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).\n\nWEAKNESSES:\n* 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\n\nQUESTIONS:\nNothing outside of the listed weaknesses.\n\nLIMITATIONS:\nYes", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-14"}