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: This paper is motivated by the desire to circumvent a troublesome assumption in Normalizing Flows (NFs) on the function $f$ that produces the pushforward measure --- that $f$ must be strictly invertible. This is difficult for generative models, since the dimension of the latent space $d$ is usually much smaller than that of the target data space $D$, i.e. d << D.  To address this, the authors train a transformation $f:\mathbb{R}^D \rightarrow \mathbb{R}^D$ as usually done in a NF approach, but separate the latent variable $z = f^{-1}(x)$ into two orthogonal components $u, v$, consisting of the first $d$ and remaining $D-d$ coordinates of $z$. Intuitively, $u$ should contain the "manifold" information of $x$, and $v$ should capture any "off-manifold" noise. This is encouraged by adding a reconstruction loss to the $u$ component of the latent space. To allow for greater complexity of the latent variable $u$, a second NF transformation $h$ is applied to u, i.e. $u = h(\tilde{u})$. In addition, the authors add noise robustness by adding Gaussian noise to $x$, i.e. $\tilde(x) = x + \epsilon$, and reconstruct $x$ --- much like with Denoising Autoencoders, hence the name.