Equivariant Manifold Flows
Authors: Isay Katsman
Aaron Lou
Derek Lim
Qingxuan Jiang
Ser-Nam Lim
Christopher De Sa
Venue: NeurIPS 2021 Poster (2021)
Keywords: manifold
normalizing flow
equivariant
invariant

Abstract: Tractably modelling distributions over manifolds has long been an important goal in the natural sciences. Recent work has focused on developing general machine learning models to learn such distributions. However, for many applications these distributions must respect manifold symmetries—a trait which most previous models disregard. In this paper, we lay the theoretical foundations for learning symmetry-invariant distributions on arbitrary manifolds via equivariant manifold flows. We demonstrate the utility of our approach by learning quantum field theory-motivated invariant SU(n) densities and by correcting meteor impact dataset bias.

Equivariant Manifold Flows
Isay Katsman* *, Aaron Lou", Derek Lim", Qingxuan Jiang* Cornell University
{isk22, a1968, d1772, qj46)@cornell.cm edu
Ser-Nam Lim Facebook AI sernan@gmail.com
Christopher De Sa Cornell University www.eeccc.oneee.eeee..
Abstract
Tractably modelling distributions manifolds has long been important goal in the natural sciences. Recent work has focused on developing general machine learning models learn such distributions. However, for many applications these distributions must respect manifold symmetries-a trait which most previous models disregard. In this paper, we lay the theoretical foundations for learning -nnnaitr-invariant distributions on arbitrary manifolds via equivariant manifold flows. We demonstrate the utility of our approach by learning quantum field theory- motivated invariant SU(n) densities and by correcting meteor impact dataset bias.
Introduction
Learning probabilistic models for data has long been the focus of many problems in machine learning and statistics. Though much effort has gone into learning models over Euclidean space 6 20 MT less attention has been allocated to learning models over non- Eun-uucldeann spaces, despite the fact that many prob- lems require manifold structure. Density learning over non- Euclidean spaces has applications ranging from quantum field
theory in physics 44 to motion estimation in robotics [16] to
protein-structur prediction in computational biology [22
Continuous normalizing flows (CNFs)
powerful generative models for learning structure in complex data due to their tractability and theoretical guarantees. Recent work 29 30 has extended the framework of continuous normalizing flows the setting of density learning on Riemannian mani- folds. However, for many applications in the natural sciences, this construction is insufficient
cannot properly model nec- essary symmetries. For example, such symmetry requirements arise when sampling coupled particle systems in physical chem- (n]"Ilatiice latticce gauge theories
Figure 1: An example of density on SU(3) that is invariant conju- gation by SU(3). The æ-axis and y-axis are the angles 0, and fg for eigenvalues e and e of a ma- trix in SU).). The axis range
istry [26] or sampling for use in SU (n theoretical physics
More precisely, these symmetries are invariances with respect to action by an isometry subgroup of the underlying manifold.
n to
indicates equal contribution (S)))) denotes the special unitary group SU (n) = {X E Caxe X = det(X) = 1}.
35th Conference on Neural Information Processing Systems (NcurIPS 2021).

---

For example, consider the task of learning a density on the sphere that is invariant to rotation around an axis; this an example of learning an isometry subgroup invarian density. For a less trivial example, note that when learning a flow-based sampler for SU(n) in the context of lattice QFT ( the learned density must be invariant to conjugation by SU(n) (see Figure for a density on (U(3) that exhibits the requisite symmetry).
One might naturally attempt work with the quotient of the manifold by the relevant isometry subgroup in order to model the invariance. First, note that this structure is not always a manifold, and additional restrictions are needed on the action to ensure the quotient will have a manifold structure Assuming the quotient is in fact a manifold, one then asks whether an invariant density may be modelled by learning over this quotient with a general manifold density learning method such as NMODE [29] Though this seems plausible, it is problematic approach for several reasons:
1. First, it is often difficult to realize necessary constructs (charts, exponential maps, tangent spaces) on the quotient manifold (e.g. this is the case for RP", a quotient of S"
2. Second, even if the above constructs can be realized, the quotient manifold often has boundary, which precludes the use of a manifold CNF. To illustrate this point, consider the simple case of the sphere invariant rotation about an axis; the quotient manifold is closed interval, and a CNF would "flow out" on the boundary.
3. Third, even if the quotient is a manifold without boundary for which we have a clear characterization, it may have a discrete structure that induces artifacts in the learned distri- bution. This the case for Boyda et al. : the flow construction over the quotient induces abnormalities in the density.
Motivated by the above drawbacks, design manifold continuous normalizing flow the original manifold that maintains the requisite symmetry invariance. Since vanilla manifold CNFs do not maintain said symmetries, we instead construct equivariant manifold flows and show they induce the desired invariance. To construct these flows, we present the first general way of designing equivariant vector fields manifolds. A summary of our paper's contributions is follows:
We present a general framework and the requisite theory for learning equivariant manifold flows: in our setup, the flows can be learned over arbitrary Riemannian manifolds while explicitly incorporating symmetries inherent to the problem. Moreover, we prove that the equivariant flows we construct can universally approximate distributions on closed manifolds.
We demonstrate the efficacy of our approach by learning gauge invariant densities over (U(n) in the context of quantum field theory. In particular, when applied to the densities in Boyda et al. [3) we adhere more naturally to the target geometry and avoid the unnatural artifacts of the quotient construction.
We highlight the benefit of incorporating symmetries into manifold flow models by compar- ing directly against previous general manifold density learning approaches. We show that when a general manifold learning model is not aware of symmetries inherent to the problem, the learned density is of considerably worse quality and violates said symmetries. Prior to our work, there did not exist literature that demonstrated the benefits of incorporating isometry group symmetries for learning flows on manifolds, yet we achieve these benefits, and do so through novel equivariant vector field construction.
2 Related Work
Our work builds directly on pre-existing manifold normalizing flow models and enables them to leverage inherent symmetries through equivariance. In this section we cover important developments from the relevant fields: manifold normalizing flows and equivariant machine learning.
Thii specific isometry subgroup known as isotropy group at a point of sphere intersecting the axis. 'In particular, the isometry subgroup action needs to be smooth, free, and proper to ensure the quotient will be manifold by the Quotient Manifold Theorem 28

---

Normalizing Flows on Manifolds Normalizing flows on Euclidean space have long been touted as powerful generative models K T0 21. Similar to GANs 20 and VAEs [24, normalizing flows learn to map samples from tractable prior density to target density. However, unlike the aforementioned models, normalizing flows account for changes in volume, enabling exact evaluation of the output probability density. In rather concrete sense, this makes them theoretically principled. As such, they are ideal candidates for generalization beyond the Euclidean setting, where a careful, theoretically principled modelling approach is necessary.
Motivated by recent developments in geometric deep learning many methods have extended normalizing flows to Riemannian manifolds. Rezende et al. 38 introduced constructions specific to tori and spheres, while Bose et al. [ introduced constructions for hyperbolic space, Following this work, Falorsi and Forré [SI. Lou et al. (29]. Mathieu and Nickel [30] concurrently introduced a general construction by extending Neural ODEs
6 to the setting of Riemannian manifolds. Our
work takes inspiration from the methods of Lou et al. 29), Mathieu and Nickel [0 and generalizes
them further enable learning that takes into account symmetries of the target density.
Equivariant Machine Learning Motivated by the observation that many classic neural network architectures incorporate symmetry as an inductive bias, recent work has leveraged symmetries inherent in data through the concept of equivariance αن [8 27 377. Köhler et al. [26], in particular, used equivariant normalizing flows enable learning symmetric densities over Euclidean space. The authors note their approach is better suited to density learning in some physical chemistry settings (when compared to general purpose normalizing flows), since they take into account the symmetries of the problem.
Symmetries also appear naturally in the context of learning densities over manifolds. While in many cases symmetry can be a good inductive bias for learning for certain test tasks it is a strict requirement. For example, Boyda et al. 3 introduced equivariant flows on (U(n) for use in lattice gauge theories, where the modelled distribution must be conjugation invariant. However, beyond conjugation invariant learning on SU(n) (n2 B not much other work has been done for learning invariant distributions over manifolds. Our work bridges this gap by introducing the first general equivariant manifold normalizing flow model for arbitrary manifolds and symmetries.
3 Bacckrroudd
In this section, provide terse overview of necessary concepts for understanding our paper. In particular, we address fundamental notions from Riemannian geometry as well as the basic set-up of normalizing flows on manifolds. For more detailed introduction to Riemannian geometry, we refer the reader to textbooks such as Lee [2 and Kobyzev et al. 2
3.1 Riemannian Geometry
Riemannian manifold (M,h) is an n-dimensional manifold with smooth collection of inner products ( for every tangent space T.M. The Riemannian metric h induces a distance dy on the manifold.
diffeomorphism f . + M is a differentiable bijection with differentiable inverse. diffeomor- phism f: : :M M is called an isometry if (v) Dz)( h(u, v) for all tangent vectors u, U.E.E E TrM where D&f is the differential of f. Note that isometries preserve the manifold distance function. The collection of all isometries forms a group G, which we call the isometry group of the manifold M.
Riemannian metrics also allow for a natural analogue of gradients on R". For function f M - R. we define the Riemannian gradient Vzf to be the vector on T,M such that h(V,f, v) = Dzf(v) for T.M.
3.2 Normalizing Flows on Manifolds
Manifold Normalizing Flow Let hM, be a Riemannian manifold. A normalizing flow on M is a diffeomorphism fo M + M (parametrized by f) that transforms a prior density p to model
For example, asteroid impacts on the sphere can be modelled as being approximately invariant to rotation about the Earth's axis.
3

---

density Pfo. The model distribution can be computed via the Riemannian change of variables
Pfo (fo = ( det,, DzJø
Manifold Continuous Normalizing Flow A manifold continuous normalizing flow with base point a function : [0, 00) - M that satisfies the manifold ODE
(1 dt = (()(t),)
+(0) =
We define Fx,r: : : + = (x) to map any base point = E M to the value of the CNF starting at 2, evaluated time T. This function is known the (vector field) flow of X.
3.3 Equivariance and Invariance
Let G be an isometry subgroup of M. We notate the action of an element g E G on M by the map Lg -M M.
Equivariant and Invariant Functions
We say that a function f FM + A is equivariant if, for
all isometries Ia M -> .M and By :N N, we have o 9z By f. We say function f:M + N is invariant if f o 9x f.
Equivariant Vector Fields Let X M [0, oo) + TM, (m,, t) E T be time-dependent vector field on manifold M, with base point zo E M. X is a G-equivariant vector field if V(m, t) E M x [0, o0), X (L () (0mLg))(mm)) t).
Equivariant Flows A flow f M + M Geeuivvrrintt if it commutes with actions from i... we have Lg o = fo Lg.
Invariance of Density A density a on manifold M is G-invariant if, for all g G and M.
p(Tgzz) p(x), where Tg is the action of g on
4 Invariant Densities from Equivariant Flows
Our goal in this section is to describe a tractable way to learn density over a manifold that obeys a symmetry given by an isometry subgroup G. Since this cannot be done directly and it is not clear how a manifold continuous normalizing flow can be altered to preserve symmetry, we will derive the following implications to yield tractable solution:
1. G-invariant potential G-equivariant vector field (Theorem We show that given G-invariant potential function M R. the vector field V is G-equivariant.
2. G-equivariant vector field = G-equivariant flow (Theorem 2 We show that a G- equivariant vector field on M uniquely induces G-equivariant flow.
3. G-equivariant flow G-invariant density (Theorem 3
We show that given G- invariant prior p and a G-equivariant flow f. the flow density Py is G-invariant.
These are constructed in the same spirit as the theorems in Köhler et al. 26 (which also appeared in Papamakarios et al. (34 ) although note that results significantly
general. In addition to extending the domain to Riemannian manifolds, we consider arbitrary symmetry groups while Köhler al. 26 only considers the linear Lie group SO(n). As result, our proof techniques based on heavy geometric machinery instead of straightforward linear algebra techniques.
If we have a prior distribution the manifold that obeys the requisite invariance, then the above implications show that we can use a G-invariant potential to produce a flow that, tandem with the CNF framework, learns output density with the desired invariance. We claim that constructing G- G-invariant potential function on a manifold far simpler than directly parameterizing a G-invariant density G-equivariant flow. We shall give explicit examples of G-invariant potential constructions in Section 5.2 that induce a desired density invariance.
$ Herr, det,, is the determinant function with volume induced the Riemannian metric h.

---

Moreover, we show in Theorem that considering equivariant flows generated from invariant potential functions suffices to learn any smooth distribution over a closed manifold, as measured by Kullback-I eibler
eibber divergence.
We defer the proofs of all theorems to the appendix.
4.1 Equivariant Gradient of Potential Function
We start by showing how to construct Geequivariant vector fields from G-invariant potential functions.
To design an equivariant vector field X, it is sufficient to set the vector field dynamics of X as the gradient some G-invariant potential function + M -> This is formalized the following theorem.
Theorem 1. Let (M, h.) be Riemannian manifold and G be its group of isometries (or an isometry subgroup). t MM- R is smooth G-invariant fincction, then the following diagram commutes for any E G:
Lg
M
M
44
DL.
TM
TM
or -
2) Hence 4 a G-equivariant vector field. This condition is also tight in the sense that it only occurs if G is the isometry subgroup.
Hence, as long one can construct a G-invariant potential function, one can obtain the desired equivariant vector field. By this construction, parameterization of G-invariant potential functions yields parameterization of (some) G-equivariant vector fields.
4.2 Constructing Equivariant Manifold Flows from Equivariant Vector Fields
To construct equivariant manifold flows, we will use tools from the theory of manifold ODEs. In particular, there exists a natural correspondence between equivariant flows and equivariant vector fields. We formalize this in the following theorem:
Let (M,h) be a Riemannian manifold, and G be its isometry group (or one of its subgroups). Let X be any time-dependent vector field M, and Fx,r be the flow of X. Then X is a G-equivariant vector field if and only if Fx,r is a G-equivariant vector field flow.
Theorem 2.
Hence obtain equivariant flow from equivariant vector field, and vice versa.
4.3 Invariant Manifold Densities from Equivariant Flows
We now show that G-equivariant flows induce G-invariant densities. Note that we require the group G' to be isometry subgroup in order to control the density of ps, and the following theorem does not hold for general diffeomorphism subgroups.
Let (M,h) be a Riemannian manifold, and G be its isometry group (or one of its subgroups). If p is G-invariant density on M, and is G-equivariant diffeomorphism, then P is also G-invariant.
Theorem 3.
In the context of manifold normalizing flows, Theorem B Iimpies that if the prior density on M is G-invariant and the flow is G-equivariant, the resulting output density will be G-invariant. In the context of the overall set-up, this reduces the problem of constructing G-invariant density to the problem of constructing a G-invariant potential function.
4.4 Sufficiency of Flows Generated via Invariant Potentials
It is unclear whether equivariant flows induced by invariant potentials can learn arbitrary invariant distributions over manifolds. particular, it is reasonable to have some concerns about limited expressivity, since it is unclear whether any equivariant flow can be generated in this way. We alleviate
5

---

these concerns for our use cases by proving that equivariant flows obtained from invariant potential functions suffice to learn any smooth invariant distribution over closed manifold, as measured by Kullback-I eibler (KI.) divergence.
Theorem 4.
Let (M, ht) be closed Riemannian manifold. Let be smooth, non-vanishing distribution over M, which will act as our target distribution. Let pt be a distribution over said manifold parameterized by real time variable t. with po acting the initial distribution. Let | TT denote the Kullback-Leible divergence between distributions and n. If we choose g M R such that
D (pel
n (xx)
g(x) = log
(t(x)
and if Pt. evolves with as the distribution of flow according to 9. it follows that
① -- expla)) dæ = -- J(a - V|ll dæ
DKe(PelIT)=
implying convergence of Pt to π in KI. Moreover, the exact diffeomorphism that takes us from Po + as follows. Given some initial point I E M, let u(t) be the solution the initial value problem given by:
du(t) dt Vg(t),
u(0) =
The desired diffeomorphism maps x to lim u(t).
Hence if the target distribution т. the current distribution is Po. and g as defined above is the potential from which the flow controlling the evolution of Pt is obtained, then Pt converges to æ in KL. This means that considering flows generated by invariant potential functions is sufficient to learn any smooth invariant target distribution closed manifold (as measured by KL divergence).
5 Learning Invariant Densities with Equivariant Flows
In this section, we discuss implementation details of the methodology given in Section In particular, we describe the equivariant manifold flow model, provide two examples of invariant potential constructions on different manifolds, and discuss how training is performed depending on the target task.
5.1 Equivariant Manifold Flow Model
For our equivariant flow model, we first construct a G-invariant potential function 4 M + R (we show how construct these potentials Section 5.2 The equivariant flow model works by using automatic differentiation (35] on & to obtain V using this V for the vector field, and integrating in step-wise fashion over the manifold. Specifically, forward integration and thange-in-densee (divergence) computations utilize the Riemannian Continuous Normalizing Flows (30) framework. This flow model is used in tandem with a specific training procedure (described in Section 5.3 to obtain a G-invariant model density that approximates some target.
5.2 Constructing G-invariant Potential Functions
In this subsection, present two constructions of invariant potentials on manifolds. Note that symmetry of a manifold (i.e. action by an isometry subgroup) will leave part of the manifold free. The core idea of invariant potential construction is parameterize neural network the free portion of the manifold. While the two constructions we give below are certainly not exhaustive, they illustrate the versatility of method, which is applicable to general manifolds and symmetries.
5.2.1 Isotropy Invariance on S2
Consider the sphere S which is the Riemannian manifold (u E R |||||| =1] with the induced pullback metric. The isotropy group for a point v is defined as the subgroup of the isometry group which fixes v. i.e. the set of rotations around an axis that passes through u. practice, we let = (0, 0, 1), so the isotropy group is the group of rotations on the pyane An isotropy invariant density would invariant to such rotations, and hence would look like a borizontally striped density on the sphere (see Figure 4a).

---

Invariant Potential Parameterization We design an invariant potential by applying neural network the free parameter. the case of our specific isotropy group listed above, the free parameter is the z-coordinate. The invariant potential is simply a 2-input neural network with the spatial input being the z-coordinate and the time input being the time during integration. a result of this design, we see that the only variance in the learned distribution that uses this potential will be along the z-axis, as desired.
Prior Distributions For proper learning with normalizing flow, we need a prior distribution on the sphere that respects the isotropy invariance. There are many isotropy invariant potentials on the sphere. Natural choices include the uniform density (which is invariant to all rotations) and the For our experiments, we use the uniform density.
wrapped distribution with the center at o 33 40).
5.2.2 Conjugation Invariance on (U)
For many applications in physics (specifically gauge theory and lattice quantum field theory), one works with the Lie Group SÜ(n) - the group of unitary matrices with determinant 1. In particular, when modelling probability distributions on (n) for lattice QFT, the desired distribution must be invariant under conjugation by (U)n) [.. Conjugation is an isometry on SU (m) (see Appendix AA.5). so we can model probability distributions invariant under this action with our developed theory.
Invariant Potential Parameterization We want to construct a conjugation invariant potential function & (n) + R. Note that matrix conjugation preserves eigenvalues. Thus, for a function 4: : (n) + IR to be invariant to matrix conjugation, it has to act on the eigenvalues of x E SU(n) as multi-set.
We can parameterize such potential functions & by the DeepSet network from 45.. permutation invariant neural network that acts on the eigenvalues, so the mapping of z (U) (x)) (({A) (z)...... An (z))) for some set function 4 We append the integration time to the input of the standard neural network layers in the DeepSet network.
DeepSet is a
As result of this design, non-similar matrices, while all similar matrices will be assigned the same density value.
that the only variance in the learned distribution will be amongst
Prior Distributions For the prior distribution of the flow, we need distribution that respects the matrix conjugation invariance. We use the Haar measure on SUnn), which is the uniform density over this manifold that is symmetric under gauge symmetry B The volume element of the Haar measure is given for an & E (U)n) as Haar(z) = Iiss (^(()) - (x)))) We can sample from and compute the log probabilities with respect to this distribution efficiently with standard matrix computations
32
5.3 Training Paradigms for Equivariant Manifold Flows
There are two notable ways in which we can use the model described in Section 5.1
Namely, we can use to learn to sample from distribution for which we have density function, or we can use it to learn the density given a way to sample from the distribution. These training paradigms are useful in different contexts, as we will see in Section 6
Learning to sample given an exact density. In certain settings, we are given an exact density and the task is to learn tractable sampler for the distribution. For example Boyda et al. [ we are given conjugation-invarian (without knowledge of any normalizing constants).
densities on SU SU(m) for which we know the exact density function
contrast to procedures for normalizing flow training that use negative lig-liielihood based losses, we do not have access to samples from the target distribution. Instead, we train our models by sampling from the Haar distribution on SU(n), computing the KL divergence between the probabilities that our model assigns to these samples and the probabilities of the target distribution evaluated at these samples, and backpropagating from this KI. divergence loss. When this loss is minimized, we can sample from the target distribution by sampling the prior, then forwarding the prior samples through our model. In the context of Boyda et al. (3) such flow- sampler is important for modelling gauge theories.
7

---

Ours
Boyda et al.
Target
Ours
Boyda
10a
10
10a
10-2
10
10-2
10
10
10" 10"
(a) (U()) learned densities from (I.eft) our model and (Right) the Boyda et al. 3 model. The target densities in orange. while the model densities are blue. The r-axis is an eigenvalue e a matrix in SU (2) (note the second eigenvalue determined as "). Our model has much better he- havior low-density regions (Boyda
(b) (U)(3) learned densities from (Middle) our model and (Right) the Boyda et al. [3] model for different target densities (Left). The r-axis and g-axis the angles O and O₂ eigenvalues eh and 12 of matrix in SU (3) (note the third eigenvalue is determined as e -1))). and the probabilities correspond colors logarith- mie scale. Our model better captures the geometry of the target densities and does not exhibit the discrete artifacts of the Boyda et al. R model.
[ fails to eliminate mass around iT) and more smoothly captures the targets in high-density regions.
Figure 2: Comparison learned densities on (a) SU()) and (b) SU All densities are normalized to have maximum value 1.
Learning the density given a sampler. In other settings, we are given way to sample from a target distribution and want to learn the precise density for downstream tasks. For this setting, we sample the target distribution, use our flow to map it to tractable prior, and use negative bog-likelihood-base loss. The flow will eventually learn to assign higher probabilities sampled regions, and in doing so, will learn to approximate the target density.
Experiments
In this section, we utilize instantiations of equivariant manifold flows to learn densities over various manifolds of interest that invariant certain symmetries. First, construct flows SU that
are invariant to conjugation by SU(n): these are useful for lattice quantum field theory [ In this setting, our model outperforms the construction of Boyda al. 3
As a second application, we model asteroid impacts on Earth by constructing flow models on S2 that invariant to the isotropy group that fixes the north pole. Our approach is able to overcome dataset bias, as only land impacts are reported in the dataset.
Finally, to demonstrate the need for enforcing equivariance of flow models, we directly compare our flow construction with general purpose flow while learning density with an inherent symmetry. The densities we decided to use for this purpose are sphere densities that are invariant to action by isotropy group. Our model able to learn these densities much better than previous manifold ODE models that do not enforce equivariance of flows 20 thus showing the ability of our model to leverage the desired symmetries. In fact, even on simple isotropy-invariant densities, our model succeeds while the free model without equivariance fails.

---

6.1 (n) Gauge Equivariant Neural Network Flows
Learning (n) gauge equivariant neural network flows is important for obtaining good flow-based samplers of densities on (U(n) useful for lattice quantum field theory [ We compare model
For the sake of staying true to the application area, we follow the framework of Boyda et al. K
for (n) gauge equivariant flows (Section 5.2.2) with that of Boyda et K
in learning densities
on SU (n) that are invariant to conjugation by (U)... particular, our goal is to learn flow to model a target distribution that may efficiently sample from it.
As mentioned above in Section 5.3 this setting follows the first paradigm in which we are given exact density functions and learn how sample.
For the actual architecture of equivariant manifold flows, parameterize potentials as DeepSet networks on eigenvalues as detailed in Section 5.2.2 The prior distribution for our model is also the Haar (uniform) distribution SU()). Further training details given in Appendix C.1 C.1
6.1.1
(2)
Figure 2a displays learned densities for our model and the model of Boyda et al. S in the case of three particular densities on (2) described in Appendix C.2.1 While both models match the target distributions well in high-density regions, we find that our model exhibits considerable improvement in lower-density regions, where the tails of our learned distribution decay faster. By contrast, the model of Boyda et al. K seems to be unable to reduce mass near = possible consequence of their construction. Even in high-density regions, our model appears to vary smoothly, with fewer unnecessary bumps and curves when compared to the densities of the model in Boyda et al. B
6.1.2
(U)
Figure 2b displays learned densities for our model and the model of Boyda et al. [3] in the case of three particular densities on (U) described in Appendix C.2.2 In this case, we see that our models fit the target densities more accurately and better respect the geometry of the target distribution. Indeed, while the learned densities of Boyda et al. [3] are often sharp and have pointed corners, our models learn densities that vary smoothly and curve in ways that are representative of the target distributions.
Figure 3: Our modelled distribution of meteor impacts in Meteorite Landings 31 The true impacts marked in blue and isotropy invariant density is shown in the background. Note that a regular manifold normalizing flow would instead model impacts only on land as the dataset does not include any ocean impacts.
6.2 Asteroid Impact Dataset Bias Correction
We also showcase our model's ability to correct for dataset bias. particular, we consider the test case of modelling asteroid impacts on Earth. Towards this end, many preexisting works have compiled locations of previous asteroid impacts [14 AT but modelling these datasets challenging since they are inherently biased. In particular, all recorded impacts found on land. However,

---

ocean impacts are also dangerous 2 and should be properly modelled. To correct for this bias, we note that the distribution of asteroid impacts should be invariant with respect to the rotation of the Earth. We apply our isotropy invariant S flow (described in Section 5.2.1 to model the asteroid impact locations given by the dataset Meteorite Landings 31 Training happens in the setting of the second paradigm described in Section 5.3] since we can easily sample the target distribution and aim to learn the density. visualize our results in Figure 5
6.3 Modelling Invariance Matters
We also show that our equivariant condition on the manifold flow matters for efficient and accurate training when the target distribution is invariant. In particular, we again consider the sphere under the action of the isotropy group. We try to learn the isotropy invariant density given in Figure 4a and compare the results of our equivariant flow against those of predefined manifold flow that does not explicitly model the symmetry [29]. While other manifold flow models have been proposed for the sphere [8), NMODE outperforms them (2) we use NMODE as strong baseline. We train for 100 epochs with a learning rate of 0.001 and a batch size of 200; our results are shown in Figure
(a) Ground Truth
(b) Isotropy Equivariant Flow
(c) NMODE
Figure 4: We compare the equivariant manifold flow and a regular manifold flow (implemented with NMODE (29] on an invariant dataset. Note that our model is able to accurately capture the ground truth data distribution while NMODE struggles.
Despite our equivariant flow having fewer parameters (as both flows have the same width and the equivariant flow has input dimension of 1), model is able capture the distribution much better than NMODE [2)] This is due to the inductive bias of our equivariant model which explicitly leverages the underlying symmetry.
7 Conclusion
In this work, we introduce equivariant manifold flows in fully general context and provide the necessary theory ensure our construction principled. We also demonstrate the efficacy of our approach in the context of learning conjugation invariant densities over SU (2) and SU))). which is an important task for sampling SU(n) lattice gauge theories in quantum field theory. In particular, we show that our method can more naturally adhere to the geometry of the target densities when compared prior work while being more generally applicable. We also present an application to modelling asteroid impacts and demonstrate the necessity of modelling existing invariances by comparing against regular manifold flow.
Further considerations. While our theory and implementations have utility in very general settings, there still some limitations that could be addressed in future work. Further research may focus finding other ways to generate equivariant manifold flows that do not rely on the construction of an invariant potential, and perhaps additionally on showing that such methods are sufficiently expressive to learn over open manifolds. Our models also require a fair bit of tuning to achieve results as strong we demonstrate. Finally, note that our theory and learning algorithm are too abstract for to be sure of the future societal impacts. Still, we advance the field of deep generative models, which is known to have potential for negative impacts through malicious generation of fake images and text. Nevertheless, we do not expect this work to have negative effects in this area, as our applications are not this domain.
This dataset released by NASA without specified license.
10

---

Acknowledgements
We would like thank Facebook Al for funding equipment that made this work possible. addition, we thank the National Science Foundation for awarding Prof. Christopher de Sa grant that helps fund this research effort (NSF IIS (00- 102) and for supporting Aaron Lou with graduate student fellowship. We would also like to acknowledge Jonas Köhler and Denis Boyda for their useful insights.
References
[1] Niels Henrik Abel. Mémoire sur les équations algébriques, où on demontre l'impossibilité de la résolution de l'équation générale du cinquième dégré. 1824.
[2] Joey Bose, Ariella Smofsky, Renjie Liao, Prakash Panangaden, and Will Hamilton. Latent variable modelling with hyperbolic normalizing flows. In Proceedings of the 37th International Conference on Machine Learning, pages 1045-1055, 2020.
[3] Denis Boyda, Gurtej Kanwar, Sébastien Racanière, Danilo Jimenez Rezende, Michael S Albergo, Kyle Cranmer, Daniel C Hackett, and Phiala E Shanahan. Sampling using su(n) gauge equivariant flows. arXiv preprint wiXiv:2008.05455 2020.
[4] Michael M. Bronstein, Joan Bruna, Taco Cohen, and Petar Velivckovié. Geometric deep learning: Grids, groups, graphs, geodesics, and gauges. arXiv preprint arXiv:2104.13478 2021.
[5] Daniel Bump. Lie groups. Springer, 2004.
[6] Ricky T. Q. Chen, Yulia Rubanova, Jesse Bettencourt, and David Duvenaud. Neural ordinary differential equations. In Advances in Neural Information Processing Systems, volume 31,
pages 6571-6583 6583, 2018.
[7] Taco Cohen and Max Welling, Group equivariant convolutional networks. In Proceedings of The 33rd International Conference on Machine Learning, pages 2990-2999, 2016.
[8] Taco Cohen, Maurice Weiler, Berkay Kicanaoglu, and Max Welling, Gauge equivariant convolutional networks and the icosahedral CNN. In Proceedings of the 36th International Conference on Machine Learning, pages 1321-13300 2019.
[9] Taco S. Cohen, Mario Geiger, Jonas Köhler, and Max Welling, Spherical CNNs. In International Conference on Learning Representations, 2018.
[10] Laurent Dinh, Jascha Sohl-Dickstein, and Samy Bengio. Density estimation using real nvp. In International Conference on Learning Representations, 2017.
[11] Manfredo Perdigão do Carmo. Riemannian geometry / Manfredo do Carmo translated by Francis Flaherty. Mathematics. Theory and applications. Birkhäuser, Boston, 1992. ISBN 08 0817634908.
[12] Harold Donnelly. Eigenfunctions of the laplacian on compact riemannian manifolds. Asian Journal of Mathematics, 10:1 2006.
[13] Conor Durkan, Artur Bekasov, Iain Murray, and George Papamakarios. Neural spline flows, In Advances in Neural Information Processing Systems, volume 32, 2019.
[14] Earth Impact Database. Earth impact database, 2011. Retrieved from http::/paase: net/ Ear EarthhmpactDatabaass
[15] Luca Falorsi and Patrick Forré. Neural ordinary differential equations manifolds. arXiv preprint arXiv:2006.0666, 2020.
[16] Wendelin Feiten, Muriel Lang, and Sandra Hirche. Rigid motion estimation using mixtures of projected gaussians. Proceedings of the 16th International Conference on Information Fusion, pages 1465-1472, 2013.

---

[17] MJ Field. Equivariant dynamical systems. Transactions of the American Mathematical Society, 259(1):185-205, 1980.
[18] Marc Finzi, Samuel Stanton, Pavel Izmailov, and Andrew Gordon Wilson Generalizing convolutional neural networks for equivariance to lie groups arbitrary continuous data. In International Conference on Machine Learning, pages 3165-3176. PMLR, 2020.
[19] Jean Gallier and Jocelyn Quaintance. Differential Geometry and Lie Groups: Computational Perspective, volume 12. Springer, 2020.
[20] Ian J Goodfellow, Jean Pouget-Abadie Mehdi Mirza, Bing Xu, David Warde-Farley Sherjil Ozair, Aaron Courville, and Yoshua Bengio. Generative adversarial nets. Proceedings of the 27th International Conference on Neural Information Processing Systems Volume 2, page 2672-2680, 2014.
[21] Will Grathwohl, Ricky TQ Chen, Jesse Bettencourt, and David Duvenaud. Scalable reversible generative models with free-form continuous dynamics. In International Conference on Learn- ing Representations, 2019.
[22] Thomas Hamelryck, John T Kent, and Anders Krogh. Sampling realistic protein conformations using local structural bias. PLoS Computational Biology, 2(9), 2006.
[23] Diederik P Kingma and Jimmy Ba. Adam: A method for stochastic optimization. In Interna- tional Conference on Learning Representations, 2015.
[24] Diederik P Kingma and Max Welling. Auto-encoding variational bayes. In International Conference on Learning Representations, 2014.
[25] Ivan Kobyzev, Simon Prince, and Marcus Brubaker. Normalizing flows: An introduction and review of current methods. IEEE Transactions on Pattern Analysis and Machine Intelligence, 2020.
[26] Jonas Köhler, Leon Klein, and Frank Noe. Equivariant flows: Exact likelihood generative learning for symmetric densities. In Proceedings of the 37th International Conference on Machine Learning, pages 3361- 2020.
[27] Risi Kondor and Shubhendu Trivedi. On the generalization of equivariance and convolution in neural networks to the action of compact groups. In Proceedings of the 35th International Conference on Machine Learning, pages 2747-2755, 2018
[28] John M Lee. Introduction to Smooth Manifolds. Graduate Texts in Mathematics. Springer New York, 2013.
[29] Aaron Lou, Derek Lim, Isay Katsman, Leo Huang, Qingxuan Jiang, Ser-Nam Lim, and Christopher De Sa. Neural manifold ordinary differential equations. In Advances in Neural Information Processing Systems, volume 33, pages 17548-17558, 2020.
[30] Emile Mathieu and Maximilian Nickel. Riemannian continuous normalizing flows. In Advances in Neural Information Processing Systems, volume 33, pages 2503-2515, 2020.
[31] Meteorite Landings. Meteorite landings dataset, March 2017. Retrieved from htttps:/ddatt porida/nasa/mtteorite---- landings
[32] Francesco Mezzadri. How to generate random matrices from the classical compact groups. Notices of the American Mathematical Society, 54:592-604, 2007.
[33] Yoshihiro Nagano, Shoichiro Yamaguchi, Yasuhiro Fujita, and Masanori Koyama. A wrapped normal distribution on hyperbolic space for gradient-based learning. In Proceedings of the 36th International Conference on Machine Learning, pages 4693-4702, 2019.
[34] George Papamakarios, Eric Nalisnick, Danilo Jimenez Rezende, Shakir Mohamed, and Balaji Lakshminarayanan. Normalizing flows for probabilistic modeling and inference. Journal of Machine Learning Research, 22(57): 2021.
12

---

[35] Adam Paszke, Sam Gross, Soumith Chintala, Gregory Chanan, Edward Yang, Zach Vito,o Zeming Lin, Alban Desmaison, Luca Antiga, and Adam Lerer. Automatic differentiation in pytorch. In Neural Information Processing System Autodiff Workshop, 2017.
[36] Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, et al. Pytorch: An imperative style, high-performance deep learning library. Advances in neural information processing systems, 32:8026-800 8037, 2019.
[37] Danilo Jimenez Rezende, Sébastien Racanière, Irina Higgins, and Peter Toth. Equivariant hamiltonian flows. arXiv preprint arXiv: 13739, 2019.
[38] Danilo Jimenez Rezende, George Papamakarios, Sebastien Racaniere, Michael Albergo, Gurtej Kanwar, Phiala Shanahan, and Kyle Cranmer. Normalizing flows on tori and spheres. In Proceedings of the 37th International Conference on Machine Learning, pages 5083-8092, 2020.
[39] Hannes Risken and Till Frank. The Fokker-Planck Equation: Methods of Solution and Applica- tions. Springer Series Synergetics. Springer Berlin Heidelberg, 1996.
[40] Ondrej Skopek. Octavian-Eugen Ganea, and Gary Bécigneul. Mixed-curvature variational autoencoders. In International Conference Learning Representations, 2020.
[41] Wei Wang, Zheng Dang, Yinlin Hu, Pascal Fua, and Mathieu Salzmann. Backpropagation- friendly eigendecomposition. In Advances in Neural Information Processing Systems, vol- ume 32, 2019.
[42] Steven N Ward and Erik Asphaug. Asteroid impact tsunami of 2880 March 16. Geophysical Journal International, F53(3):F--F10, 2003.
[43] Arthur G Wasserman. Equivariant differential topology. Topology, 8(2):127-150, 1969.
[44] Peter Wirnsberger, Andrew J Ballard, George Papamakarios, Stuart Abercrombie, Sébastien Racanière, Alexander Pritzel, Danilo Jimenez Rezende, and Charles Blundell. Targeted free energy estimation via learned mappings. The Journal of Chemical Physics, 144112) 2020.
[45] Manzil Zaheer, Satwik Kottur, Siamak Ravanbakhsh, Barnabas Poczos, Russ R Salakhutdinov, and Alexander J Smola. Deep sets. Advances in Neural Information Processing Systems, volume 30, pages 3391-3401, 2017.
13


Reviews:
[{"review_id": "lEc5tiBD_r9", "reviewer": "Reviewer_e1R7", "summary": "The paper develops a new method for learning densities over SU(n), and evaluates the performance for learning some idealized densities over SU(2) and SU(3), showing improvement over prior work of Boyda et al.\n", "questions": "", "limitations": "", "rating": 5, "confidence": 1, "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": "Unfortunately the methods of this paper lie far outside my area of expertise so there is very little of substance I can offer in review.  However my guess is that, as currently written, it will be accessible and hence of interest only to a tiny fraction of the NeurIPS community.  This sentence for example seems to convey the main significance of the paper:\n\"Learning SU(n) gauge equivariant neural network flows is important for obtaining good flow-based samplers of densities on SU(n) useful for lattice quantum field theory [2].\" \nI have no idea what this means, and my guess is that many others will be in the same boat.  If the authors wish to reach more than just the theoretical physics/quantum field theory community, I would suggest that they put more work into writing this in a way that others can appreciate and learn something from that they can possibly apply to their own work.\n", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "Yes", "review_text": "", "extra_scores": {"ethical_concerns": "none", "needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "1 hour"}, "extra_text": {}}, {"review_id": "XIQJlE9a2S7", "reviewer": "Reviewer_N2xA", "summary": "This paper proposes an extension Continuous Manifold Normalizing Flows to Equivariant Manifolds where the the equivariance is with respect to an isometry group of the manifold. To this end, the paper lifts existing tools to a more manifold specific setting and empirically demonstrates the utility of the proposed approach on sampling from SU(N).  Overall, the paper is relatively well written and easy to follow. I found the main ideas, such as Theorems 1-3, interesting---albeit not extremely novel---and correct. I believe the empirical results, while a bit on the simpler end, are illuminating and see no reason to doubt the correctness of the numbers.", "questions": "", "limitations": "", "rating": 6, "confidence": 5, "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": "First I'd like to highlight that the paper is quite easy to read and I had little trouble following the main results including the proofs in the Appendix. The main highlight of this paper for me were the two invariant potentials introduced which I believe are more of a central contribution than the developed theory. \n\nThis brings me the central complaint of the paper which is limited novelty. To be perfectly honest Theorems 1-3 are not sufficiently novel. In fact, I would not even elevate them to theorems but rather Lemmas as done in [1]. In fact [2] proves much of the results in Theorems 1-3 for group actions on Euclidean spaces. Lifting these to more general manifolds requires a little bit of technical nuance but I don't believe it is that difficult (e.g. instead of gradient, we use Riemannian gradient). Theorem 4 on the other seems to be more novel and interesting. I think it pokes at a larger but not fully explored issue on the universality of learning equivariant flows on manifolds. However, I believe this should be investigated further in the context of the proposed invariant potentials. A more powerful result would be to show that for a specific class of invariant potentials we can learn any equivariant density for some isometry group. Right if my understanding is correct the current result highlights that there exists an invariant potential where the KL is minimized and as a result this is a good learning target. \n\nFinally, the experiments are also a bit light. It would've been interesting to propose new manifolds where existing approaches would definitely not work. Also, I would argue that there are missing baselines for the S^2 example. Finally, a lot of the results in the appendix are quite interesting and would've made for a stronger experimental section.\n\nOverall, I believe this paper in its current narrative leans too heavily on borrowed theory which is extended in a natural but straightforward manner. Consequently, the main results that I found interesting such as the two proposed invariant potential functions and Theorem 4 are not given the full attention that they deserve. Consequently, I'm currently leaning towards a weak reject but I am willing to change my score if a good argument can be made against the points above and below.\n\n\nDetailed Feedback:\n------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------\nline 61 - Kohler et. al formulation for equivariant vector field is fairly general.\nSection 3.2, this is technically correct but it would be more fitting and match the title if you decomposed the change of variable and explicitly show the role of the Riemannian metric\n\nTheorem 1. Is a strict generalization from Kohler[2] & Papamakarios[1] review paper\nTheorem 2-3 are direct consequences\nR_g is a representation of the group. Is it a linear represenation? You must be more precise here.\n- line 196 G-equivariant CNF, I can construct discrete NF's on manifold without this condition.\n- you could also consider the prior density introduced in Kohler et. al which is on M-1 dimensional subspace.\n- 5.2.2. you could also use uniform density here no?\n- 359, Can you elaborate on what exactly you had to tune? This is an interesting point for the general community.\n\n\nPoint 2 line 45: We can always restrict the flow to a pre-defined support. This can be done by a projection operation. For example this is the underlying principle of a Restricted/Wrapped Gaussian. I don't believe this is an insurmountable limitation.\n\nPoint 3. Line 51. Can you expand more on this unnaturalness? Give a more precise characterization and explicate why exactly this is undesirable? It is possible to design flows over discrete structures in the conventional non-Manifold sense. Can we not leverage a similar design principle here? Please argue against this point.\n\nContribution 3 line 71. This is not really a new contribution but moreso the benefit of invoking symmetries when doing generative models. Other papers have also remarked the same observations. This paper also acknowledges this fact line 103-104\n\n\nline 207. I disagree. It does not strictly reduce the problem as one can construct G-invariant densities but not going through a G-invariant potential function. This is a clear mistake. I believe what the authors intended to say that in this particular context it is sufficient to use G-invariant potential function. But indeed this is not necessary, for example I can construct an equivariant map that is not the gradient of an invariant potential function and still arrive at a G-invariant density. \n\nTheorem 4. I believe I don't fully agree with the interpretation here. I agree that there may exist a G-invariant potential for every target distribution on a closed manifold. But you must show that the construction of one is universal. This point is missing.\n\n\nFor 6.2 The obvious missing baseline is NF's on Spheres and Tori in addition to the Manifold flow.\n\nMinor Presentation issues:\nLine 141, you haven't defined what a symmetry is.\nLine 146: This is obvious, and here $f$ is not distinguished as a flow. \nOverall 3.3 is a bit sloppy in notation and could be cleaned up a little bit.\nX in line 174 is a vector field, earlier in the background it was a manifold.\n\nTypo in line 506 in proof. I believe it should be D_u not Du R_g(x), ... in middle equality.\n\n5.1 -> 5.2 This is confusing, first you assume that you're given an invariant potential function then immediately you design two new ones? From a narrative perspective this is not very clear.\n\n\nQuestions:\n1. The prior for S^2 is only on the z-axis right ---i.e. its a 1-d uniform density?\n2. 5.3 Isn't the conventional duality of Normalizing Flows. We either need to do density estimation or we need to generate high quality samples. I don't see how this is particularly unique to Manifold Flows. \n\n------------------ Post-Rebuttal Updates -----------------------------------------------------------------------------------------------------------------------------------------\nAfter fruitful discussion with the other reviewers and the detailed author comments to my questions, I have raised my score from 5->6. This is conditioned on the authors making significant changes to the writing as mentioned in my original review. \n\n1.) The paper must do a better job of putting this work in context with [1] and [2] and clearly stating these are extensions of these ideas to the Riemannian case and not that they are proposed here for the first time. This is currently ambiguous in the original draft.\n\n2.) Theorem 4 at submission time has many interesting nuances that were uncovered in the discussion. I suggest this theorem be re-written along with the proof which has a few technical issues that need to be addressed (e.g. smoothness in the limit) as pointed in the discussion.\n\nI believe this paper has the potential to be very interesting to the community but the key detractor at the moment is the writing and more generally the unclear narrative in context with prior work. Once these issues are fixed---and if accepted---this paper would be interesting to the NeurIPS community.\n\n[1] Papamakarios, G., Nalisnick, E., Rezende, D.J., Mohamed, S. and Lakshminarayanan, B., 2019. Normalizing flows for probabilistic modeling and inference. arXiv preprint arXiv:1912.02762.\n\n[2] Köhler, J., Klein, L. and Noé, F., 2020, November. Equivariant flows: exact likelihood generative learning for symmetric densities. In International Conference on Machine Learning (pp. 5361-5370). PMLR.\n", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "The authors have adequately addressed their limitations and discussed negative societal impact.", "review_text": "", "extra_scores": {"ethical_concerns": "N/A", "needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "8 hours"}, "extra_text": {}}, {"review_id": "DwLZ_CpHsDF", "reviewer": "Reviewer_RD2g", "summary": "The paper investigates the construction of equivariant flows on Riemanian manifolds. It proposes to learn such flows by using continous flows, where the vector field is the gradient of an invariant function. Defining invariant functions is usually much easier than defining equivariant vector fields. This framework is then tested on S^2 with the action rotations along a fixed axis, and on SU(N) with the action of  SU(N) on itself by conjugation, for N=2,3.", "questions": "", "limitations": "", "rating": 7, "confidence": 5, "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": "Summary of the review: On the positive side, the approach of first building invariant functions in order to get equivariant vector fields is elegant. The experiments also show clear improvement over past work. On the negative side, I find the that authors over-claim their contribution, and Theorem 4 (which is crucial to prove that this approach is general enough) is badly written.\n\n\n== Positive points ==\nMoving the problems from \"building equivariant maps to building invariant maps\" is not just elegant, it is also of high practical value since building invariant functions is often much easier than building equivariant ones. This is particularly true when we also have the requirements that our functions should be invertible. The theory justifying this approach is laid out clearly, and the proof of Theorem 1 looks correct.\n\nThe experiments provided by the authors are convincing, and show clear improvements over the reference paper of Boyda et al.\n\n== Negative points ==\nA big over-claim I have found in the paper is the claim of \"learning gauge invariant densities over SU(n) in the context of QFT\". This claim appears in the abstract, and is repeated throughout the paper. This is totally incorrect. This papers addresses the issue of learning conjugation equivariant densities on SU(n), and tests it for n=2, 3. While learning such densities is a crucial step *towards* flows for QFT in Boyda et al, it is definitely not the whole story and this paper shows no QFT experiment. In short, the authors are right to use QFT as a motivation for their work, but they should *not* claim to be learning gauge invariant densities, or be doing QFT in this paper.\n\nI tried really hard to understand Theorem 4 about the sufficency of flows generated by invariant functions, but failed to do so. The authors should rewrite this. For example (but this might not be the only problem), on line 219, how can then take a derivative with respect to time, when time is not used to define \\rho or \\pi?\n\nIn the experiment section, the authors build conjugation equivariant flows on SU(N) for N=2 or 3 only, when Boyda et al also showed results for SU(9). Did the authors find that there method was harder to implement for N > 3?\n\n\n==========================\nUpdate following rebuttal and discussions with authors/reviewers\n==========================\n\nThe discussion with other reviewers and with the authors has been particularly fruitful. I am happy to raise my rating to \"7. Good Paper - Accept\", conditioned on the few points that were raised and discussed during the reviews and rebuttal period:\n* make it clear (including in the abstract and introduction), that the paper does not contain any Lattice QCD or QFT experiment. Instead, lattice QFT is a motivation for learning SU(3) conjugation equivariant flows.\n* as pointed out by other reviewers, the relationship of Theo 1-3 with Lemma 2 in [1] and Theo 2 in [2] should be made much clearer. In particular, Theo 1-3 are extensions to the Riemannian case of ideas that were previously explored in the Euclidean case.\n* Theo 4 needs to be completely rewritten, using all the added explanations and improvements discussed during the rebuttal.\n\n[1] Papamakarios, G., Nalisnick, E., Rezende, D.J., Mohamed, S. and Lakshminarayanan, B., 2019. Normalizing flows for probabilistic modeling and inference. arXiv preprint arXiv:1912.02762.\n\n[2] Köhler, J., Klein, L. and Noé, F., 2020, November. Equivariant flows: exact likelihood generative learning for symmetric densities. In International Conference on Machine Learning (pp. 5361-5370). PMLR.", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "Continuous ODE flows are powerful but can be quite expensive, and ODE solvers do not usually deal with manifolds. The paper would benefit from having a short discussion on these points. For example, how many integration steps were used? Does this number vary over training? What ODE solver was used (Euler, some Runge-Kutta, ...)? What technique did the authors use to ensure their ODE solver stayed on SU(N) or S^2 (doing x + dt * grad_x is not going to stay on the manifold)?", "review_text": "", "extra_scores": {"ethical_concerns": "N/A", "needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "10"}, "extra_text": {}}, {"review_id": "CWLNhV0TXB1", "reviewer": "Reviewer_tx8j", "summary": "The paper proposes a general method to construct equivariant continuous normalizing flows on Riemannian manifolds. For this, the authors built on (more pedestrian) results of Koehler et al using the fact that an invariant potential can be used to obtain an equivariant vector field. This vector field can then in turn be used to define a continuous normalizing flow which is equivariant under (a certain  left or right action of) the symmetry group. Finally, the push-forward of a invariant density under the equivariant flow is a invariant density with respect to the symmetry group.  Specifically, the manuscript considers the case for which G is (a subgroup of) the isometry group of the manifold. \n\nThe developed framework is then applied to SU(N) Lie groups for which the conjugation by SU(N) is (a subgroup of) the isometry group. This application is of relevance in the context of lattice field theory, i.e. numerical simulation of (mostly strongly interacting) quantum field theories, for which ML techniques have recently gained significant attention. ", "questions": "", "limitations": "", "rating": 8, "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": "The paper is well-written, insightful and rigorous. The underlying theory is clearly explained (although the presentation necessarily requires a solid background in Differential Geometry). The proofs of the main results are derived in a transparent and pedagogical manner in the Supplement while the main text focuses on summarizing these main results and their interplay for the proposed methods which makes the paper very readable. \n\nPersonally, I find the experimental sections a bit weaker. The LQFT application is of high relevance but in this context one is mainly interested in gauge symmetries, i.e. each lattice site can be conjugated by different SU(N) group elements, while the corresponding density is invariant under such local symmetry actions. The authors mainly consider the case of a single lattice site. Of course, this is a very important building block for a fully gauge-invariant flow and thus still highly interesting. However, it would be preferable if the authors had demonstrated the superiority of their approach over existing results by Boyada et al by training a flow to approximate the density of e.g. 2d quenched SU(N) lattice gauge theory comparing for example the ESS or the error of the topological susceptibility for fixed number of samples. The reported results would suggest that their method also helps in this case but an experimental validation of this expectation would be desirable.\n\n\nApart from this weakness, I find the manuscript very convincing and insightful. I also want to encourage the authors to release their code such that the community can benefit from it.\n\nA few minor suggestions:\nL136-137: d and \\partial an the equation should be used consistently in the numerator and denumerator\nL189-190: larger spacing between the two equations would improve readability\nL214, L216-218: distribution -> density\nL506-507: Du -> D_u\n\n\n========================== Update after rebuttal and discussions with authors/reviewers ==========================\nI want to thank the authors and also the other reviewers for a very insightful discussion period. I want to strongly encourage the authors to implement the changes already suggested by the other reviewers. ", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "Both limitations and societal impact are adequately discussed.", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "3"}, "extra_text": {}}]

Author rebuttal:

Meta-review: The paper proposes a theoretical framework for constructing normalizing flows on Riemannian manifolds that are equivariant with respect to the action of an isometry group. This is an interesting contribution and highly relevant for distribution modelling under known symmetries in non-Euclidean spaces, which comes up often in problems in physics.

Overall, the reviews are positive so I'm happy to recommend acceptance. There has been a fruitful discussion between the reviewers and the authors that I hope will improve the final version further.

**Comments to the authors**: I would like to see the improvements suggested by the reviewers in the final version of the paper. For clarity, I'm copying below the changes requested by the reviewers that I'd like to see implemented:

1. Make it clear that the paper doesn't contain an application to QFT or Lattice QCD, but uses these as motivation only.
2. Moderate the novelty of Theorems 1-3 and make it clear how they relate to similar results from previous work.
3. Revise Theorem 4 and its proof following the discussion with the reviewers.

I would also like to highlight the review by Reviewer e1R7 and specifically this statement: "as currently written, [the paper] will be accessible and hence of interest only to a tiny fraction of the NeurIPS community". Indeed, I'd encourage you to reflect carefully on how this work can be communicated effectively so that it can maximize its future impact.
Decision: Accept (Poster)