On the Latent Space of Flow-based Models

1 INTRODUCTION . Normalizing flows ( Rezende and Mohamed , 2015 ; Kobyzev et al. , 2020 ) have shown considerable potential for the tasks of modelling and inferring expressive distributions through the learning of well-specified probabilistic models . Contemporary flow-based approaches define a latent space with dimensionality identical to the data space , typically by parameterizing a complex model pX ( x|θ ) using an invertible neural network fθ . Samples drawn from an initial , simple distribution pZ ( z ) ( e.g . Gaussian ) can be mapped to a complex distribution as x = fθ ( z ) . The process results in a tractable density that inhabits the full data space . However , contemporary flow models may make for an inappropriate choice to represent data that resides in a lower-dimensional manifold and thus does not populate the full ambient space . In such cases , the estimated model will necessarily have mass lying off the data manifold , which may result in under-fitting and poor generation qualities . Furthermore , principal objectives such as Maximum Likelihood Estimation ( MLE ) and Kullback-Leibler ( KL ) divergence minimization are ill-defined , bringing additional challenges for model training . In this work , we propose a principled strategy to model a data distribution that lies on a continuous manifold and we additionally identify the intrinsic dimension of the data manifold . Specifically , by using the connection between MLE and KL divergence minimization in Z space , we can address the important problem of ill-defined KL divergence under typical flow based assumptions . Flow models are based on the idea of “ change of variable ” . Assume a random variable Z with distribution PZ and probability density pZ ( z ) . We can transform Z to get a random variable X : X = f ( Z ) , where f : RD → RD is an invertible function with inverse f−1 = g. Suppose X has distribution PX and density function pX ( x ) , then log pX ( x ) will have the following form log pX ( x ) = log pZ ( g ( x ) ) + log ∣∣∣∣det ( ∂g∂x ) ∣∣∣∣ , ( 1 ) where log ∣∣∣det ( ∂g∂x ) ∣∣∣ is the log determinant of the Jacobian matrix . We call f ( or g ) a volumepreserving function if the log determinant is equal to 0 . Training of flow models typically makes use of MLE . We denote Xd as the random variable of the data with distribution Pd and density pd ( x ) . In addition to the well-known connection between MLE and minimization of the KL divergence KL ( pd ( x ) ||pX ( x ) ) in X space ( see Appendix A for detail ) , MLE is also ( approximately ) equivalent to minimizing the KL divergence in Z space , this is due to the KL divergence is invariant under invertible transformation ( Yeung , 2008 ; Papamakarios et al. , 2019 ) . Specifically , we define ZQ : ZQ = g ( Xd ) with distribution QZ and density function q ( z ) , the KL divergence in Z space can be written as KL ( q ( z ) ||p ( z ) ) = ∫ q ( z ) log q ( z ) dz − ∫ q ( z ) log p ( z ) dz ( 2 ) = − ∫ pd ( x ) ( log pZ ( g ( x ) ) + log ∣∣∣∣det ( ∂g∂x ) ∣∣∣∣ ) dx+ const. , ( 3 ) The full derivation can be found in Appendix A . Since we can only access samples x1 , x2 , . . . , xN from pd ( x ) , we approximate the integral by Monte Carlo sampling KL ( q ( z ) ||p ( z ) ) ≈ − 1 N N∑ n=1 log pX ( xn ) + const .. ( 4 ) We thus highlight the connection between MLE and KL divergence minimization , in Z space , for flow based models . The prior distribution p ( z ) is usually chosen to be a D-dimensional Gaussian distribution . However , if the data distribution Pd is singular , for example a measure on a low dimensional manifold , the induced latent distribution QZ is also singular . In this case , the KL divergence in equation 2 is typically not well-defined under the considered flow based model assumptions . This issue brings both theoretical and practical challenges that we will discuss in the following section . 2 FLOW MODELS WITH MANIFOLD DATA . We assume a data sample x ∼ Pd to be a D dimensional vector x ∈ RD and define the ambient dimensionality of Pd , denoted by Amdim ( Pd ) , to be D. However for many datasets of interest , e.g . natural images , the data distribution Pd is commonly believed to be supported on a lower dimensional manifold ( Beymer and Poggio , 1996 ) . We assume the dimensionality of the manifold to beK where K < D , and define the intrinsic dimension of Pd , denoted by Indim ( Pd ) , to be the dimension of this manifold . Figure 1a provides an example of this setting where Pd is a 1D distribution in 2D space . Specifically , each data sample x ∼ Pd is a 2D vector x = { x1 , x2 } where x1 ∼ N ( 0 , 1 ) and x2 = sin ( 2x1 ) . Therefore , this example results in Amdim ( Pd ) = 2 and Indim ( Pd ) = 1 . In flow-based models , function f is constructed such that it is both bijective and differentiable . When the prior PZ is a distribution whose support is RD ( e.g . Multivariate Gaussian distribution ) , the marginal distribution PX will also have support RD and Amdim ( PX ) = Indim ( PX ) = D. When the support of the data distribution lies on aK-dimensional manifold andK < D , Pd and PX are constrained to have different support . That is , the intrinsic dimensions of PX and Pd are always different ; Indim ( PX ) 6= Indim ( Pd ) . In this case it is impossible to learn a model distribution PX identical to the data distribution Pd . Nevertheless , flow-based models have shown strong empirical success in real-world problem domains such as the ability to generate high quality and realistic images ( Kingma and Dhariwal , 2018 ) . Towards investigating the cause and explaining this disparity between theory and practice , we employ a toy example to provide intuition for the effects and consequences resulting from model and data distributions that possess differing intrinsic dimension . Consider the toy dataset introduced previously ; a 1D distribution lying in a 2D space ( Figure 1a ) . The prior density p ( z ) is a standard 2D Gaussian p ( z ) = N ( 0 , IZ ) and the function f is a nonvolume preserving flow with two coupling layers ( see Appendix C.1 ) . In Figure 1b we plot samples from the flow model ; the sample x is generated by first sampling a 2D datapoint z ∼ N ( 0 , IZ ) and then letting x = f ( z ) . Figure 1c shows samples from the prior distributions PZ and QZ . QZ is defined as the transformation of Pd using the bijective function g , such that QZ is constrained to support a 1D manifold in 2D space , and Indim ( QZ ) = Indim ( Pd ) = 1 . Training of QZ to match PZ ( which has intrinsic dimension 2 ) , can be seen in Figure 1c to result in curling up of the manifold in the latent space , contorting it towards satisfying a distribution that has intrinsic dimension 2 . This ill-behaved phenomenon causes several potential problems for contemporary flow models : 1 . Poor sample quality . Figure 1b shows examples where incorrect assumptions in turn result in the model generating bad samples . 2 . Low quality data representations . The discussed characteristic that results in “ curling up ” of the latent space may cause degradations of the representation quality . 3 . Inefficient use of network capacity . Neural network capacity is spent on contorting the distribution QZ to satisfy imposed dimensionality constraints . A natural solution to the problem of intrinsic dimension mismatch is to select a prior distribution PZ with the same dimensionality as the intrinsic dimension of the data distribution such that : Indim ( PZ ) = Indim ( Pd ) . However , since we do not know Indim ( Pd ) explicitly , one option involves to instead learn it from the data distribution . In the following section , we will introduce a parameterization approach that enables us to learn Indim ( Pd ) . 3 LEARNING A MANIFOLD PRIOR . Consider a data vector x ∈ RD , then a flow based model prior PZ is usually given by a Ddimensional Gaussian distribution or alternative simple distribution that is also absolutely continuous ( a.c. ) in RD . Therefore , the intrinsic dimension Indim ( PZ ) = D. To allow a prior to have intrinsic dimension strictly less than D , we let PZ have the ‘ generalized density ’ 1 p ( z ) with the form p ( z ) = N ( 0 , AAT ) , ( 5 ) where z ∈ RD and A is a D×D lower triangular matrix with D ( D+1 ) 2 parameters , such that AA T is constrained to be a positive semi-definite matrix . When AAT has full rank D , then PZ is a ( nondegenerate ) multivariate Gaussian on RD . When Rank ( AAT ) = K and K < D , then PZ will degenerate to a Gaussian supported on a K-dimensional manifold , such that the intrinsic dimension Indim ( PZ ) = K. Figure 2 illustrates a sketch of this scenario . In practice , we initialize A to be an identity matrix , so AAT will also be an identity matrix and PZ is initialized as a standard Gaussian . When Rank ( AAT ) < D , the degenerate covariance AAT is no longer invertible and we are unable to evaluate the density value of p ( z ) for a given random vector z . Furthermore , when the data distribution Pd is supported on a K-dimensional manifold , QZ will also be supported on a Kdimensional manifold and no longer has a valid density function . Using equation 2 to train the flow model then becomes impossible as the KL divergence between PZ and QZ is not well defined2 . Recent work by Zhang et al . ( 2020 ) proposed a new family of divergence to address this problem . In the following section we briefly review the key concepts pertaining to this family of divergences . 1We use the generalized density to include the case that AAT is not full rank . 2The KL divergence KL ( Q||P ) is well defined when Q and P have valid densities and their densities have the same support ( Ali and Silvey , 1966 ) . 4 FIX THE ILL-DEFINED KL DIVERGENCE . random Let ZQ and ZP be two random variables with distribution QZ and PZ . The KL divergence between QZ and PZ is not defined if QZ or PZ does not have valid density function . Let K be an a.c. random variable that is independent of ZQ and ZP and has density pK , We define ZP̃ = ZP +K ; ZQ̃ = ZQ +K with distributions P̃Z and Q̃Z respectively . Then P̃Z and Q̃Z are a.c. with density functions q ( z̃ ) = ∫ z pK ( z̃− z ) dQZ p ( z̃ ) = ∫ z pK ( z̃− z ) dPZ . ( 6 ) We can thus define the spread KL divergence between QZ and PZ as the KL divergence between Q̃Z and P̃Z as : K̃L ( QZ ||PZ ) ≡ KL ( Q̃Z ||P̃Z ) ≡ KL ( q ( z̃ ) ||p ( z̃ ) ) . ( 7 ) In this work we letK be a Gaussian with diagonal covariance σ2ZI to satisfy the sufficient conditions such that K̃L is a valid divergence ( see Zhang et al . ( 2020 ) for details ) and has the properties : K̃L ( QZ ||PZ ) ≥ 0 , K̃L ( QZ ||PZ ) = 0⇔ QZ = PZ . ( 8 ) Since QZ and PZ are transformed from Pd and PX using an invertible function g , we have QZ = PZ ⇔ Pd = PX . ( 9 ) Therefore , the spread KL divergence can be used to train flow based models with a manifold prior in order to fit a dataset that lies on a lower-dimensional manifold .

Summary: This paper proposes a new method of training flow models, instead of minimizing KL divergence in the data space X, the paper proposes to minimize the KL divergence in the latent space Z. However, the problem of this is dimension mismatch, so the KL divergence is ill-defined. The proposed solution is to add noise, such that the KL divergence can be well defined again. In addition, this paper proposes learning a low dimensional prior p(Z) by a Gaussian distribution with low rank covariance matrix.