Denoising Normalizing Flow

1 Introduction . Given samples from the data-density p ( x ) , key objectives in probabilistic Machine Learning are 1. estimating p ( x ) ( density estimation ) , 2. generating new data points from p ( x ) ( sampling ) , and 3. finding low-dimensional representations of the data ( inference ) . The three main methods used to perform these tasks are Normalizing Flows ( NFs ) [ 34 ] , Generative Adversarial Networks ( GANs ) [ 17 ] , and Variational Autoencoders ( VAEs ) [ 26 ] . Among those methods , only the VAE does check all desired objectives by default . However , it does so by sacrificing the sample quality ( compared to GANs ) , and only learning a lower bound on p ( x ) ( rather than the exact value as NFs do ) . Finding new ways to meet these key objectives , based on either the known main methods or new ones , is an active and important research area . How can we infer low-dimensional representations using NFs ? The standard NF requires the datadensity p ( x ) to have a support diffeomorphic to the entire data-space RD . However , the manifold hypothesis conjectures that the data-manifoldM lies close to a d-dimensional manifold embedded in RD , d < D. Thus , unfortunately , a standard NF can not be used to infer the latent space . Recently , Brehmer et al . proposed to overcome this limitation by first projecting the data into a d-dimensional space using the first d−components of a standard NF , and then using another NF to learn the latent distribution π ( u ) [ 10 ] . For their method to work they propose different learning schemes , separating the manifold learning from the density estimation . In this paper , we propose an easy and new way to learn low-dimensional representations with NFs . Our idea is based on the theoretical work derived in [ 20 ] where it was shown that by inflating the data-manifold with Gaussian noise ε ∼ N ( 0 , σ2ID ) , i.e . x̃ : = x + ε , the data-density p ( x ) can be well approximated by learning the inflated distribution qσ ( x̃ ) . More concretely , the main result in [ 20 ] states sufficient conditions on the choice of noise qσ ( x̃|x ) and type of manifold such that 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . qσ ( x ) = p ( x ) qσ ( x|x ) holds . Here , by adding a penalty term to the usual KL-divergence used to learn qσ ( x̃ ) , we ensure that the first d-components of the corresponding flow are noise insensitive and thus encode the manifold . This penalty term is essentially the objective function of a Denoising autoencoder ( DAE ) , and thus we call our method Denoising Normalizing Flow ( DNF ) . In summary , our contributions are the following : 1 . We propose a new method , the Denoising Normalizing Flow , which combines two previously well-known methods ( DAE and NF ) , and is able to • approximate p ( x ) , • sample new data x ∼ p ( x ) as a non-linear transformation of u ∼ N ( u ; 0 , Id ) , • infer low-dimensional latent variables u ∼ p ( u|x ) given x ∼ p ( x ) . 2 . We demonstrate on naturalistic data that our method learns meaningful latent representations without sacrificing the sample quality . Notations : We adapt the notation used in [ 10 ] . To further simplify it and avoid clutter , we denote the Gram matrix of g evaluated at g−1 ( x ) as Gg ( x ) : = Jg ( g −1 ( x ) ) TJg ( g −1 ( x ) ) ( 1 ) where Jg ( g−1 ( x ) ) T is the transpose of the Jacobian of g : Rd → RD , d ≤ D , evaluated at g−1 ( x ) . 2 Problem statement . In the following , we are going to show why standard NFs are not suited to infer low-dimensional representations of the given data . We end the section with the research question we are going to study in Section 3 . General setting : In generative modeling , it is assumed that the data x ∈M ⊂ RD is generated by a non-linear transformation g of some latent variables u ∈ U ⊂ Rd , i.e . x = g ( u ) where u ∼ π ( u ) . Typically , d < D and the latent distribution π ( u ) is assumed to be Gaussian , π ( u ) = N ( u ; 0 , Id ) . Hence , the latent random variable u generates x , and the data-density p ( x ) evaluated at x ∈ RD is given by p ( x ) = ∫ U π ( u ) δ ( x− g ( u ) ) du , ( 2 ) where δ denotes the Dirac function , see [ 3 ] . If d = D and g is a diffeomorphism , we have for x ∈M that p ( x ) = |detGg ( x ) |− 1 2π ( g−1 ( x ) ) . ( 3 ) Then , the target density p ( x ) can be learned , in principle , exactly using an NF [ 21 ] . In general , an NF is an embedding mapping U ⊂ Rd toM⊂ RD , see [ 32 ] or [ 27 ] for some recent reviews . Denoting this mapping as gθ : U →M and its parameters as θ , the induced density onM is given by pθ ( x ) = |detGgθ ( x ) |− 1 2 pu ( g −1 θ ( x ) ) , ( 4 ) where pu ( u ) is a known reference density ( usually set to be standard Gaussian ) . 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 . ( 5 ) is minimized . Thus , to learn θ efficiently , one needs to evaluate T1 ( x ) : = log pu ( g−1θ ( x ) ) and T2 ( x ) : = 1 2 log |detGgθ ( x ) | efficiently . Topological constraints : The evaluation of T1 ( x ) is efficient since we are free to choose the reference measure pu , and g−1θ ( x ) is the forward pass of a neural network constructed to be bijective [ 13 , 21 ] . Therefore , the majority of the NF literature focuses on designing clever flow architectures to be able to calculate T2 ( x ) efficiently without sacrificing the flow ’ s expressiveness ( i.e . the size of the space of embeddings able to learn ) . However , so far these architectures are constructed for d = D since in this case , the Jacobian of gθ is a square matrix , and thus T2 ( x ) becomes 1 2 log |detGgθ ( x ) | = log |det Jgθ ( g −1 θ ( x ) ) | . ( 6 ) Hence , calculating the Gram determinant of gθ efficiently amounts to calculating the determinant of the Jacobian of gθ efficiently . Popular choices are to construct gθ such that Jgθ is lower triangular , as in this case , the determinant is simply the product of diagonal elements . Unfortunately , for d < D , Jgθ is not a square matrix , and the full Gram determinant needs to be calculated . This makes NFs unsuitable for finding low-dimensional representations u of high-dimensional data points for large d as the computational complexity to calculate detGgθ is O ( d2D ) +O ( d3 ) . Research question : As mentioned in [ 10 ] , finding ways to design gθ such that T2 ( x ) can be efficiently calculated for d < D is an interesting research question . Here , we address it from a different angle . LetM be a d−dimensional manifold embedded in RD through g. Can we construct a mapping gθ : Rd →M to 1. generate x ∼ p ( x ) in 2 steps : ( a ) generate u ∼ N ( u ; 0 , Id ) and ( b ) set x = gθ ( u ) ? 2. infer u ∈ Rd such that x = gθ ( u ) ? 3. approximate detGg ( x ) efficiently ? 3 Denoising Normalizing Flow . We answer the research question based on the theoretical work developed in [ 20 ] . First , we briefly review this work , and then introduce the DNF . Preliminaries : In Section 2 , we discussed why classical NFs are not suited to infer low-dimensional representations of the given data . Also , if one is only interested in the value of p ( x ) , standard NF can not be used . Intuitively , a standard NF fψ with parameters ψ simply squeezes or expands a volume element where the net change is given by its Jacobian determinant . A volume element of a d−dimensional manifold is d−dimensional and thus has D−dimensional Lebesgue measure 0 . Thus , we are asking fψ to expand a d−dimensional volume to a D−dimensional one which will lead to a degeneration of |detGfψ ( x ) | and manifest in numerical instabilities . Therefore , [ 20 ] inflated the manifold by adding Gaussian noise to the data points.1 This inflated manifold is D−dimensional and thus a usual flow can be used to learn the corresponding density . Their main result states sufficient conditions on the choice of noise and type of manifoldM , such that the learned inflated distribution can be deflated , and p ( x ) is exactly retrieved . More precisely , given a random variable x ∼ p ( x ) with probability measure PX and taking values in a d−dimensional manifoldM embedded in RD , if we add some noise ε to it , the resulting new random variable x̃ = x+ ε has the following density : qσ ( x̃ ) = ∫ M qσ ( x̃|x ) dPX ( x ) . ( 7 ) In [ 20 ] , it was shown that if ( a ) the noise is only added in the ( D− d ) −dimensional normal space Nx in x , ( b ) the noise magnitude σ is sufficiently small , and ( c ) the manifoldM is sufficiently smooth and disentangled , the resulting inflated distribution evaluated at x̃ = x takes the following product form : qσ ( x ) = p ( x ) qσ ( x|x ) , ( 8 ) where qσ ( x|x ) is the normalization constant of the noise density . More concretely , ( a ) and ( b ) need to ensure that x is almost surely uniquely determined by x̃ as the orthogonal projection of x̃ onM . For this projection to be well-defined , a sufficient condition is that the manifolds reach number2 is finite [ 6 ] . A manifold where almost every point x̃ ∼ q ( x̃ ) in the inflated set has a unique projection on M was called Q−normally reachable in [ 20 ] , where Q denotes the collection of noise distributions qσ ( x̃|x ) . Their main Theorem proves that for anyQ−normally reachable manifold equation ( 8 ) holds . Therefore , if qσ ( x̃ ) can be learned exactly using a standard NF , the on-manifold density p ( x ) can be retrieved exactly . It was also shown that for the case where d D ( as it is generally assumed for high-resolution images ) , full Gaussian noise is an excellent approximation for a Gaussian in the normal space . 1Adding noise to circumvent the aforementioned degeneracy problem was also proposed in [ 24 ] . 2Informally , this reach condition ensures that a manifold is learnable through samples . Main idea : We use a standard flow fψ to learn qσ ( x̃ ) such that the corresponding density has the product form of equation ( 8 ) , and the first d−components u of the flow ’ s output f−1ψ ( x̃ ) are noiseinsensitive whereas the remaining ( D − d ) −components v remain noise-sensitive . Thus , intuitively , we want the first d−components to denoise the inflated data . DNF : Let fψ : RD → RD be a standard flow with reference measure pz ( z ) . We denote the first d−components of the flows output as u , and the remaining ones as v , i.e . ( u , v ) T = f−1ψ ( x̃ ) . More formally , u = u ( x̃ ) = Proju ( f −1 ψ ( x̃ ) ) , v = v ( x̃ ) = Projv ( f −1 ψ ( x̃ ) ) ( 9 ) with Proju ( z ) = ( z1 , . . . , zd ) and Projv ( z ) = ( zd+1 , . . . , zD ) for z ∈ RD . As reference measure pz , we choose pz ( ( u , v ) T ) = pu ( u ) pv ( v ) with pv ( v ) modelling the noise-sensitive part , and pu ( u ) the noise-insensitive part . In particular , if qσ ( x̃|x ) is a ( D− d ) −dimensional Gaussian distribution with covariance σ2ID−d , we set pv ( v ) = N ( v ; 0 , σ2ID−d ) . For pu ( u ) to model the noise-insensitive part , we want the image fψ ( u , 0 ) to be in the manifold . Therefore , we embed u back in RD by padding the missing coordinates with 0 , Pad ( u ) = ( u , 0 , . . . , 0︸ ︷︷ ︸ ( D−d ) -times ) T , ( 10 ) such that the operation Pad ( Proju ( ( u , v ) ) ) = ( u , 0 ) T ignores the noise-sensitive part v in the latent space . This operator allows us to define a denoising function rψ ( x̃ ) as rψ ( x̃ ) : = ( fψ ◦ Pad ◦ Proju ◦ f −1 ψ ) ( x̃ ) ( 11 ) and we regularize ψ by minimizing C ( ψ ) : = Ex∼p ( x ) Ex̃∼qσ ( x̃|x ) ||x− rψ ( x̃ ) || 2 ( 12 ) where || · || denotes the L2 norm . We have not specified the reference measure pu ( u ) . To facilitate the disentanglement of noiseinsensitivity and noise-sensitivity in the u and v variables , we transform u with yet another flow hφ with paramters φ and reference measure pu′ ( e.g . standard Normal ) . Now , our sampling procedure looks as follows : 1. sample u′ ∼ pu′ ( u′ ) , 2. apply hφ to obtain a sample u from pu ( u ) , i.e . u = hφ ( u′ ) , and 3. set x = fψ ( Pad ( u ) ) . Denoting θ = ( ψ , φ ) , our model to learn qσ ( x̃ ) is qθ ( x̃ ) =|detGfψ ( x̃ ) |− 1 2 pu ( u ( x̃ ) ) pv ( v ( x̃ ) ) =|detGfψ ( x̃ ) |− 1 2 |detGhφ ( u ( x̃ ) ) |− 1 2 pu′ ( h −1 φ ( u ( x̃ ) ) ) pv ( v ( x̃ ) ) . ( 13 ) Note that the reconstruction loss , equation ( 12 ) , is essentially the objective function for the Denoising Autoencoder introduced in [ 1 ] . Therefore , we call our method Denoising Normalizing Flow ( DNF ) and it is trained on LDNF ( θ ) : =DKL ( qσ ( x̃ ) ||qθ ( x̃ ) ) + λC ( ψ ) ( 14 ) where λ > 0 is the penalty hyperparameter and is trading the density estimation with the manifold learning . A graphical description of the DNF model is given in Figure 1 ( a ) , and an algorithmic description in the DNF Algorithm below.3 Answer to research question : If LDNF ( θ ) = 0 , the reconstruction error expressed in equation ( 12 ) is 0 . Thus , the generative story of the DNF is exactly the one described in point 1. of our research question with gθ : = fψ ◦ Pad ◦ hφ . The inverse , h−1φ ◦ Proju ◦ f −1 ψ , can be used to infer u ′ s.t . point 2. holds . Finally , to show the third claim , we additionally assume that U ( which is the domain of the manifold generating function g ) is diffeomorphic to Rd , s.t . without loss of generality we set π ( u ) = N ( u ; 0 , Id ) . Then , we exploit equation ( 8 ) and calculate |detGg ( x ) | efficiently with the help of fψ and hφ , see Proposition 1 and its proof in the supplementary . 3We show the algorithm for the general scenario where the manifold is unknown and thus noise can not be added to the normal space . We also ignore terms independent of θ in the calculation of LDNF , and denote this loss function as L′DNF . DNF Algorithm : Training of Denoising Normalizing Flow for qσ ( x̃|x ) = N ( x̃ ; x , σ2ID ) . For simplicity , we show a stochastic gradient descent with a constant learning rate . Alternative optimization methods and learning rate schedules can be easily adapted . Require : Manifold dimension d , Learning rate α , penalty parameter λ , inflation variance σ2 , batch size n , number of epochs E. Initialize : Parameters ψ and φ for flows fψ and hφ . while θ = ( ψ , φ ) has not converged do for e = 1 to E do for i = 1 to n do Sample : xi ∼ p ( x ) # sample data Inflate : x̃i = xi + εi , where εi ∼ N ( 0 , σ2ID ) # add noise ( ui , vi ) ← f−1ψ ( x̃i ) # project on u ∈ Rd and v ∈ RD−d x̂i ← fψ ( ui , 0 ) # reconstruct x u′i ← h−1φ ( ui ) # transform u end for L′DNF ← 1n ∑n i log pu′ ( u ′ i ) − log |det Jhφ ( u′i ) |+ log ( pv ( vi ) ) − log |det Jfψ ( ui , vi ) | +λ||xi − x̂i||2 # calculate logqθ ( x̃ ) and add reconstruction error θ ← θ − α∇θL′DNF # update model parameters end for end while Proposition 1 LetM be a d−dimensional manifold embedded in RD through g : Rd →M . Let x ∼ p ( x ) be generated by g ( u ) , where u ∼ N ( u ; 0 , Id ) , i.e . x = g ( u ) . Assume that we can learn the inflated distribution qσ ( x̃ ) using an NF and that for x̃ = x it holds that qσ ( x ) = p ( x ) qσ ( x|x ) . If LDNF ( θ ) = 0 , then |detGg ( x ) | = |detGfψ ( x ) ||detGhφ ( x ) | . ( 15 )

Summary: Authors propose a novel normalizing-flow-based method for simultaneous manifold and probability-density-on-the-manifold learning. Similarly to previous methods addressing this problem, the dimensionality of the manifold learned needs to be set by the user before training and is not learned as part of the process, albeit there exist other methods for estimating it. The advantage of the proposed method is the single-step training procedure, in contrast to the multi-step training procedure of the other competitive method in this domain. I consider it a good contribution towards building a normalizing flow based representation learning toolkit.