Density estimation on low-dimensional manifolds: an inflation-deflation approach

Normalizing Flows ( NFs ) are universal density estimators based on Neuronal Networks . However , this universality is limited : the density ’ s support needs to be diffeomorphic to a Euclidean space . In this paper , we propose a novel method to overcome this limitation without sacrificing the universality . The proposed method inflates the data manifold by adding noise in the normal space , trains an NF on this inflated manifold and , finally , deflates the learned density . Our main result provides sufficient conditions on the manifold and the specific choice of noise under which the corresponding estimator is exact . Our method has the same computational complexity as NFs , and does not require to compute an inverse flow . We also show that , if the embedding dimension is much larger than the manifold dimension , noise in the normal space can be well approximated by some Gaussian noise . This allows using our method for approximating arbitrary densities on non-flat manifolds provided that the manifold dimension is known . 1 INTRODUCTION . Many modern problems involving high dimensional data are formulated probabilistically . Key concepts , such as Bayesian Classification , Denoising or Anomaly Detection , rely on the data generating density p∗ ( x ) . Therefore , a main research area and of crucial importance is learning this data generating density p∗ ( x ) from samples . For the case where the corresponding random variable X ∈ RD takes values on a manifold diffeomorphic to RD , a Normalizing Flow ( NF ) can be used to learn p∗ ( x ) exactly ( Huang et al. , 2018 ) . Recently , a few attempts have been made to overcome this topological constraint . However , to do so , all of these methods either need to know the manifold beforehand ( Gemici et al . ( 2016 ) , Rezende et al . ( 2020 ) ) or they sacrifice the exactness of the estimate ( Cornish et al . ( 2019 ) , Dupont et al . ( 2019 ) ) . Our goal in this paper is to overcome both the aforementioned limitations of using NFs for density estimation on Riemannian manifolds . Given data points from a d−dimensional Riemannian manifold embedded in RD , d < D , we first inflate the manifold by adding a specific noise in the normal space direction of the manifold , then train an NF on this inflated manifold , and , finally , deflate the trained density by exploiting the choice of noise and the geometry of the manifold . See Figure 1 for a schematic overview of these points . Our main theorem states sufficient conditions on the manifold and the type of noise we use for the inflation step such that the deflation becomes exact . To guarantee the exactness , we do need to know the manifold as in e.g . Rezende et al . ( 2020 ) because we need to be able to sample in the manifold ’ s normal space . However , as we will show , for the special case where D d , the usual Gaussian noise is an excellent approximation for a noise in the normal space component . This allows using our method for approximating arbitrary densities on Riemannian manifolds provided that the manifold dimension is known . In addition , our method is based on a single NF without the necessity to invert it . Hence , we don ’ t add any additional complexity to the usual training procedure of NFs . Notations : We denote the determinant of the Gram matrix of f as gf ( x ) : = |det ( Jf ( x ) TJf ( x ) ) | where Jf ( x ) is the Jacobian of f . We denote the Lebesque measure in Rn as λn . Random variables will be denoted with a capital letter , say X , and its corresponding state space with the calligraphical version , i.e . X . Small letters correspond to vectors with dimensionality given by context . The letters d , D , n , and N are always natural numbers . 2 BACKGROUND AND PROBLEM STATEMENT . An NF transforms a known auxiliary random variable by using bijective mappings parametrized by Neuronal Networks such that the given data points are samples from this transformed random variable , see Papamakarios et al . ( 2019 ) . Formally , an NF is a diffeomorphism Fθ : U → X and induces a density on X through pθ ( x ) = ( gF ( u ) − 1 2 pU ( u ) where pU ( u ) is known and u = F−1θ ( x ) . The parameters θ are updated such that the KL-divergence between p∗ ( x ) and pθ ( x ) , DKL ( p ∗ ( x ) ||pθ ( x ) ) = Ex∼p∗ ( x ) [ log pθ ( x ) ] + const . ( 1 ) is minimized . If Fθ is expressive enough , it was proven that in the limit of infinitely many samples , updating θ to minimize this objective function converges to a θ∗ such that it holds PX -almost surely p∗ ( x ) = pθ∗ ( x ) , see ( Huang et al. , 2018 ) . More generally , let X ∈ X ⊂ RD be generated by an unobserved random variable Z ∈ Z ⊂ Rd with density π ( z ) , that is X = f ( Z ) for some function f : Z → X where typically d < D. In Gemici et al . ( 2016 ) , f is an embedding1 , and it was shown that one can calculate probabilities such as PX ( A ) for measurable A ⊂ X using a density p∗ ( x ) with respect to the volume form dVf induced by f , that is P ( X ∈ A ) = ∫ f−1 ( A ) π ( z ) dz = ∫ A p∗ ( x ) dVf ( x ) ( 2 ) with p∗ ( x ) = π ( z ) gf ( z ) − 1 2 and dVf ( x ) = √ gf ( z ) dz where z = f−1 ( x ) . Hence , given an explicit mapping f and samples from p∗ ( x ) , we can learn the unknown density π ( z ) using a usual NF in Rd . However , in general , the generating function f is either unknown or not an embedding creating numerical instabilities for training inputs close to singularity points . In Brehmer & Cranmer ( 2020 ) , f and the unknown density π are learned simultaneously . The main idea is to define f as a level set of a usual flow in RD and train it together with the flow in Rd used to learn π ( z ) . To evaluate the density , one needs to invert f and thus this approach may be very slow for high-dimensional data . Besides , to guarantee that f learns the manifold they proposed several ad hoc training strategies . We tie in with the idea to use an NF for learning p∗ ( x ) with unknown f and study the following problem . 1Thus , a regular continuously differentiable mapping ( called immersion ) which is , restricted to its image , a homeomorphism . Problem 1 Let X be a d−dimensional manifold embedded in RD . Let X = f ( Z ) be a random variable generated by an embedding f : Rd → RD and a random variable Z ∼ π ( z ) in Rd . Given N samples from p∗ ( x ) as described above , find an estimator p̂ of p∗ such that in the limit of infinitely many samples we have that p̂ ( x ) = p∗ ( x ) , PX − almost surely . 3 METHODS . To solve Problem 1 , we want to exploit the universality of NFs . We want to inflate X such that the inflated manifold X̃ becomes diffeomorphic to a set U on which a simple density exists . Doing so , allows us to learn the inflated density p ( x̃ ) exactly using a single NF , see Section 2 . Then , given such an estimator for the modified density , we approximate p∗ ( x ) and give sufficient conditions when this approximation is exact . 3.1 THE INFLATION STEP . Given a sample x of X , if we add some noise ε ∈ RD to it , the resulting new random variable X̃ = X + ε has the following density q ( x̃ ) = ∫ X q ( x̃|x ) dPX ( x ) . ( 3 ) Denote the tangent space in x as Tx and the normal space asNx . By definition , Nx is the orthogonal complement of Tx . Therefore , we can decompose the noise ε into its tangent and normal component , ε = εt + εn . In the following , we consider noise in the normal space only , i.e . εt = 0 , and denote the density of the resulting random variable as qn ( x̃ ) . The corresponding noise density qn ( x̃|x ) has mean x and domain Nx . We denote the support of qn ( ·|x ) by Nqn ( ·|x ) . The random variable X̃ = X + εn is now defined on X̃ = ⋃ x∈X Nqn ( ·|x ) . We want X̃ to be diffeomorphic to a set U on which a known density can be defined . Example 1 ( a ) Let X = S1 = { x ∈ R2 | ||x|| = 1 } be the unit circle . For each x ∈ S1 there exists z ∈ [ 0 , 2π ) such that x = er ( x ) = ( cos ( z ) , sin ( z ) ) T . To sample a point x̃ in Nx , which is spanned by er ( x ) , we sample a scalar value γ and set x̃ = x + γer ( x ) . With γ ∼ Uniform [ −1 , 1 ) , we have that X̃ = ⋃ x∈X { x+ γer ( x ) |γ ∈ [ −1 , 1 ) } = { x ∈ R2 | ||x||2 < 2 } ( 4 ) which is the open disk with radius 2 . The open disk is diffeomorphic to ( 0 , 1 ) × ( 0 , 1 ) . Thus , qn ( x̃ ) can be learned by a single NF F−1 and pU ( u ) = Uniform ( ( 0 , 1 ) × ( 0 , 1 ) ) as reference . ( b ) As in ( a ) , we consider the unit circle . Now we set γ to be a χ2− distribution with support [ −1 , ∞ ) . Then X̃ = ⋃ x∈X { x+ γer ( x ) |γ ∈ [ −1 , ∞ ) } = R2 . ( 5 ) Thus , qn ( x̃ ) can be learned by a single NF F−1 and pU ( u ) = N ( u ; 0 , ID ) as reference . Both cases can be analogously extended to higher dimensions . Remark 1 To be precise , the random variable εn is generated by a random variable in RD−d , say Γ , with measure PΓ . Then , qn ( x̃|x ) is the density of the pushforward of the noise measure PΓ with regard to the mapping h : RD−d → Nx . Hence , formally , the density q ( x̃|x ) is with respect to the induced volume form dVh , see Section 2 . However , if we choose an orthonormal basis for Nx , say n ( 1 ) , . . . , n ( D−d ) , then we have that x̃ = h ( γ ) = Aγ + x where the columns of A ∈ RD× ( D−d ) are given by these basis vectors , i.e . A = [ n ( 1 ) , . . . , n ( D−d ) ] . Thus , the Gram determinant of h is gh = det ( ATA ) = 1 and we have that dVh ( x̃ ) = dγ where dγ denotes the volume form with respect to λD−d . In this case , we can think of qn ( x̃|x ) as a density with respect to λD−d . If qn ( x̃|x ) depends only on ||x̃ − x|| , as it is for the Gaussian distribution , we have that qn ( x̃|x ) = qn ( ||x̃ − x|| ) = qn ( ||γ|| ) because h is an isomorphism . Thus , for this case it holds that qn ( x̃|x ) dVh ( x̃ ) = qn ( ||γ|| ) dγ . Then , for convenience , we may abuse notation by writing γ ∼ q ( x̃|x ) or εn ∼ r ( γ ) where r ( γ ) is the density of PΓ with respect to λD−d .

Summary: In this paper, the authors address main limitations of Normalizing Flows (NFs) method for estimation of density functions on manifolds. Since NFs requires the support of density function to cover the whole Euclidean space, they propose to add noise (inflate) to apply NF. The paper is nicely written with clear introduction of basic concepts very useful for non-expert readers. They provide theoretical guarantees on the variance and type of added noise that make the method work and illustrate with synthetic experiments.