Equivariant Manifold Flows

1 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 , 21 ] , less attention has been allocated to learning models over non-Euclidean spaces , despite the fact that many problems require a manifold structure . Density learning over nonEuclidean spaces has applications ranging from quantum field theory in physics [ 44 ] to motion estimation in robotics [ 16 ] to protein-structure prediction in computational biology [ 22 ] . Continuous normalizing flows ( CNFs ) [ 6 , 21 ] are 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 to the setting of density learning on Riemannian manifolds . However , for many applications in the natural sciences , this construction is insufficient as it can not properly model necessary symmetries . For example , such symmetry requirements arise when sampling coupled particle systems in physical chemistry [ 26 ] or sampling for use in SU ( n ) 1 lattice gauge theories in theoretical physics [ 3 ] . More precisely , these symmetries are invariances with respect to action by an isometry subgroup of the underlying manifold . * indicates equal contribution 1SU ( n ) denotes the special unitary group SU ( n ) = { X 2 Cn⇥n | X⇤X = I , det ( X ) = 1 } . 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . For example , consider the task of learning a density on the sphere that is invariant to rotation around an axis ; this is an example of learning an isometry subgroup invariant2 density . For a less trivial example , note that when learning a flow-based sampler for SU ( n ) in the context of lattice QFT [ 3 ] , the learned density must be invariant to conjugation by SU ( n ) ( see Figure 1 for a density on SU ( 3 ) that exhibits the requisite symmetry ) . One might naturally attempt to 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 structure3 . 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 a 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 RPn , a quotient of Sn [ 28 ] ) . 2 . Second , even if the above constructs can be realized , the quotient manifold often has a boundary , which precludes the use of a manifold CNF . To illustrate this point , consider the simple case of the sphere invariant to rotation about an axis ; the quotient manifold is a 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 distribution . This is the case for Boyda et al . [ 3 ] : the flow construction over the quotient induces abnormalities in the density . Motivated by the above drawbacks , we design a manifold continuous normalizing flow on 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 on manifolds . A summary of our paper ’ s contributions is as 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 SU ( 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 comparing 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 a 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 . 2This specific isometry subgroup is known as the isotropy group at a point of the sphere intersecting the axis . 3In particular , the isometry subgroup action needs to be smooth , free , and proper to ensure the quotient will be a manifold by the Quotient Manifold Theorem [ 28 ] . Normalizing Flows on Manifolds Normalizing flows on Euclidean space have long been touted as powerful generative models [ 6 , 10 , 21 ] . Similar to GANs [ 20 ] and VAEs [ 24 ] , normalizing flows learn to map samples from a tractable prior density to a target density . However , unlike the aforementioned models , normalizing flows account for changes in volume , enabling exact evaluation of the output probability density . In a 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 [ 4 ] , many methods have extended normalizing flows to Riemannian manifolds . Rezende et al . [ 38 ] introduced constructions specific to tori and spheres , while Bose et al . [ 2 ] introduced constructions for hyperbolic space . Following this work , Falorsi and Forré [ 15 ] , 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 [ 30 ] and generalizes them further to 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 [ 7–9 , 18 , 27 , 37 ] . Köhler et al . [ 26 ] , in particular , used equivariant normalizing flows to 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 learning4 , for certain test tasks it is a strict requirement . For example , Boyda et al . [ 3 ] introduced equivariant flows on SU ( n ) for use in lattice gauge theories , where the modelled distribution must be conjugation invariant . However , beyond conjugation invariant learning on SU ( n ) [ 3 ] , 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 Background . In this section , we provide a 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 a more detailed introduction to Riemannian geometry , we refer the reader to textbooks such as Lee [ 28 ] and Kobyzev et al . [ 25 ] . 3.1 Riemannian Geometry . A Riemannian manifold ( M , h ) is an n-dimensional manifold with a smooth collection of inner products ( hx ) x2M for every tangent space TxM . The Riemannian metric h induces a distance dh on the manifold . A diffeomorphism f : M ! M is a differentiable bijection with differentiable inverse . A diffeomorphism f : M ! M is called an isometry if h ( Dxf ( u ) , Dxf ( v ) ) = h ( u , v ) for all tangent vectors u , v 2 TxM where Dxf 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 Rn . For a function f : M ! R , we define the Riemannian gradient rxf to be the vector on TxM such that h ( rxf , v ) = Dxf ( v ) for v 2 TxM . 3.2 Normalizing Flows on Manifolds . Manifold Normalizing Flow Let ( M , h ) be a Riemannian manifold . A normalizing flow on M is a diffeomorphism f✓ : M ! M ( parametrized by ✓ ) that transforms a prior density ⇢ to model 4For example , asteroid impacts on the sphere can be modelled as being approximately invariant to rotation about the Earth ’ s axis . density ⇢f✓ . The model distribution can be computed via the Riemannian change of variables5 : ⇢f✓ ( x ) = ⇢ ⇣ f 1 ✓ ( x ) ⌘ dethDxf 1✓ . Manifold Continuous Normalizing Flow A manifold continuous normalizing flow with base point z is a function : [ 0,1 ) ! M that satisfies the manifold ODE d ( t ) dt = X ( ( t ) , t ) , ( 0 ) = z . We define FX , T : M ! M , z 7 ! FX , T ( z ) to map any base point z 2 M to the value of the CNF starting at z , evaluated at time T . This function is known as 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 2 G on M by the map Lg : M ! M. Equivariant and Invariant Functions We say that a function f : M ! N is equivariant if , for all isometries gx : M ! M and gy : N ! N , we have f gx = gy f . We say a function f : M ! N is invariant if f gx = f . Equivariant Vector Fields Let X : M⇥ [ 0,1 ) ! TM , X ( m , t ) 2 TmM be a time-dependent vector field on manifold M , with base point x0 2 M. X is a G-equivariant vector field if 8 ( m , t ) 2 M⇥ [ 0,1 ) , X ( Lgm , t ) = ( DmLg ) X ( m , t ) . Equivariant Flows A flow f : M ! M is G-equivariant if it commutes with actions from G , i.e . we have Lg f = f Lg . Invariance of Density A density ⇢ on a manifold M is G-invariant if , for all g 2 G and x 2 M , ⇢ ( Lgx ) = ⇢ ( x ) , where Lg is the action of g on x .

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.