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: The paper introduces a new type of continuous normalizing flow (CNF) for manifold-valued data. Unlike in previous works, this Moser Flow (MF) models a density as a base density minus the divergence of a neural network. The authors prove that under certain assumptions this defines a probability distribution over the manifold. This approach has a different trade-off between inference and generation from other approaches: while the model density can be evaluated very cheaply, sampling from this model requires an ODE solver. The approach is qualitatively demonstrated on toy datasets on tori and the Stanford bunny, and quantitatively evaluated on a few datasets on 2-spheres.