E(n) Equivariant Normalizing Flows
Authors: Victor Garcia Satorras
Emiel Hoogeboom
Fabian Bernd Fuchs
Ingmar Posner
Max Welling
Venue: NeurIPS 2021 Oral (2021)
Keywords: equivariance
normalizing flows
molecule generation
generative models
graph neural networks

Abstract: This paper introduces a generative model equivariant to Euclidean symmetries: E(n) Equivariant Normalizing Flows (E-NFs). To construct E-NFs, we take the discriminative E(n) graph neural networks and integrate them as a differential equation to obtain an invertible equivariant function: a continuous-time normalizing flow. We demonstrate that E-NFs considerably outperform baselines and existing methods from the literature on particle systems such as DW4 and LJ13, and on molecules from QM9 in terms of log-likelihood. To the best of our knowledge, this is the first flow that jointly generates molecule features and positions in 3D.

E(n) Equivariant Normalizing Flows
Victor Garcia Satorras', Emiel Hoogeboom" !, Fabian B. Fuchs",
"agreboomm",Fabbiin Ingmar Posner", Max Welling'
UvA-Bosch Delta Lab, University of Amsterdam', Department of Engineering Science, University of Oxford
v. w.garciasatoorrasuuaa.aa nl, e. hoogeboom@uva.nng wtw.::/waaaa....iaaaaaaaaaooooea fabian@robots ox ac uk
Abstract
This paper introduces generative model equivariant to Euclidean symmetries: E(n) Equivariant Normalizing Flows (E-NFs). To construct E-NFs, we take the discriminative E(n) graph neural networks and integrate them as differential equation to obtain invertible equivariant function: a continuous-time normalizing flow. demonstrate that E-NFs considerably outperform baselines and existing methods from the literature particle systems such DW4 and LJ13, and on molecules from QM9 in terms of bog-liielhhood. To the best of our knowledge, this is the first flow that jointly generates molecule features and positions in 3D.
Introduction
Leveraging the structure of the data has long been core design principle for building neural networks. Convolutional layers for images famously do so by being translation equivariant and therefore incorporating the symmetries of pixel grid. Analogously, for discriminatory machine learning tasks on 3D coordinate data, taking into account the symmetries of data has significantly improved performance (Thomas al. 2018 Anderson al. 2019: Finzi et al. 2020 Fuchs al.) 2020; Klicpera et al., 2020). One might say that equivariance has been proven an effective tool to build inductive bias into the network about the concept of 3D coordinates. However, for generative tasks, e.g., sampling new molecular structures, the development of efficient yet powerful rotation equivariant approaches-wee while having made great progress is still in its infancy.
recent method called E(n) Equivariant Graph Neural Networks (EGNNs) (Satorras et al. 2021 is both computationally cheap and effective in regression and classification tasks for molecular data, while being equivariant to Euclidean symmetries. However, this model is only able to discriminate features on nodes, and cannot generate new molecular structures.
In this paper, we introduce E(n) Equivariant Normalizing Flows (E-NFs): A generative model for E(n) Equivariant data such molecules in 3D. To construct E-NFs parametrize continuous-tima flow, where the first-order derivative is modelled by an EGNN. We adapt EGNNs so that they are stable when utilized as derivative. In addition, literature to lift the discrete features of nodes to continuous space. We show that our proposed flow model significantly outperforms its non-equivariant variants and previous equívariant generative methods (Köhler et al. 2020). Additionally, we apply our method to molecule generation and we show that method is able generate realistic molecules when trained the QM9 dataset.
recent advances in the dequantization
Equal contribution.
35th Conference on Neural Information Processing Systems (NcurIPS 2021).

---

1)
pv h)
00:22: 2/ 1 x,
Figure 1: Overview of method in the sampling direction. An equivariant invertible function go has learned to map samples from a Gaussian distribution to molecules 3D, described by x, h.
2 Background
Normalizing Flows A learnable invertible transformation x = go( and a simple base distribution Pz (such as normal distribution) yield a complex distribution px. Let fe = 9g then the likelihood of datapoint x under px can exactly computed using the change of variables formula:
px (x) pz(() (y()), where f(x),
(I)
where Jr is the Jacobian of fa- A particular type of normalizing flows are continuous- normalizing
flows (Chen al. 2017 2018b: Grathwohl al. 2018). These flows use conceptual time direction to specify an invertible transformation as a differential equation. The first order differential is
predicted by a neural network é, referred to as the dynamics. The continuous-tiin change of
variables formula is given by:
log (x) log p2))) + (.'' 'Ir (x) where - x+ f. (x(t))dt,
(2)
where x(0) and x(1) In practice, Chen al. 201 8b):
Grathwohl al. (0018 estimate
the trace using Hutchinson's 's trace estimator, and the integrals can be straightforwardly computed using the torchdiffeg package written by [Cen al. 2018b). the dynamics for faster training and more stable solutions (Finlay et al. 2020). Continuous-time normalizing flows are desirable because the constraints that need to be enforced on & relatively mild: á only needs to be high order differentiable and Lipschitz continuous, with possibly large Lipschitz constant.
often desired to regularize
Equivariance in Normalizing Flows tions is often invariance whereas for transformations it is equivariance. Concretely, when function go is equivariant and a base distribution is invariant, then the distribution px given by x ga (z) where z pz also invariant (Köhler et al. 2020). As a result, one can design expressive invariant and an equivariant function ge. Furthermore,
Note that in this context, the desirable property for distribu-
distributions using an invariant base distribution when ge is restricted to be bijective and fo = g then equivariance of go implies equivariance of fo. In addition, the likelihood px be directly computed using the change of variables formula.
2.1 Equivariance
Equivariance of a function under group G defined T,(f(x) f())(x)) for all g E G, where Sy. T, are transformations related to the group element 9. where in our case S, T, will be the same. In other words, f being equivariant means that first applying the transformation S, on and then / yields the same result as first applying / on and then transforming using Tg.
In this paper we focus symmetries of the n-dimensional Fuclidean group, referred to E(n). An important property of this group is that its transformations preserve Euclidean distances. The transformations can be described as rotations, reflections and translations. Specifically, given set of points x = (XI...., xM) E R M embedded in an n-dimensional space, an orthogonal matrix R E _.x* and a translation vector I we can say f is rotation (and reflection) equivariant if, for z fxx we have Rz f (Rx), where Rx is the shorthand for (RxI,.... Rx,) ). Similarly f
2To be precise. in matrix multiplication notation (Rxi,.... RXM) = XR for simplicity we use the notation Rx = (Rxi,..., Rx and see R as operation instead of a direct multiplication on the entire

---

is translation equivariant if z +t= ft) where x + t is the shorthand for (x1 +t....... XM t). In this work we consider data defined on vertices of graph = (x. h). which in addition to the position coordinates x, also includes node features h E R 1a (for example temperature or atom type). Features h have the property that they are invariant to E(n) transformations, while they do affect x. In other words, rotations and translations of x do not modify h. In summary, given function f x. h + Zz: Z), we require equivariance of f with respect to the Euclidean group E(n) so that for all orthogonal matrices R and translations t:
R2, + t, 2/. [[Rx + h)
(3)
E(n) Equivariant Graph Neural Networks (EGNN) (Saooraas et al. 2021) consider a graph 9 = (V,E) with nodes oy E V and edges eij. Each node u, associated with position vector x; and node features h, as the ones defined in previous paragraphs. Then, an E(n) Equivariant Graph Convolutional Layer (EGCL) takes as input the set of node embeddingg h (h/..... - I. coordiante embeddings x' = (x....... XM-1} at layer and edge information E = (eg) and outputs a transformation on h1+1 and x Concisely: h xx/1_=+_____ = EGCL(h', x', EJ This layer satisfies the equivariant constraint defined in Equation g The equations that define this layer are the following:
mig = De (h/, h2, ||x] - -=11()) and
m, Σ
Mij
(4)
jtt
(xc xc)) Oz (min) and
hft on (h), m.)
(5)
A neural network h², x = EGNN(h", x"] composed of L layers will define the dynamics of our ODE flow. In experiments are not provided with an adjacency matrix, but only node information. Therefore, we use the edge inference module etj = Qint (m:y) introduced in the EGNN paper that outputs a soft estimation of the edges where @iat : gnf + [0, 1|' resembles a linear layer followed by sigmoid function. This behaves attention mechanism over neighbor nodes.
3 Related Work
Group equivariant neural networks (Cohen and Welling 2016 2017 Dieleman et al. 2016) have
demonstrated their effectiveness in wide variety of tasks. A growing body of literature is finding neural networks that are equivariant to transformations in Euclidean space Thomas et al. 2018
Fuchs et al. 2020 Horie al. 2020 Finzi al.) 2020 Hutchinson al. 2020 Satorras al.
2021
These methods have proven their efficacy in discriminative tasks and modelling dynamical sytems.
On the other hand, graph neural networks (Bruna al. 2013 Kipf and Welling 2016) can be seen as networks equivariant to permutations. Often, the methods that study Euclidean equivariance operate on point clouds embedded in Fuclidean space, and they also incorporate permutation equivariance.
Normalizing Flows (Rippel and Adams 2013, Rezende and Mohamed 2013: Dinh et al. 2015) are an attractive class of generative models since they admit exact likelihood computation and can be designed for fast inference and sampling. Notably Chen et al. (2018a b); Grathwohl et al. 2018) introduced continuous-ting normalizing flows, flow that is parametrized via a first-order derivative over time. This class of flows is useful because of the mild constraints on the parametrization function compared to other flow approaches (Dinh et al. 2017| Kingma and Dhariwal 2018).
There are several specialized methods for molecule generation: Gebauer et al. (019) generate 3D molecules iteratively via an autoregressive approach, but discretize positions and use additional focus tokens which makes them incomparable in log-likclihood, Klomez-Bombarede et al. (016)); You et al.
22018 Lia et al. (2019) generate discrete graph structures instead of coordinate space and Xu ct al. (202 generate molecule positions only, Some flows in literature Noč et al. (2018); Li et al.
(2019) y Köhler et al. 2020) model positional data, but not in combination with discrete properties.
Recently various forms of normalizing flows for equivariance have been proposed: Köhler al.)
(2019 2020) propose flows for positional data with Euclidean symmetries, |Rezndee et al. (211)) introduce a Hamiltonian flow which can be designed to be equivariant using an invariant Hamiltonian function. Liu et al. (2019); Bilos and Günnemann (2020) model distributions over graphs, with flows equivariant to permutations. [oodd et al. (2020) introduce equivariant flows for SU(n) symmetries. For adversarial networks, Dey et al, (2021 introduced a generative network equivariant to 90 degree
3

---

lift to continuous
ode integrate
+
st
h
loog g(Ah h3
log P2 Rd)
Jog (ug.Pu((M)
Figure 2: Overview of the training procedure: The discrete h is lifted to continuous h. Then the variables x, h are transformed by an ODE 2 To get a lower bound on log (v h) we sum the variational term log g/)/h), the volume term from the ODE Jo Tr (z( the log- likelihood of the latent representation on Gaussian log pz (Ze Xx). and the log- oog-likelhood of the molecule size log pMo (M). To train the model, the sum of these terms maximized.
image rotations and reflections. Our work differs from these approaches in that design a general- purpose normalizing flow for sets of nodes that contain both positional and invariant features while remaining equivariant E(n) transformations. This combination of positional and non-positional information results in a more expressive message passing scheme that defines the dynamics of our ODE, A by-product of our general approach is that also allows us to jointly model features and structure of molecules in 3D space without any additional domain knowledge.
4 Method: E(n) Equivariant Normalizing Flows
In this section propose our E(n) Equivariant Normalizing Flows, probabilistic model for data with Euclidean symmetry. The generative process defined by a simple invariant distribution pz (Z.. 2n (such Gaussian) over latent representation of positions 2₂ R Mx and a latent representation of invariant features Z/, E IR M x from which an invertible transformation ge (z.: z,.) = (x. h) generates V (x. h). These nodes consist of node features h e R M xnf and position coordinates x E MR embedded in a n-ddiennionall space. Together, pz and go define a distribution pv over node positions and invariant features.
Since we will restrict to be invertible, fe 9o exists and allows us to map from data space to latent space. This makes it possible directly optimize the implied likelihood py of a datapoint utilizing the change of variables formula:
(v(V) Pv h) = Pz(ff det Jyl Pz(2,, 2n)1 det Jyl,
(6)
where J, a(2%, is the Jacobian, where all tensors are vectorized for the Jacobian computation. dv
The goal is to design the model pv such that translations, rotations and reflections of x do not change pv pv(x. (x. h). meaning that py is E(n) invariant with respect to Additionally, want py to also be invariant to permutations of (x, h). A rich family of distributions can be learned by choosing pz to be invariant, and fo to be equivariant.
The Normalizing Flow As consequence, multiple constraints need to be enforced on fo. From probabilistic modelling perspective 1) we require fo be invertible, and from symmetry perspective 2) require fe be cquivariant. One could utilize the EGNN (Satorras et al. 2021) which is E(n) cquivariant, and adapt the transformation so that it is also invertible. The difficulty is that when both these constraints are enforced naively, the space of Iearnable functions might be small. For this reason we opt for a method that only requires mild constraints to achieve invertibility; neural ordinary differential equations. These methods require functions to be Lipschitz (which most neural networks are in practice, since they operate on a restricted domain) and to be continuously differentiable (which most smooth activation functions are).
To this extent, we define Jo to be the differential equation, integrated over a conceptual time variable using differential predicted by the EGNN Φ. We redefine x, h as functions depending time, where x(t 0) = x and h(t = 0) = h in the data space. Then x(1) = Ze. h(1) = 2j, are the latent representations. This admits formulation of the flow as the solution to ODE defined as:
Zz 2/ = f(x, h) = ((00)b()) + L. b(x()).h(t)dd
(7)

---

The solution to this equation can be straightforwardly obtained by using the torchdiffeq package, which also supports backpropagation. The Jacobian term under the ODE formulation is log |det J Jo Tr (x((), h(tt)dtt as explained in Section 2 Equation The Trace of Jo has been approximated
with the Hutchinson's 's trace estimator.
The Dynamics The dynamics function & in Equation 7 models the first derivatives of x and with respect to time, over which we integrate. Specifically: derivative is modelled by the EGNN of L layers introduced in Section 2.1
- x h(t - (x((t), h()). This
di x
di h(1)
x h(t))=
x(t), (t)
where
x (t),
(GNN(X((), h()).
(8)
Notice that we directly consider the output h from the last layer L of the EGNN as the differential d de htt) of the node features, since the representation is invariant. In contrast, the differential of the node coordinates is computed as the difference between the EGNN output and intput ai xt)
d
x" (t). This choice is consistent with the nature of velocity-typp equivariance: although de x(t)
rotates exactly like x, it is unaffected by translations as desired.
The original EGNN from (Satorras et al. 202T) is unstable when utilized in an ODE because the coordinate update from Equation 5 would easily explode. Instead, we propose an extension in Equation 5 that normalizes the relative difference of two coordinates by their norm plus a constant C to ensure differentiability. In practice we set C' = 1 and found this to give stable results.
(xl x-xx) x²)
= xxx++ +
OR (m.y)
(9)
||x x/|1
Translation Invariance Recall that we want the distribution pv pv(V) to be translation invariant with respect to the overall location and orientation of positional coordinates x. For simplicity, let's assume only distribution px (x) over positions and an invertible function z = fxx). Translation invariance is defined as px (x +t) = px (x) for all t: a constant function. However, this cannot be a distribution since it cannot integrate to one. Instead, we have to restrict px to subspace.
To construct translation invariant px, we can restrict the data, flow So and prior pz translation invariant linear subspace, for instance by centering the nodes so that their center of gravity is zero Then the positions x M lie on the (M 1) x n-diiennional linear subspace defined by Ei= X; = O. However, from a modelling perspective it is easier to represent node positions as M sets of coordinates that are n-dimensional in the ambient space. In short, we desire the distribution to be defined on the subspace, but with the representation of the nodes in the ambient space.
To limit the flow to the subspace. in practice only the mean of the output of the dynamics network ø removed, following (Köhler et al. 2020). Expanding their analysis, we derive that the Jacobian determinant in the ambient space is indeed equal to the Jacobian determinant in the subspace under this condition. Intuîtively, the transformation fo does not change orthogonal to the subspace, and as result there is no volume change in that direction. For this reason the determinant can safely be computed in the ambient space, which conveniently allows the use of existing libraries without modification. Additionally, we can find the proper normalization constant for the base distribution.
For a more formal argument, let P be a dimensional subspace with orthonormal rows. Consider a collection of points
M-n matrix that projects points to the (M-1) n
R RM where x = 0 and z = fx). Define x Px, ž Pz, where x, z are considered to be vectorized and signifies that the variable is defined in the coordinates of the subspace. The Jacobian in the subspace J, is:
dā dz dx dx dx j, dx dz
de
PJ
(10)
To connect the determinant of y to Jy. let Q E R M- be the orthogonal extension of P using
orthonormal vectors qi,.... .a.......... so that QT = [ qi qi... q.. Then QJ,Q² / I
0
where
TP=PJ,PP and In x identity matrix. From this observe that det J det QJ,Q"= det Jr. which proves the claim. As a result of this argument, we are able to utilize existing methods to compute volume changes in the subspace without modification, as det. J, = det Jy, under the constraint that fe is an identity orthogonal to the subspace.
5

---

The base distribution Next we need to define base distribution pz- This base distribution can be divided in two parts: the positional part (zzz) and the feature part p(Zz), which we will choose to be independent so that p(Z,, 2x,) ((xx)) p Z..) The feature part is straightforward because the features are already invariant with respect to ((n)) symmetries, and only need to be permutation invariant. A common choice is standard Gaussian p(z,) = (z,,, (0, T). For the positional part recall that Zy lies on an (M- - 1)n subspace, and we need to specify the distribution over this space. Standard Gaussian distributions are reflection and rotation invariant since || ||zm ||2 for any rotation or reflection R. Further, observe that for our particular projection 2g = Pz, it is true that || 2 ||z since 2₂ lies in the subspace. More formally this can be seen using the orthogonal
% 2
extension Q of P as defined earlier and observing that: ||z||| = ||Qz,11 =
1|2,11
0
Therefore, a valid choice for rotation invariant base distribution on the subspace is given by:
p(2.) = (2
(2. 10, 1) = (2T)(M- 1)n/2 exp
(11)
which can be directly computed in the ambient space using |z= with the important property that the normalization constant uses the dimensionality of the subspace: (M - 1)n, so that the distribution is properly normalized.
Modelling discrete properties Normalizing flows model continuous distributions. However, the node features h may contain both ordinal (e.g. charge) and categorical (e.g. atom type) features. To train a normalizing flow on these, the values need to be lifted to a continuous space. Let h (hord, heas) ) be divided in ordinal and categorical features. For this we utilize variational dequantization (Ho et al. 2019) for the ordinal features and argmax flows (Hoogeboom al. 2021 for the categorical features. For the ordinal representation hord, interval noise u ~ aord( (
is used to lift to the continuous representation hord = hord + u. Similarly, heat is lifted using a distribution heat Qeat hhaat) where geat is the probabilistic inverse to an argmax function. Both
ford and Geat are parametrized using normal distributions where the mean and standard deviation are learned using an EGNN conditioned on the discrete representation. This formulation allows training on the continuous representation = (hord, Reat) as it lowerbounds an implied log- oog-liclihood of the discrete representation h using variational inference:
log PH (h) 2 E Ene-gord.eae (- (b) (o) log qord, (hi]h)
(12)
To sample the discrete h ~ PH. first sample the continuous h ~ PH via flow and then compute = (round(hord) argmax(hhcat) to obtain the discrete version. In short, instead training directly on the discrete properties h, the properties are lifted to the continuous variable h. The lifting method depends on whether a feature is categorical or ordinal. On this lifted continuous variable h the flow learns рH, which is guaranteed to be lowerbound via Equation 12 the discrete PH- To avoid clutter, in the remainder of this paper no distinction is made between h and h as one can easily transition between them using qord 4eat and the round, argmax functions.
Finally, the number of nodes M may differ depending on the data. In this case we extend the model using a simple one dimensional categorical distribution pM of M categories. This distribution PM is constructed by counting the number of molecules and dividing by the total. The likelihood of a set of nodes is then pv(x, h, .) =PV, (x, h [PM(MM where PVu, (x, h MM is modelled by the flow as before and the same dynamics can be shared for different sizes as the EGNN adapts to the number of nodes. In notation we sometimes omit the M conditioning for clarity. To generate sample, we first sample M ~ PM. then Zz, Z/,- ~ pz Ez. Z/, M) and finally transform to the node features and positions via the flow.
5 Experiments
5.1
DW4 and LJJ3
In this section we study two relatively simple sytems, DW-4 and LJ-13 presented in (Köhler et al.) 2020) where E(n) symmetries are present. These datasets have been synthetically generated by sampling from their respective energy functions using Markov Chain Monte Carlo (MCMC).

---

DW4: This system consists of only M-4 particles embedded in a 2-dimensional space which are governed by an energy function that defines a coupling effect between pairs of particles with multiple metastable states. More details are provided in Appendix A.
IJ-13: This is the second dataset used in (Köhler al., 2020) which is given by the Leonnard-Jones potential. It is an approximation of inter-molecular pair potentials that models repulsive and attractive interactions. captures essential physical principles and it has been widely studied model solid, fluid and gas states. The dataset consists of M=13 particles embedded in 3-dimensional state. More details are provided in Appendix A.I
Both energy functions (DW4 and LJ13) are equivariant to translations, rotations and reflections which makes them ideal to analyze the benefits of equivariant methods when E(n) symmetries are present on the data. We use the same MCMC generated dataset from (Köhler et al. 2020). For both datasets we use 1,000 validation samples, and 1,000 test samples. We sweep over different numbers of training samples (10?, 10 10", 109) and (10, 102, 109, 10") for DW4 and LJ13 respectively to analyze the performance in different data regimes.
We compare to the state-of-the art E(n) equivariant flows "Simple Dynam- ics" and "Kernel Dynamics" presented in (Köhler et al. 2020). We also compare to non-equivariant variants of our method, Graph Normalizing Flow (GNF), GNF with attention (GNF-att) and GNF with attention and data augmentation (GNF-att-aug), i.c. augmenting the data with rotations. Our E-NF method and its non won-equivariant variants (GNF, GNF-att, GNF-att-aug) consist of 3 layers each, 32 features per layer, and SiLU activation functions. All reported numbers have been averaged over runs. Further implementation details are provided in the Appendix
Implementation details:
A A.1
Table 1: Negative Log Lik Likelihood comparison on the test partition over different methods on DW4 and IJ13 datasets for different amount of training samples averaged over 3 runs.
DW4
LJ13
# Samples
10a
10
10
10
10
10
10
10
GNF GNF-att
11.93
11.31
10.38
7.95
43.56
42.84
37.17
36.49
11.65
11.13
9.34
7.83
43.32
36.22
33.84
32.65
GNF-att-aug
8.81
8.31
7.90
7.61
41.09
31.50
30.74
30.93
Simple dynamics
9.58
9.51
9.53
9.47
33.67
33.10
32.79
32.99
Kernel dynamics
8.74
8.67
8.42
8.26
35.03
31.49
31.17
31.25
E-NF
8.31
8.15
7.69
7.48
33.12
30.99
30.56
30.41
we report the cross-validated Negative Log Likelihood for the test partition. Our E-NF outperforms its non-equivariant variants (GNF, GNF-att and GNF-att-aug) and (Köhler et al. 2020) methods in all data regimes. It is interesting to see the significant increase in performance when including data augmentation (from GNt- to GNF-att-aug) in the non-equivariant models.
Results: In Table
5.2
QM9 Positional
We introduce QM9 Positional as a subset of QM9 that only considers positional information and does not encode node features. The aim of this experiment is to compare our method to those that only operate on positional data (Köhler et al. 2020) while providing a more challenging scenario than synthetically generated datasets. QM9 Po- sitional consists only of molccules with 19 atoms/nodes, where each node only has a 3-dimensional positional vec- tor associated. The likelihood of a molecule should be invariant to translations and rotations on a 3-dimensional space which makes equivariant models very suitable for this type of data. The dataset consists of 13,831 training samples, 2,501 for validation and 1,813 for test.
504
2:53-
0:2-
Figure 3: Normalized histogram of rel- ative distances between atoms for QM9 Positional and E-NF generated samples.
In this experiment, in addition to reporting the estimated Negative Log Likelihood, we designed metric to get an additional insight into the quality the generated samples. More specifically, we produce histogram of relative distances between all pairs of nodes within each molecule and
7

---

compute the Jensen Shann-nn divergence (Lin 1991 JSaiv (P Pdata) between the normalized histograms from the generated samples and from the training set. See Figure B for an example.
Implementation details: As in the previous experiment, we compare our E-NF to its non-equivariant variants GNF, GNF-att, GNF-att-auu and to the equivariant methods from (Köhler et al. 2020 Simple Dynamics and Kernel Dynamics. The dynamics of our E-NF, GNF, GNF-att and GNF-att-aug consist of 6 convolutional layers each, the number of features is The learning rate is set to 5 10 for all experiments except for the E-NF and Simple dynamics which was set to 2- 10-4 for stability reasons. All experiments have been trained for 160 epochs. The JS divergence values have been averaged over the last 10 epochs for all models.
64 and all activation layers are SiLUs.
# Metrics
NLL
JS(rel. dist)
Simple dynamics
73.0
.086
Kernel dynamics
38.6
.029
GNF
-00.9
.011
GNF-att
-26.6
.007
GNF-att-aug
-33.5
.006
E-NF (ours)
-70.2
.006
Figure 4: The table on the left presents the Negative Log Likelihood (NLL) log Pv (x) for the QM9
Positional dataset on the test data. The figure on the right shows the training curves for all methods.
Results: In the table from Figure 4 we report the cross validated Negative Log Likelihood - log pv (x) for the test data and the
divergence. Our E-NF outperforms all other algorithms in terms of NLL of the dataset. Additionally, the optimization curve with respect to the number of iterations converges much quicker for our E-NF than for the other methods as shown the right in Figure Regarding the JS divergence, the E-NF and GNF-att-aup achieve the best performance.
5.3
QM9 Molecules
(Ramakrishnan et al. 2014) is molecular dataset standarized in machine learning as chemical property prediction benchmark. consists small molecules (up to 29 atoms per molecule). Atoms contain positional coordinates embedded in 3D space, one-hot encoding vector that defines the type of molecule (H, C, N, O, F) and an integer value with the atom charge. Instead of predicting properties from molecules, We desire the likelihood of a molecule to be invariant to translations and rotations, therefore, our E(n) equivariant normalizing flow is very suitable for this task. To summarize, model distribution over 3D positions x, and atom properties h. These atom properties consist of the atom type (categorical) and atom charge (ordinal).
QM9
the QM9 dataset to learn a generative distribution
molecules.
We use the dataset partitions from (Anderson et al. 2019), 100K/18K/13K for training/validationnnees respectively. To train the method, the nodes (x, h) put into Equation as x(0), h(0) at time 0 and integrated to time 1. Using the continuous-timg change of variables formula and the base distribution, the (negative) log-likelihood of a molecule is computed - log pv pv (x, h, M). Since molecules differ in size, this term includes - log PM (M) which models the number of atoms as a simple 1D categorical distribution and is part of the generative model as described in Section B!
Implementation details: We compare our E-NF to the non-equivariant GNF-att and GNF-att-auu introduced in previous experiments. Notice in this experiment we do not compare to Köhler et al. 2020) since this dataset contains invariant feature data in addition to positional data. Each model composed of 6 layers, 128 features in the hidden layers and SilU activation functions. We use the same learning rates as in QM9 Positional. Note that the baselines can be seen as instances of permutation equivariant flows (Liu et al. 2019 Bilos and Günnemann 2020) but where the GNN architectures have been chosen to be as similar as possible to the architecture in our E-NFs.
Results (quantitative): Results are reported in Table 2 As in previous experiments, our E-NF significantly outperforms the non-equivariant models GNF and GNF-aug. terms of negative log-like ligllikel the E-NF performs much better than its non-equivariant counterparts. One factor that increases this difference the E-NFs ability to capture the very specific distributions of inter- atomic distances. Since the E-NF is able to better capture these sharp peaks in the distribution, the
8

---

Table 2: Neg. log- log-liiklihood - logpoogpp h, M), atom stability and mol stability for the QM9 dataset.
Metrics
NLL
Atom stability
Mol stable
GNF-attention
-28.2
72%
0.3%
GNN-attention-augmeain
-29.3
75%
0.5%
E-NF (ours)
-59.7
85%
4.9%
Data
99%
95.2%
Figure 5: Sampled molecules by our E-NF. The top row contains random samples, the bottom row also contains samples but selected to be stable. Edges are drawn depending inter-atomic distance.
negative log-likelihood becomes much lower. This effect is also seen when studying the number of stable atoms and molecules, which is very sensitive the inter-atomic distances. This stablity metric was computed over 10.000 samples from the model, for a detailed explanation of stability Appendix A..3 Observe that it might also be possible to utilize post-processing increase the molecule stability using prior knowledge. However. here we are mainly using this metric to see how many stable atoms and molecules the E-NF able to generate one shot, only by learning the molecule distribution. The E-NF on average produces 85% valid atoms, whereas the best baseline only produces 75% valid atoms. An even stronger improvement visible when comparing molecule stability: where the E-NF produces 4.9% stable molecules versus 0.5% by the best baseline. Interestingly, the percentage of stable molecules much lower than that of atoms. This is not unexpected: if even one atom in a large molecule is unstable, the entire molecule is considered to be unstable.
Finally, we evaluate the Validity, Uniqueness and Novelty as defined in (Smonnvsky and Komodakis 2018) for the generated molecules that are stable. For this purpose, we map the 3-dimensional representation of stable molecules to a graph structure and then to a SMILES notation using the rdkit toolkit. All our stable molecules are already defined as valid, therefore we only report the Novelty and Uniqueness since the Validity of those molecules that are already stable is 100%. The Novelty is defined as the ratio of stable generated molecules not present in the training set and Uniqueness is the ratio of unique stable generated molecules. Notice that different generated point clouds 3-dimensional space may lead the same molecule in the graph space or SMILES notation. Therefore, even if in the 3D space, all our generated samples were unique and novel, the underlying molecule that they represent doesn't have to be. Using our E-NF, we generated 10.000 examples to compute these metrics. We obtained 491 stable molecules (4.91 %), from these subset 415 molecules were novel (84.4%) and 352 were unique (71.8%). Further details and analyses are provided in Appendix A.4
Results (qualitative): In Figure 5 samples from our model shown. The top row contains random samples that have not been cherry-picker in any way. Note that although the molecule structure mostly accurate, sometimes small mistakes are visible: Some molecules are disconnected, and some atoms in the molecules do not have the required number of bonds. In the bottom row, random samples have been visualized but under the condition that the sample stable. Note that atoms might be double-bonded or triple-bonded, which is indicated in the visualization software. For example, in the molecule located in the bottom row 4th column, an oxygen atom double bonded with the carbon atom at the top of the molecule.

---

Limitations and Conclusions
Limitations In spite of the good results there are some limitations in our model that are worth mentioning and could be addressed in future work: 1) The ODE type of flow makes the training computationally expensive since the same forward operation has be done multiple times sequentially in order to solve the ODE equation. 2) The combination of the ODE with the EGNN exhibited some instabilities that we had to address (Equation 9 Despite the model being stable in most the experiments, we still noticed some rare peaks in the loss of the third experiment (QM9) that in one very rare case made diverge. 3) Different datasets may also contain edge data for which E-NFs should be extended. 4) Our likelihood estimation is invariant to reflections, therefore our model assigns the same likelihood to both a molecule and its mirrored version, which for chiral molecules may not be the same.
Societal Impact Since this work presents a generative model that can be applied to molecular data, it may advance the research the intersection of chemistry and deep learning. In the long term, more advanced molecule generation methods can benefit applications in drug research and material discovery. Those long term applications may have a positive impact in society, for example, creating new medications or designing new catalyst materials that permit cheaper production of green energy. Negative implications of our work may be possible when molecules materials are created for immoral purposes, such as the creation of toxins or illegal drugs.
Conclusions In this paper we have introduced E(n) Equivariant Normalizing Flows (E-NFs): generative model equivariant to Euclidean symmetries. E-NFs are continous- normalizing flows that utilize an EGNN with improved stability as parametrization. We have demonstrated that E-NFs considerably outperform existing normalizing flows in log-likelihood on DW4, LJ13, QM9 and also in the stability of generated molecules and atoms.
References
Anderson, B., Hy, T.-S., and Kondor, R. (2019). Cormorant: Covariant molecular neural networks. arXiv preprint arXiv: 0rXiv:1906.04015
Biloš, M. and Günnemann, S. (2020). Equivariant normalizing flows for point processes and sets. arXiv preprint arXiv:2010.03442.
Boyda, D., Kanwar, G., Racanière, S., Rezende, D. J., Albergo, M. S., Cranmer, K., Hackett, D. C., and Shanahan, P. E. (2020). Sampling using SU(N) gauge equivariant flows. CoRR. abs/2008.05456.
Bruna, J., Zaremba, W., Szlam, A., and LeCun, Y. (2013). Spectral networks and locally connected networks on graphs. arXiv preprint arXiv:1312.620:
Chen, C., Li, C., Chen, L., Wang, W., Pu, Y., and Carin, L. (2017). Continuous-time flows for deep generative models. arXiv preprint arXiv:1709.00117
Chen, C., Li, C., Chen, L., Wang, W., Pu, Y., and Carin, L. (2018a). Continuous-time flows for efficient inference and density estimation. In Dy, G. and Krause, A., editors, Proceedings of the 35th International Conference on Machine Learning, ICML.
Chen, T. Q., Rubanova, Y., Bettencourt, J., and Duvenaud, D. K. (2018b). Neural ordinary differential equations. In Advances in Neural Information Processing Systems, pages 6572-6583.
Cohen, T. and Welling, M. (2016). Group equivariant convolutional networks. In Balcan, M. and Weinberger, K. Q., editors, Proceedings the 33nd International Conference on Machine Learning, ICML.
Cohen, T. S. and Welling, M. (2017), Steerable cnns. In 5th International Conference on Learning Representations, ICLR.
Dey, N., Chen, A., and Ghafurian, S. (2021). Group equivariant generative adversarial networks, In International Conference Learning Representations.
10

---

Dieleman, S., Fauw, J. D., and Kavukeuoglu, K. (2016). Exploiting cyclic symmetry in convolu- tional neural networks. Balcan, M. and Weinberger, K. Q.. editors, Proceedings of the 33nd International Conference Machine Learning, ICML,, volume 48, pages 1889-1898. JMLR.org
Dinh, L., Krueger, D., and Bengio, Y. (2015). NICE: Non-linear independent components estimation. 3rd International Conference on Learning Representations. JCLR, Workshop Track Proceedings.
Dinh, L.,, Sohl-Dickstein, J., and Bengio, (2017). Density estimation using Real NVP. 5th International Conference on Learning Representations, ICLR.
Finlay, C., Jacobsen, J., Nurbekyan, L., and Oberman, A. M. (2020). How to train your neural ODE: the world of jacobian and kinetic regularization. In Proceedings of the 37th International Conference on Machine Learning, ICML.
Finzi, M., Stanton, S., Izmailov. P., and Wilson, A. G. (2020). Generalizing convolutional neural networks for equivariance to lie groups on arbitrary continuous data. In Proceedings of the 37th International Conference Machine Learning, ICMI..
Fuchs, F. B., Worrall, D. E., Fischer, V., and Welling, M. (2020). transforroso 3d roto- translation equivariant attention networks. CoRR, abs/2006. 10503.
Gebauer, N. W. A., Gastegger, M., and Schütt, K. (2019). Symmetry-adapter generation of 3d point sets for the targeted discovery of molecules. In Advances in Neural Information Processing Systems 32: Annual Conference on Neural Information Processing Systems 2019, NeurIPS.
BomezzBombareeli,, R., Duvenaud, D.., Hernández-LLobtto, J. M.,
J., Hirzel, T. D., Adams, R. P., and Aspuru-Guzik, A. (2016). Automatic chemical design using data-driven continuous representation of molecules. CoRR, abs/1610.02415..
Grathwohl, W.. Chen, R. T., Bettencourt, J., Sutskever, I., and Duvenaud, D. (2018). Ffjord: Free-form continuous dynamics for scalable reversible generative models. arXiv preprint arXiv: 1810.01367.
Ho, J., Chen, X., Srinivas, A., Duan, Y., and Abbeel, P. (2019). Flow++: Improving flow-based generative models with variational dequantization and architecture design. 36th International Conference on Machine Learning.
Hoogeboom, E., Nielsen, D., Jaini, P., Forré, P., and Welling, M. (2021). Argmax flows and multinomial diffuuson: Towards non-autoregressive language models. CoRR, abs/2102.05379.
Horie, M., Morita, N., Thara, Y., and Mitsume, N. (2020). Isometric transformation invariant and equivariant graph convolutional networks. CoRR, abs/2005.06316.
Hutchinson, M., Lan, C. Zaidi, S., Dupont, E., Teh, Y. W., and Kim, H. (2020). Lietransformer: Equivariant self-attention for lie groups. CoRR, abs/2012.10885.
Kingma, D. and Dhariwal, (2018). Glow: Generative flow with invertible 1x1 convolutions, In Advances in Neural Information Processing Systems, pages 10236-10245.
Kipf, T. N. and Welling, M. (2016). Semi-supervised classiiiction with graph convolutional networks. arXiv preprini arXiv: (609.029007
Klicpera, J.. GroB, J., and Günnemann, S. (2020). Directional message passing for molecular graphs. In 8th International Conference Learning Representations, ICLR 2020, Addis Ababa, Ethiopia, April 26-30, 2020. OpenReview.net.
Köhler, J., Klein, L., and Noé, F. (2019). Equivariant flows: sampling configurations for multi-body systems with symmetric energies. CoRR, abs/1910.00753.
Köhler, J., Klein, I.., and Noé, F. (2020). 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.
Li, S., Dong, C.,, Zhang, I.,, and Wang, I. (2019). Neural canonical transformation with symplectic flows. CoRR, abs/ 110.//90.00004..

---

Liao, R., Li, Y., Song, Y., Wang, S., Hamilton, W. L.. Duvenaud, D., Urtasun, R., and Zemel, R. S. (2019). Efficient graph generation with graph recurrent attention networks. In Advances Neural Information Processing Systems 32: Annual Conference on Neural Information Processing Systems 2019, NeurIPS.
Lin, J. (1991). Divergence measures based on the shannon entropy. IEEE Transactions on Information theory, 37(1):145-151
Liu, J., Kumar, A., Ba, J., Kiros, J., and Swersky, K. (2019). Graph normalizing flows. arXiv preprint 1905:1190
Nielsen, D., Jaini, P., Hoogeboom, E., Winther, O., and Welling, M. (2020). Survae flows: Surjections to bridge the gap between vaes and flows. CoRR, abs/2007.0273..
Noé, F., Olsson, S., and Wu, K. H. (2018). Boltzmann generators sampling equilibrium states of many-body systems with deep learning. CoRR, abs/1812.01729.
Nwankpa, C., Ijomah, W., Gachagan, A., and Marshall, S. (2018). Activation functions: Comparison of trends in practice and research for deep learning. arXiv preprint arXiv:1811.03378.
Ramakrishnan, R., Dral, P. O., Rupp, M., and Von Lilienfeld, O. A. (2014). Quantum chemistry structures and properties of 134 kilo molecules. Scientific data, 1(1):1-7.
Rezende, D. and Mohamed, (2015). Variational Inference with Normalizing Flows. In Proceedings of the 32nd International Conference on Machine Learning, volume 37 of Proceedings of Machine Learning Research, pages 1530-1538 1538. PMLR.
Rezende, J., Racanière, S., Higgins, I,, and Toth, P. (2019). Equivariant hamiltonian flows CoRR, abs/1909. 13739.
Rippel, O. and Adams, R. P. (2013). figh--dmmensionl probability estimation with deep density models. arXiv preprint arXiv:
Satorras, V. G., Hoogeboom, E., and Welling, M. (2021). E (n) equivariant graph neural networks. arXiv preprini arXiv:2102.0984:
Simonovsky, M. and Komodakis, N. (2018). Graphvae: Towards generation of small graphs using variational autoencoders. In International conference on artificial neural networks, pages 412 422. Springer.
Thomas, N., Smidt, T., Kearnes, S., Yang, L.., Li, L.,, Kohlhoff, K., and Riley, P. (2018). Tensor field networks: Rotation-and translation-cquivariaat neural networks for 3d point clouds. arXiv preprint arXiv:1802.08219
Xu, M., Luo, S., Bengio, Y., Peng, J., and Tang. J. (2021). Learning neural generative dynamics for molecular conformation generation. CoRR, abs/2102. 10240.
You, J., Ying, R., Ren, X., Hamilton, W. L.,, and Leskovec, J. (2018). Graphran: Generating realistic graphs with deep auto-regressive models. In Dy, G. and Krause, A., editors, Proceedings of the 35th International Conference on Machine Learning, ICML.
12


Reviews:
[{"review_id": "uyUJX8gTq92", "reviewer": "Reviewer_fcXk", "summary": "The paper introduces a new type of equivariant normalizing flow by using the CNF framework with a modified equivariant graph NN as a dynamics function. The equivariance discussed is wrt Euclidean symmetries, which are important for physical systems. They show that their flow beats prior equivariant baselines on synthetic data. Moreover, the model is capable of transforming spacial, categorical and ordinal features simultaneously. This allows the sampling of molecular configurations with positions, atom types, and charges. ", "questions": "", "limitations": "", "rating": 7, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "Merits of paper\n- The paper is clearly written, well organized and easy to understand. \n- They improve on previous work of Köhler et al. (2019, 2020) by utilizing equivariant graph NN as a dynamic function instead of their \"Kernel Dynamics\" function.\n- Incorporating equivariant graph NN into the (equivariant) continuous normalizing flow framework is non-trivial and novel.\n- The theory and related work are well presented. All previous work is acknowledged when necessary. \n- They compare to prior baseline of Köhler et al. (2020) and show that their model performs superior. \n- They introduce a new benchmark dataset, based on the QM9 dataset, and show that their model outperforms other non-equivariant as well as equivariant models.\n- The theory of equivariant flows is heavily based on the derivation from Köhler et al. (2019, 2020) and Rezende et al. (2019), but they contribute a proof for equivalence of the Jacobian of the used subspace and the whole space.  \n- They are able to sample new molecular configurations of different size with corresponding atom types and charges simultaneously. Meaning, they were (at least to some extent) able to learn the distribution of the subset of QM9 molecules with their E-NF. \n- This is a valuable contribution, sets a new baseline for molecular structure generation, and is very interesting for researchers in that area.\n\nCriticism\n- Although the training process and network architecture is outlined quite well, I could not find information about the Jacobian evaluation. This is of interest as it has proven to be a bottleneck in prior work, e.g. Köhler et al. (2020).  I would assume that some sort of approximation, such as the Hutchinson estimator, is used at least in the training process. I think this information should be included in the paper (or at least in the appendix).\n- The most interesting experiment is the generation of molecules. However, their stability measure is only based on distances. I think a better metric for stability measure would be the energy of the generated molecules. In fact, it might be that far less than 5% of the generated molecules are stable considering the energy instead of the distances, as already small deviations of the distances as well as the atom types can have huge impact on the energies. This should at least be discussed in the final paper.\n\nAll in all, a good paper, especially when the energy of the generated molecules is included. ", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "As mentioned above, a further limitation could be that the generated molecules are unstable considering their energies. So far, this is not mentioned in the paper.\nThe authors adequately addressed potential negative societal impact of their work.\n", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "8"}, "extra_text": {}}, {"review_id": "onOvuHvktx", "reviewer": "Reviewer_aifD", "summary": "The authors introduce a equivariant model for one-shot 3D molecule generation, building upon continuous time normalizing flows and equivariant graph networks. They demonstrated the utility of their approach based on several benchmark applications where they consistently outperform other flow-based models.", "questions": "", "limitations": "", "rating": 7, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "The proposed model present a sensible combination of existing methods, namely E(n) equivariant graph neural networks, variational dequantization and argmax flows. This enables the model to lean distributions over molecules considering positions as well as atom types. This is a clear improvement over previous related normalizing flows, which were focused on positional data.\n\nThe reported molecular generation experiments serve as a clear demonstration for the benefits of their method over other normalizing flows. However, one criticism is the strong focus on this particular class of methods  which is particularly noticeable in the final QM9 experiment. This makes it hard to evaluate how the proposed model behaves with respect to other generative approaches for 3D molecular structure, many of which have even been applied to the same dataset. Examples include iterative approaches that place atom after atom [1,2,3] as well as other one-shot approaches, either GAN- [4] or VAE-based [5]. A comparison to other methods is also made difficult by using the negative log-likelihood as primary evaluation metric. In this context, it would be worthwhile to analyze other statistics of the structures sampled in the QM9 experiment, e.g. how many of the stable molecules (which pass a valency check of choice) are training structures, test structures, or even novel structures and how many are unique? Additional measures of interest would be how the ring/atom/bond count of the generated structures compare to their distribution in QM9. One additional question would be how suitable the model is for practical applications, as the number of valid sampled molecules appear low at a first glance? Since QM9 provides relaxed structures of the molecules, one measure to consider in future experiments is the median RMSD between atomic positions of generated and relaxed structures, i.e.  how close are the sampled structures to actual equilibrium configurations (a drawback of this metric is that it is computationally intensive due to the need for quantum chemical computations).\n\nThe paper itself is well written providing a thorough description of the formal/theoretical background as well as the architecture.\n\nOn the whole, the proposed approach build upon and extends an interesting combination of machine learning approaches and represents a viable contribution with respect to one-shot generation. What remains unclear how it performs in comparison to other, non flow, generative approaches.\n\n1 Simm, Gregor, Robert Pinsler, and José Miguel Hernández-Lobato. \"Reinforcement learning for molecular design guided by quantum mechanics.\" International Conference on Machine Learning. PMLR, 2020.\n\n2 Simm, Gregor NC, et al. \"Symmetry-aware actor-critic for 3d molecular design.\" arXiv preprint arXiv:2011.12747 (2020).\n\n3 Gebauer, N. W., Gastegger, M., & Schütt, K. T. (2019). Symmetry-adapted generation of 3d point sets for the targeted discovery of molecules. arXiv preprint arXiv:1906.00957.\n\n4 Hoffmann, Moritz, and Frank Noé. \"Generating valid Euclidean distance matrices.\" arXiv preprint arXiv:1910.03131 (2019).\n\n5 Nesterov, Vitali, Mario Wieser, and Volker Roth. \"3DMolNet: a generative network for molecular structures.\" arXiv preprint arXiv:2010.06477 (2020).\n", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "Potential limitations and societal impact have been sufficiently addressed.\n", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "3"}, "extra_text": {}}, {"review_id": "JAArzTgpG8X", "reviewer": "Reviewer_iDHj", "summary": "The authors introduce a novel approach for generating data with intrinsic Euclidean symmetries. The proposed model uses recent advances in continuous normalizing flows, variational dequantization, and equivariant neural networks to construct a generative algorithm E(n)-equivariant along some modeled variables. The authors tackled many theoretical and technical problems to make the model stable and obtain desired equivariance in Euclidean space. The generative process looks as follows: first, the two groups of latent codes are sampled from the Gaussian distribution, which is rotation and reflection invariant; these samples are passed through the continuous normalizing flow guided by an Equivariant Graph Neural Networks. To force the translation equivariance, the authors propose to center the latent codes and work with (M-1) x n linear subspace. The target distribution is given by applying the variational dequantization for the sampled latent codes. The authors validate their approach on synthetic data and molecular conformations produced with DFT (density functional theory). ", "questions": "", "limitations": "", "rating": 7, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "**Originality:** Although all key parts are well-known, combining everything in one framework seems novel and original.\n\n**Quality:** The authors motivate this work well, describe all important details of their method, and propperly discuss limitations and impact of their work. The experiments on molecular data are quite simple and miss several important metrics to evaluate whether the proposed method can compete with state-of-the-art approaches in drug discovery.\n\n**Clarity:** The paper is well written and is easy to follow. The authors cover all key details of the proposed method and discuss all problems that were solved while developing the method.\n\n**Significance:** The paper provides simple experiments that do not evaluate the practical significance of the method in real drug discovery problems. The missing standard deviation in metrics also makes it difficult to validate significance of the results.\n\nUpdate: The authors clarified my concerns, I tend to keep the original score", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "The limitation of the work lies in the simplicity of the provided experiments on molecular data. First, important to say that model tries to solve two distinct drug discovery problems simultaneously — learn the distribution of molecular compounds [1] and model the distribution of conformational space for the molecular compounds [2]. Both of these problems have a wide variety of developed methods (including the neural-based) and a set of metrics on which the methods compete. I would recommend comparing the proposed model with the state-of-the-art model from these areas on the set of distribution learning and conformation modeling metrics.\n\n[1] Molecular Sets (MOSES): A Benchmarking Platform for Molecular Generation Models, Polykovskiy et al.\n[2] Learning Neural Generative Dynamics for Molecular Conformation Generation, Xu et al.\n", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "6"}, "extra_text": {}}]

Author rebuttal:

Meta-review: All reviewers agree on acceptance. This work is a highly significant contribution by introducing equivariant graph neural networks into the
(equivariant) continuous normalizing flow framework to obtain invertible equivariant functions. They show that their flow beats prior equivariant baselines and allows the sampling of molecular configurations with positions, atom types, and charges. A clear acceptance decision.

Pros:

- High significant contribution using equivariant graph neural networks and continuous normalizing flows to obtain invertible equivariant functions.
- The model can lean distributions over molecules considering positions as well as atom types, a clear improvement over previous related normalizing flows, which were focused on positional data.
- The reported molecular generation experiments serve as a clear demonstration for the benefits of their method over other normalizing flows.
- The theory of equivariant flows is heavily based on the derivation from Köhler et al. (2019, 2020) and Rezende et al. (2019), but they contribute a proof for equivalence of the Jacobian of the used subspace and the whole space. 
- Sets a new baseline for molecular structure generation.
- Very well-written paper. The authors motivate this work well, describe all important details of their method, and properly discuss limitations and impact of their work.

Limitations:

- Remains unclear how it performs in comparison to other, nonflow, generative approaches.
- Energies of the generated molecules are not reported.

The above list of pros and few existing limitations make this work highly relevant and worth an oral presentation.
Decision: Accept (Oral)