Moser Flow: Divergence-based Generative Modeling on Manifolds

1 Introduction . The major successes of deep generative models in recent years are primarily in domains involving Euclidean data , such as images ( Dhariwal and Nichol , 2021 ) , text ( Brown et al. , 2020 ) , and video ( Kumar et al. , 2019 ) . However , many kinds of scientific data in the real world lie in non-Euclidean spaces specified as manifolds . Examples include planetary-scale data for earth and climate sciences ( Mathieu and Nickel , 2020 ) , protein interactions and brain imaging data for life sciences ( Gerber et al. , 2010 ; Chen et al. , 2012 ) , as well as 3D shapes in computer graphics ( Hoppe et al. , 1992 ; Kazhdan et al. , 2006 ) . Existing ( Euclidean ) generative models can not be effectively applied in these scenarios as they would tend to assign some probability mass to areas outside the natural geometry of these domains . An effective way to impose geometric domain constraints for deep generative modeling is to design normalizing flows that operate in the desired manifold space . A normalizing flow maps a prior ( source ) distribution to a target distribution via the change of variables formula ( Rezende and Mohamed , 2015 ; Dinh et al. , 2016 ; Papamakarios et al. , 2019 ) . Early work in this direction proposed invertible architectures for learning probability distributions directly over the specific manifolds defined over spheres and tori ( Rezende et al. , 2020 ) . Recently , Mathieu and Nickel ( 2020 ) proposed to extend continuous normalizing flows ( CNF ) ( Chen et al. , 2018 ) for generative modeling over Riemannian manifolds wherein the flows are defined via vector fields on manifolds and computed as the solution to an associated ordinary differential equation ( ODE ) . CNFs have the advantage that the neural 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . network architectures parameterizing the flow need not be restricted via invertibility constraints . However , as we show in our experiments , existing CNFs such as FFJORD ( Grathwohl et al. , 2018 ) and Riemannian CNFs ( Mathieu and Nickel , 2020 ) can be slow to converge and the generated samples can be inferior in capturing the details of high fidelity data densities . Moreover , it is a real challenge to apply Riemannian CNFs to complex geometries such as general curved surfaces . To address these challenges , we propose Moser Flows ( MF ) , a new class of deep generative models within the CNF family . An MF models the desired target density as the source density minus the divergence of an ( unrestricted ) neural network . The divergence is a local , linear differential operator , easy to approximate and calculate on manifolds . By drawing on classic results in differential geometry by Moser ( 1965 ) and Dacorogna and Moser ( 1990 ) , we can show that this parameterization induces a CNF solution to the change-of-variables formula specified via an ODE . Since MFs directly parameterize the model density using the divergence , unlike other CNF methods , we do not require to explicitly solve the ODE for maximum likelihood training . At test-time , we use the ODE solver for generation . We derive extensions to MFs for Euclidean submanifolds that efficiently parameterize vector fields projected to the desired manifold domain . Theoretically , we prove that Moser Flows are universal generative models over Euclidean submanifolds . That is , given a Euclidean submanifold M and a target continuous positive probability density µ overM , MFs can push arbitrary positive source density ν overM to densities µ̄ that are arbitrarily close to µ . We evaluate Moser Flows on a wide range of challenging real and synthetic problems defined over many different domains . On synthetic problems , we demonstrate improvements in convergence speed for attaining a desired level of details in generation quality . We then experiment with two kinds of complex geometries . First , we show significant improvements of 49 % on average over Riemannian CNFs ( Mathieu and Nickel , 2020 ) for density estimation as well as high-fidelity generation on 4 earth and climate science datasets corresponding to global locations of volcano eruptions , earthquakes , floods , and wildfires on spherical geometries . Next and last , we go beyond spherical geometries to demonstrate for the first time , generative models on general curved surfaces . 2 Preliminaries . Riemannian manifolds . We consider an orientable , compact , boundaryless , connected ndimensional Riemannian manifoldM with metric g. We denote points in the manifold by x , y ∈M . At every point x ∈ M , TxM is an n-dimensional tangent plane toM . The metric g prescribes an inner product on each tangent space ; for v , u ∈ TxM , their inner product w.r.t . g is denoted by 〈v , u〉g . X ( M ) is the space of smooth ( tangent ) vector fields toM ; that is , if u ∈ X ( M ) then u ( x ) ∈ TxM , for all x ∈ M , and if u written in local coordinates it consists of smooth functions . We denote by dV the Riemannian volume form , defined by the metric g over the manifoldM . In particular , V ( A ) = ∫ A dV is the volume of the set A ⊂M . We consider probability measures overM that are represented by strictly positive continuous density functions µ , ν : M→ R > 0 , where µ by convention represents the target ( unknown ) distribution and ν represents the source ( prior ) distribution . µ , ν are probability densities in the sense their integral w.r.t . the Riemannian volume form is one , i.e. , ∫ M µdV = 1 = ∫ M νdV . It is convenient to consider the volume forms that correspond to the probability measures , namely µ̂ = µdV and ν̂ = νdV . Volume forms are differential n-forms that can be integrated over subdomains ofM , for example , pν ( A ) = ∫ A ν̂ is the probability of the event A ⊂M . Continuous Normalizing Flows ( CNF ) on manifolds operate by transforming a simple source distribution through a map Φ into a highly complex and multimodal target distribution . A manifold CNF , Φ : M→M , is an orientation preserving diffeomorphism from the manifold to itself ( Mathieu and Nickel , 2020 ; Lou et al. , 2020 ; Falorsi and Forré , 2020 ) . A smooth map Φ : M→M can be used to pull-back the target µ̂ according to the formula : ( Φ∗µ̂ ) z ( v1 , . . . , vn ) = µ̂Φ ( z ) ( DΦz ( v1 ) , . . . , DΦz ( vn ) ) , ( 1 ) where v1 , . . . , vn ∈ TzM are arbitrary tangent vectors , DΦz : TzM→ TΦ ( z ) M is the differential of Φ , namely a linear map between the tangent spaces toM at the points z and Φ ( z ) , respectively . By pulling-back µ̂ according to Φ and asking it to equal to the prior density ν , we get the manifold version of the standard normalizing equation : ν̂ = Φ∗µ̂ . ( 2 ) If the normalizing equation holds , then for an event A ⊂M we have that pν ( A ) = ∫ A ν̂ = ∫ A Φ∗µ̂ = ∫ Φ ( A ) µ̂ = pµ ( Φ ( A ) ) . Therefore , given a random variable z distributed according to ν , then x = Φ ( z ) is distributed according to µ , and Φ is the generator . One way to construct a CNF Φ is by solving an ordinary differential equation ( ODE ) ( Chen et al. , 2018 ; Mathieu and Nickel , 2020 ) . Given a time-dependent vector field vt ∈ X ( M ) with t ∈ [ 0 , 1 ] , a one-parameter family of diffeomorphisms ( CNFs ) Φt : [ 0 , 1 ] ×M→M is defined by d dt Φt = vt ( Φt ) , ( 3 ) where this ODE is initialized with the identity transformation , i.e. , for all x ∈ M we initialize Φ0 ( x ) = x . The CNF is then defined by Φ ( x ) = Φ1 ( x ) . Example : Euclidean CNF . Let us show how the above notions boil down to standard Euclidean CNF for the choice ofM = Rn , and the standard Euclidean metric ; we denote z = ( z1 , . . . , zn ) ∈ Rn . The Riemannian volume form in this case is dz = dz1∧dz2∧· · ·∧dzn . Furthermore , µ̂ ( z ) = µ ( z ) dz and ν̂ ( z ) = ν ( z ) dz . The pull-back formula ( equation 1 ) in coordinates ( see e.g. , Proposition 14.20 in Lee ( 2013 ) ) is Φ∗µ̂ ( z ) = µ ( Φ ( z ) ) det ( DΦz ) dz , where DΦz is the matrix of partials of Φ at point z , ( DΦz ) ij = ∂Φ i ∂zj ( z ) . Plugging this in equation 2 we get the Euclidean normalizing equation : ν ( z ) = µ ( Φ ( z ) ) det ( DΦz ) . ( 4 ) 3 Moser Flow . Moser ( 1965 ) and Dacorogna and Moser ( 1990 ) suggested a method for solving the normalizing equation , that is equation 2 . Their method explicitly constructs a vector field vt , and the flow it defines via equation 3 is guaranteed to solve equation 2 . We start by introducing the method , adapted to our needs , followed by its application to generative modeling . We will use notations introduced above . 3.1 Solving the normalizing equation . Moser ’ s approach to solving equation 2 starts by interpolating the source and target distributions . That is , choosing an interpolant αt : [ 0 , 1 ] ×M→ R > 0 , such that α0 = ν , α1 = µ , and ∫ M αtdV = 1 for all t ∈ [ 0 , 1 ] . Then , a time-dependent vector field vt ∈ X ( M ) is defined so that for each time t ∈ [ 0 , 1 ] , the flow Φt defined by equation 3 satisfies the continuous normalization equation : Φ∗t α̂t = α̂0 , ( 5 ) where α̂t = αtdV is the volume form corresponding to the density αt . Clearly , plugging t = 1 in the above equation provides a solution to equation 2 with Φ = Φ1 . As it turns out , considering the continuous normalization equation simplifies matters and the sought after vector field vt is constructed as follows . First , solve the partial differential equation ( PDE ) over the manifoldM div ( ut ) = − d dt αt , ( 6 ) where ut ∈ X ( M ) is an unknown time-dependent vector field , and div is the Riemannian generalization to the Euclidean divergence operator , divE = ∇· . This manifold divergence operator is defined by replacing the directional derivative of the Euclidean space with its Riemannian version , namely the covariant derivative , div ( u ) = n∑ i=1 〈∇eiu , ei〉g , ( 7 ) where { ei } ni=1 is an orthonormal frame according to the Riemannian metric g , and ∇ξu is the Riemannian covariant derivative . Note that here we assume thatM is boundaryless , otherwise we need ut to be also tangent to the boundary ofM . Second , define vt = ut αt . ( 8 ) Theorem 2 in Moser ( 1965 ) implies : Theorem 1 ( Moser ) . The diffeomorphism Φ = Φ1 , defined by the ODE in equation 3 and vector field vt in equation 8 solves the normalization equation , i.e. , equation 2 . The proof of this theorem in our case is provided in the supplementary for completeness . A simple choice for the interpolant αt that we use in this paper was suggested in Dacorogna and Moser ( 1990 ) : αt = ( 1− t ) ν + tµ . ( 9 ) The time derivative of this interpolant , i.e. , ddtαt = µ − ν , does not depend on t. Therefore the vector field can be chosen to be constant over time , ut = u , and the PDE in equation 6 takes the form div ( u ) = ν − µ , ( 10 ) and consequently vt takes the form vt = u ( 1− t ) ν + tµ . ( 11 ) Figure 1 shows a one dimensional illustration of Moser Flow .

Summary: This paper presents Moser Flow a new class of Continuous Normalizing Flows for Riemannian Manifolds. The main approach exploits results in differential geometry (Moser 1965) to construct a vector field that satisfies the normalizing equation given by the pullback of the volume form under $\Phi$. The approach is both sensible and provides direct advantages over Riemmanian CNF's as training does not require solving the ODE using a solver. The authors also consider Moser Flows for Euclidean Submanifolds under the Induced metric and show that Moser Flows can learn arbitrary distributions on manifolds of interest (---i.e. compact, boundryless etc...). Finally, the experiments conducted are both interesting in their variety and illuminating in that they show the effectiveness of the proposed method. Overall, this paper represents a non-trivial step forward---manifold or otherwise---for Continuous Normalizing Flow research.