Quantitative Understanding of VAE as a Non-linearly Scaled Isometric Embedding

1 INTRODUCTION . Variational autoencoder ( VAE ) ( Kingma & Welling , 2014 ) is one of the most successful generative models , estimating posterior parameters of latent variables for each input data . In VAE , the latent representation is obtained by maximizing an evidence lower bound ( ELBO ) . A number of studies ( Higgins et al. , 2017 ; Kim & Mnih , 2018 ; Lopez et al. , 2018 ; Chen et al. , 2018 ; Locatello et al. , 2019 ; Alemi et al. , 2018 ; Rolı́nek et al. , 2019 ) have tried to reveal the property of latent variables . However , quantitative behavior of VAE is still not well clarified . For example , there has not been a theoretical formulation of the reconstruction loss and KL divergence in ELBO after optimization . More specifically , although the conditional distribution pθ ( x|z ) in the reconstruction loss of ELBO is predetermined such as the Gaussian or Bernoulli distributions , it has not been discussed well whether the true conditional distribution after optimization matches the predetermined distribution . Rate-distortion ( RD ) theory ( Berger , 1971 ) , which is an important part of Shannon information theory and successfully applied to image compression , quantitatively formulates the RD trade-off optimum in lossy compression . To realize a quantitative data analysis , Rate-distortion ( RD ) theory based autoencoder , RaDOGAGA ( Kato et al. , 2020 ) , has been proposed with isometric embedding ( Han & Hong , 2006 ) where the distance between arbitrary two points of input space in a given metrics is always the same as L2 distance in the isometric embedding space . In this paper , by mapping VAE latent space to an implicit isometric space like RaDOGAGA on variable-by-variable basis and analysing VAE quantitatively as a well-examined lossy compression , we thoroughly clarify the quantitative properties of VAE theoretically and experimentally as follows . 1 ) Implicit isometric embedding is derived in the loss metric defined space such that the entropy of data representation becomes minimum . A scaling factor between the VAE latent space and implicit isometric space is formulated by the posterior for each input . In the case of β-VAE , the posterior variance of each dimensional component in the implicit isometric embedding space is a constant β/2 , which is analogous to the rate-distortion optimal of transform coding in RD theory . As a result , the reconstruction loss and KL divergence in ELBO can be quantitatively formulated . 2 ) From these properties , VAE can provide a practical quantitative analysis of input data . First , the data probabilities in the input space can be estimated from the prior , loss metric , and posterior parameters . In addition , the quantitative importance of each latent variable , analogous to the eigenvalue of PCA , can be evaluated from the posterior variance of VAE . This work will lead the information theoretic generative models in the right direction . 2 RELATED WORKS . 2.1 VARIATIONAL AUTOENCODER AND THEORETICAL ANALYSIS . In VAE , ELBO is maximized instead of maximizing the log-likelihood directly . Let x ∈ Rm be a point in a dataset . The original VAE model consists of a latent variable with fixed prior z ∼ p ( z ) = N ( z ; 0 , In ) ∈ Rn , a parametric encoder Encφ : x ⇒ z , and a parametric decoder Decθ : z ⇒ x̂ . In the encoder , qφ ( z|x ) = N ( z ; µ ( x ) , σ ( x ) ) is provided by estimating parameters µ ( x ) and σ ( x ) . Let Lx be a local cost at data x . Then , ELBO is described by ELBO = Ex∼p ( x ) [ Ez∼qφ ( z|x ) [ log pθ ( x|z ) ] −DKL ( qφ ( z|x ) ‖p ( z ) ) ] . ( 1 ) In Ex∼p ( x ) [ · ] , the first term Ez∼qφ ( z|x ) [ · ] is called the reconstruction loss . The second term DKL ( · ) is a Kullback–Leibler ( KL ) divergence . Let µj ( x ) , σj ( x ) , and DKLj ( x ) be j-th dimensional values of µ ( x ) , σ ( x ) , and KL divergence . Then DKL ( · ) is derived as : DKL ( · ) = n∑ j=1 DKLj ( x ) , where DKLj ( x ) = 1 2 ( µj ( x ) 2 + σj ( x ) 2 − log σj ( x ) 2 − 1 ) . ( 2 ) D ( x , x̂ ) denotes a metric such as sum square error ( SSE ) and binary cross-entropy ( BCE ) as loglikelihoods of Gaussian and Bernoulli distributions , respectively . In training VAE , the next objective is used instead of Eq . 1 , where β is a parameter to control the trade-off ( Higgins et al. , 2017 ) . Lx = Ez∼qφ ( z|x ) [ D ( x , x̂ ) ] + βDKL ( · ) . ( 3 ) However , it has not been fully discussed whether the true conditional distribution matches the predetermined distribution , or how the value of KL divergence is derived after training . There have been several studies to analyse VAE theoretically . Alemi et al . ( 2018 ) introduced the RD trade-off based on the information-theoretic framework to analyse β-VAE . However , they did not clarify the quantitative property after optimization . Dai et al . ( 2018 ) showed that VAE restricted as a linear transform can be considered as a robust PCA . However , their model has a limitation for the analysis on each latent variable basis because of the linearity assumption . Rolı́nek et al . ( 2019 ) showed that the Jacobian matrix of VAE at each latent variable is orthogonal , which makes latent variables disentangled implicitly . However , they do not uncover the orthonormality and quantitative properties because they simplify KL divergence as a constant . Dai & Wipf ( 2019 ) also showed that the expected rate of VAE for the r-dimensional manifold is close to − ( r/2 ) log γ +O ( 1 ) at γ → 0 when pθ ( x̂|x ) = N ( x̂ ; x , γIm ) holds . The remaining challenge is to clearly figure out what latent space is obtained at a given dataset , a loss metric , and β in the model . 2.2 RATE-DISTORTION THEORY , TRANSFORM CODING , AND ISOMETRIC EMBEDDING . RD theory ( Berger , 1971 ) formulated the optimal transform coding ( Goyal , 2001 ) for the Gaussian source with square error metric as follows . Let x ∈ Rm be a point in a dataset . First , the data are transformed deterministically with the orthonormal transform ( orthogonal and unit norm ) such as Karhunen-Loève transform ( KLT ) ( Rao & Yip , 2000 ) . Let z ∈ Rm be a point transformed from x . Then , z is entropy-coded by allowing equivalent stochastic distortion ( or posterior with constant variance ) in each dimension . A lower bound of a rate R at a distortion D is denoted by R ( D ) . The derivation of R ( D ) is as follows . Let zj be the j-th dimensional component of z and σzj2 be the variance of zj in a dataset . It is noted that σzj2 is the equivalent to eigenvalues of PCA for the dataset . Let d be a distortion equally allowed in each dimensional channel . At the optimal condition , the distortion Dopt and rate Ropt on the curve R ( D ) is calculated as a function of d : Ropt = 1 2 m∑ j=1 max ( log ( σzj 2/d ) , 0 ) , Dopt = m∑ j=1 min ( d , σzj 2 ) . ( 4 ) The simplest way to allow equivalent distortion is to use a uniform quantization ( Goyal , 2001 ) . Let T be a quantization step , and round ( · ) be a round function . Quantized value ẑj is derived as kT , where k = round ( zj/T ) . Then , d is approximated by T 2/12 as explained in Appendix H.1 . To practically achieve the best RD trade-off in image compression , rate-distortion optimization ( RDO ) has also been widely used ( Sullivan & Wiegand , 1998 ) . In RDO , the best trade-off is achieved by finding a encoding parameter that minimizes a cost L = D + λR at given Lagrange parameter λ . Recently , deep image compression ( Ballé et al. , 2018 ) has been proposed . In these works , instead of an orthonormal transform with sum square error ( SSE ) metric in the conventional lossy compression , a deep autoencoder is trained with flexible metrics , such as structural similarity ( SSIM ) ( Wang et al. , 2001 ) for RDO . Recently , an isometric autoencoder , RaDOGAGA ( Kato et al. , 2020 ) was proposed based on Ballé et al . ( 2018 ) . They proved that the latent space to be isometric to the input space if the model is trained by RDO using a parametric prior and posterior with constant variance . By contrast , VAE uses a fixed prior with a variable posterior . In section 3 , we explain that VAE can be quantitatively understood as the rate-distortion optimum as in Eq . 4 by mapping VAE latent space to implicit isometric embedding on a variable-to-variable basis as in Fig . 1 . 3 UNDERSTANDING OF VAE AS A SCALED ISOMETRIC EMBEDDING . This section shows the quantitative understanding of VAE . First , we present the hypothesis of mapping VAE latent space to an implicit isometric embedding space . Second , we reformulate the objective of β-VAE for easy analysis . Third , we prove the hypothesis from the minimum condition of the objective . Then , we show that ELBO can be interpreted as an optimized RDO cost of transform coding where the quantitative properties are well clarified , as well as discuss and correct some prior theoretical studies . Lastly , we explain the quantitative properties of VAE to validate the theory including approximations and provide a practical data analysis . 3.1 HYPOTHESIS OF MAPPING VAE TO THE IMPLICIT ORTHONORMAL TRANSFORM . Figure 1 shows the mapping of VAE to the implicit isometric embedding . Assume the data manifold is smooth and differentiable . Let Sinput ( ⊂ Rm ) be an input space of the dataset . D ( x , x́ ) denotes a metric for points x , x́ ∈ Sinput . Using the second order Taylor expansion , D ( x , x + δx ) can be approximated by tδx Gxδx , where Gx and δx are an x dependent positive definite Hermitian metric tensor and an arbitrary infinitesimal displacement in Sinput , respectively . The derivations of Gx for SSE , BCE , and SSIM are shown in Appendix H.2 . Next , an implicit isometric embedding space SIso ( ⊂ Rm ) is introduced like the isometric latent space in RaDOGAGA ( Kato et al. , 2020 ) , such that the entropy of data representation is minimum in the inner product space ofGx . Let y and yj be a point in SIso and its j-th component , respectively . Because of the isometricity , p ( x ) ' p ( y ) will hold . We will also show the posterior variance of each dimensional component yj is a constant β/2 . In addition , the variance of yj will show the importance like PCA when the data manifold has a disentangled feature by nature in the metric space ofGx and the prior covariance is diagonal . Then , SIso is nonlinearly scaled to the VAE ’ s anisometric orthogonal space SVAE ( ⊂ Rn ) on a variable-by-variable basis . Let z be a point in SVAE , and zj denotes the j-th component of z . Let p ( yj ) and p ( zj ) be the probability distribution of the j-th variable in SIso and SVAE . Each variable yj is nonlinearly scaled to zj , such that dzj/dyj = p ( yj ) /p ( zj ) to fit the cumulative distribution . dzj/dyj is σj ( x ) / √ β/2 , the ratio of posterior ’ s standard deviations for zj and yj , such that KL divergences in both spaces are equivalent . In addition , dimensional components whose KL divergences are zero can be discarded because such dimensions have no information . 3.2 REFORMULATION OF OBJECTIVE TO THE FORM USING ∂x/∂zj AND ∂x/∂z We reformulate the objective Lx to the form using ∂x/∂zj and ∂x/∂z . Here , the dimensions of x and z , i.e. , m and n , are set as the same . The condition to reduce n is shown in section 3.3 . Reformulation of D ( x , x̂ ) loss : In accordance with Kato et al . ( 2020 ) , the loss D ( x , x̂ ) can be decomposed into D ( x̆ , x̂ ) + D ( x , x̆ ) , where x̆ denotes Decθ ( µ ( x ) ) . The first term D ( x̆ , x̂ ) is a distortion between the decoded values of µ ( x ) with and without noise σ ( x ) . We call this term as a coding loss . This term is expanded as follows . δx̆ denotes x̂ − x̆ . Then , D ( x̆ , x̂ ) term can be approximated by tδx̆ Gxδx̆ . Let xzj be ∂x/∂zj at zj = µj ( x ) , and δzj ∼ N ( 0 , σj ( x ) ) be an added noise in zj . Then , δx̆ is approximated by δx̆ ' ∑m j=1 δzj xzj . Because δzj and δzk for j 6= k are uncorrelated , the average of D ( x̆ , x̂ ) over z ∼ qφ ( z|x ) can be finally reformulated by Ez∼qφ ( z|x ) [ D ( x̆ , x̂ ) ] ' Ez∼qφ ( z|x ) [ tδx̆ Gxδx̆ ] ' n∑ j=1 σj ( x ) 2 txzjGxxzj . ( 5 ) The second term D ( x , x̆ ) is a loss between the input data and Decθ ( µ ( x ) ) . We call this term a transform loss . We presume VAE is analogous to the Wiener filter ( Wiener , 1964 ; Jin et al. , 2003 ) where the coding loss is regarded as an added noise . From the Wiener filter theory , the ratio between the transform loss and coding loss is close to the ratio between the coding loss and the variance of the input data . The coding loss , approximately nβ/2 as in Eq . 14 , should be smaller than the variance of the input data to capture meaningful information . Thus the transform loss , usually small , is not considered in the following discussion . Appendix B explains the detail in a simple 1-dimensional VAE . We show the exhaustive and quantitative evaluation of coding loss and transform loss in the toy dataset in appendix E.2 to validate this approximation . Reformulation of KL divergence : When σj ( x ) 1 , σj ( x ) 2 − log σj ( x ) 2 is observed . For example , when σj ( x ) 2 < 0.1 , we have − ( σj ( x ) 2/ log σj ( x ) 2 ) < 0.05 . In such dimensions , DKLj ( x ) can be approximated as Eq . 6 by ignoring the σj ( x ) 2 term and setting p ( µj ( x ) ) to N ( zj ; 0 , 1 ) : DKLj ( x ) ' 1 2 ( µj ( x ) 2 − log σj ( x ) 2 − 1 ) = − log ( σj ( x ) p ( µj ( x ) ) ) − log 2πe 2 . ( 6 ) Eq . 6 can be considered as a rate of entropy coding for a symbol with mean µj ( x ) allowing quantization noise σj ( x ) 2 , as shown in Appendix H.3 . Thus , in the dimension with meaningful information , σj ( x ) 2 is much smaller than the prior variance 1 , and the approximation in Eq.6 is reasonable . Let p ( µ ( x ) ) be ∏n j=1 p ( µj ( x ) ) . p ( µ ( x ) ) = p ( x ) |det ( ∂x/∂z ) | holds where det ( ∂x/∂z ) is a Jacobian determinant at z = µ ( x ) . Let CDKL be a constant n 2 log 2πe . Then , DKL ( · ) is reformulated by DKL ( · ) ' − log ( p ( µ ( x ) ) n∏ j=1 σj ( x ) ) − CDKL ' − log ( p ( x ) ∣∣∣∣det ( ∂x∂z ) ∣∣∣∣ n∏ j=1 σj ( x ) ) − CDKL . ( 7 ) Final objective form : From Eqs . 5 and 7 , the objective L′x to minimise is derived as : L′x = n∑ j=1 σj ( x ) 2 txzjGxxzj − β log ( p ( x ) ∣∣∣∣det ( ∂x∂z ) ∣∣∣∣ n∏ j=1 σj ( x ) ) − CDKL . ( 8 )

Summary: The paper builds on a branch of recent works that consider and analyse Variational autoencoders (VAE) from the view point of data compression. This started from Alemi et al. 2018, where the authors consider and analyse the mutual information between data and latent codes and culminates in Kato et al. 2020, where the authors consider models in which both the encoder and the decoder are assumed as deterministic (isometric) mappings. The submitted paper aims at reconciling this type of models with standard VAEs by claiming to prove that VAEs can be obtained from the former by an non-linear component-wise scaling of the latent space. The authors claim among others that this approach allows to estimate the data probability density from the learned VAE model.