AR-ELBO: Preventing Posterior Collapse Induced by Oversmoothing in Gaussian VAE

1 INTRODUCTION . The variational autoencoder ( VAE ) framework ( Kingma & Welling , 2014 ; Higgins et al. , 2017 ; Zhao et al. , 2019 ) is a popular approach to achieve generative modeling in the field of machine learning . In this framework , a model that approximates the true posterior of observation data , is learned by a joint training of encoder and decoder , which creates a stochastic mapping between the observation data and the learned deep latent space . The latent space is assumed to follow a prior distribution . The generation of a new data sample can be done by sampling the latent space and passing the sample through the decoder . It is common to assume that both the prior on the latent space and the posterior of the observation data follow a Gaussian distribution . This setup is also known as the Gaussian VAE . In this case , the variance of the decoder output is usually modeled as an isotropic matrix σ2xI with a scalar parameter σ2x ≥ 0 . Furthermore , in order to deal with the intractable log-likelihood of the true posterior , the evidence lower bound ( ELBO ) ( Jordan et al. , 1999 ) is adopted as the objective function instead . While VAE-based generative models are usually considered to be more stable and easier to train than generative adversarial networks ( Goodfellow et al. , 2014 ) , they often suffer from the problem of posterior collapse ( Bowman et al. , 2015 ; Sønderby et al. , 2016 ; Alemi et al. , 2017 ; Xu & Durrett , 2018 ; He et al. , 2019 ; Razavi et al. , 2019a ; Ma et al. , 2019 ) , in which the latent space has little information of the input data . The phenomenon is generally mentioned as “ the posterior collapses to the prior in the latent space ” ( Razavi et al. , 2019a ) . Recently , several works have suggested that the variance parameter σ2x in the ELBO is strongly related to posterior collapse . For example , Lucas et al . ( 2019 ) analyzed posterior collapse through the analysis on a linear VAE . It revealed that an inappropriate choice of σ2x will introduce sub-optimal local optima and cause posterior collapse . Moreover , Lucas et al . ( 2019 ) reveals that contrary to the popular belief , these local optima are not introduced by replacing the log-likelihood with the ELBO , but by an excessively large σ2x . On the other hand , it can be shown that fixing σ2x to an excessively small value leads to under-regularization of the decoder , which can cause overfitting . However , in most implementations of a Gaussian VAE , the variance parameter σ2x is a fixed constant regardless of the input data and is usually 1.0 . In another work , Dai & Wipf ( 2019 ) proposed a two-stage VAE and treated σ2x as a training parameter . Besides the inappropriate choice of the variance parameter , posterior collapse can also induced by other causes . For example , Dai et al . ( 2020 ) found that , small nonlinear perturbation introduced in the network architecture can also result into extra sub-optimal local minima . However , in this work we will keep our focus on the variance parameter . We suggest that σ2x affects the strength of regulation over the gradient magnitude of the decoder . We call the expected gradient magnitude the smoothness throughout this paper . The smaller the gradient magnitude , the smoother the model . In particular , we would like to focus on the local smoothness of the model , which is the smoothness evaluated within the neighborhood of the encoded latent variable of the observation data . Thus , we begin with the following hypothesis : Main Hypothesis . The value of σ2x controls the regularization strength of the smoothness of the decoder . Therefore , an excessively large σ2x causes oversmoothness , which results in posterior collapse . Following the hypothesis , the estimation of σ2x should be related to properties of the approximated posterior of the latent space , such as its local smoothness . We will start with analyzing how σ2x regularizes the local smoothness of the stochastic decoder and then propose new objective functions that inherently determine σ2x via maximum likelihood estimation ( MLE ) . This proposed objective function is named AR-ELBO ( Adaptively Regularized ELBO ) , which controls the regularization strength via σ2x . Furthermore , several variations are derived for different parameterizations of variance parameters . Our main contributions are listed as follows : 1 . We show that our main hypothesis holds for linear approximated ELBO and empirically holds in the general case in Section 3 . This also suggests that the variance parameter σ2x should be estimated from properties of the approximated posterior instead of being treated as a hyperparameter . 2 . We propose the AR-ELBO , an ELBO-based objective function that adaptively regularizes the smoothness of the decoder by MLE of the variance parameter σ2x in Section 4 . Variations of AR-ELBO for several variance parameterizations of posterior distributions are also derived . AR-ELBO prevents the model from the posterior collapse induced by oversmoothing and improves the quality of generation , which is shown in Section 5 . The organization of this paper is as follows . In Section 2 , we propose a mathematical definition of posterior collapse in the form of mutual information . This also includes the conventional definition : “ the posterior collapses to the prior in the latent space ” . In Section 3 , the theoretical analysis and empirical support of the main hypothesis are given . We perform an analysis showing that σ2x affects the smoothness of the decoder via variance parameters of the latent space learned by the encoder . In Section 4 , we propose new AR-ELBO objective functions for various variance parameterizations of posterior distributions , which can relieve the decoder from being oversmoothed in the training and prevent posterior collapse . These objective functions no longer include any hyperparameters and can adaptively estimate the variance parameter σ2x from the observation data . It should be noted that if we adaptively determine σ2x with the proposed AR-ELBO , the strength of regularization of the decoder smoothness will gradually decrease as training progresses . In Section 5 , we conduct an experiment on the MNIST and CelebA datasets , which shows that utilizing the proposed AR-ELBO with the standard Gaussian VAE can be competitive with many other variations of VAE models in most situations . Throughout this paper , we use a , a and A for a scalar , a column vector and a matrix , and ln and log denote the natural logarithm and common logarithm , respectively . Our code is available from the following URL1 . 2 POSTERIOR COLLAPSE IN GAUSSIAN VAE . We begin with the standard formulation of the Gaussian VAE , which is the foundation of our research . A definition of posterior collapse is proposed by using mutual information ( MI ) . 1URL hidden due to blind review . 2.1 GAUSSIAN VAE . Consider a data space X ⊂ Rdx and a sample set { xi } Ni=1 ⊂ X , where xi ∼ pdata ( x ) . The empirical distribution p̃data ( x ) on X can be evaluated by p̃data ( x ) = 1N ∑N n=1 δ ( x − xn ) , where δ ( · ) denotes the Dirac delta function . In the standard VAE framework , a latent space Z ⊂ Rdz is learned and the sampled latent variables z ∈ Z are used to generate data samples x′ ∈ X . Let qφ ( z|x ) and pθ ( x|z ) denote the stochastic encoder and decoder , respectively . Trainable parameters of the two neural networks are denoted as φ and θ . The decoder generates data samples by pθ ( x ) : = Ep ( z ) [ pθ ( x|z ) ] , where p ( z ) is the prior distribution onZ . The encoder and decoder are jointly trained by minimizing the following objective function : L = −Epdata ( x ) [ ln pθ ( x ) ] + Epdata ( x ) DKL ( qφ ( z|x ) ‖ pθ ( z|x ) ) + Epdata ( x ) [ ln pdata ( x ) ] = DKL ( pdata ( x ) ‖ pθ ( x ) ) + Epdata ( x ) DKL ( qφ ( z|x ) ‖ pθ ( z|x ) ) . ( 1 ) This objective function was derived in Zhao et al . ( 2019 ) , which represents everything in the form of Kullback–Leibler divergence , and is equivalent to ELBO maximization up to an additive constant . In the context of the Gaussian VAE , the encoder and decoder are assumed to satisfy qφ ( z|x ) = N ( z|µφ ( x ) , diag ( σ2φ ( x ) ) ) and pθ ( x|z ) = N ( x|µθ ( z ) , σ2xI ) . ( 2 ) The prior p ( z ) is also assumed to be the Gaussian distribution as p ( z ) = N ( z|0 , I ) . Substituting ( 2 ) into ( 1 ) while omitting terms independent of θ and φ leads to the following objective : J̃σ2x ( θ , φ ) = Ep̃data ( x ) [ 1 2σ2x Eqφ ( z|x ) [ ‖x− µθ ( z ) ‖ 2 2 ] +DKL ( qφ ( z|x ) ‖ p ( z ) ) ] , ( 3 ) which can be interpreted as the sum of the expected values of the reconstruction loss and the regularization term . In the case of a Gaussian prior and posterior , the regularization term is equal to 1 2 ∑dz i=1 ( σ 2 φ , i ( x ) + µφ , i ( x ) 2 − log σ2φ , i ( x ) − 1 ) . 2.2 POSTERIOR COLLAPSE . Posterior collapse is a major problem , where the encoder learns to map inputs to the latent space while ignoring the data distribution . In this phenomenon , the MI between input data and reconstructed data through the encoder-decoder path is reduced because the latent space has less information about the data distribution . Here , we suggest the following definition of posterior collapse . Definition 1 . Posterior collapse is defined as the MI I ( x ; x′ ) becoming nearly zero , where x′ : = µθ ( z ) with z ∼ qφ ( z|x ) . In many works ( Bowman et al. , 2015 ; Sønderby et al. , 2016 ; Alemi et al. , 2017 ; He et al. , 2019 ; Razavi et al. , 2019a ) , the phenomenon denoted as posterior collapse has been mathematically represented as Epdata ( x ) DKL ( qφ ( z|x ) ‖ p ( z ) ) → 0 , which we hereafter refer to as KL collapse ( Xu & Durrett , 2018 ) . However , posterior collapse is not always caused by the diminished KL divergence , i.e. , posterior collapse with Epdata ( x ) DKL ( qφ ( z|x ) ‖ p ( z ) ) 6= 0 can occur . The proposed definition of posterior collapse includes KL collapse from the following theorem , which is proven in Appendix A. Theorem 2 . I ( x ; x′ ) → 0 as Epdata ( x ) DKL ( qφ ( z|x ) ‖ p ( z ) ) → 0 holds for any pθ ( x|z ) . In Appendix E , it is demonstrated that posterior collapse can happen even if the KL divergence is nonzero when the posterior variance is fixed in Z . 3 VARIANCE PARAMETER σ2x AND THE LOCAL SMOOTHNESS In this section , we provide mathematical and empirical support of the main hypothesis . Throughout this section , we adopt the following parameterization for the encoder for simplicity : qφ , σ2z ( z|x ) = N ( z|µφ ( x ) , σ2zI ) , where the variance is parameterized as an isotropic matrix unlike the conventional VAE . A similar analysis on the conventional VAE can be found in Appendix B . It begins with showing that the choice of σ2x affects the convergence point of σ 2 z , which is the variance parameter of the latent space . Then , we show that σ2z acts as the weight of the gradient penalty , which is implicitly included in ( 3 ) . This supports the main hypothesis that the over-regulation imposed by a large σ2x via σ 2 z causes the oversmoothness of the decoder and leads to posterior collapse . It is also empirically supported by observing the tendencies of the convergence point of σ2z , the smoothness and the MI I ( x , x′ ) . Ultimately , these items of evidence motivated us to develop a method that adapts σ2x to prevent oversmoothing of the decoder . 3.1 REGULARIZATION EFFECT OF σ2x IN LINEAR APPROXIMATED ELBO The effect of σ2x on the convergence point of the variance parameter σ 2 z can be observed from two extreme cases , σ2x → 0+ and σ2x → ∞ . In the first case , J̃σ2x reduces to Ep̃data ( x ) Eqφ , σ2z ( z|x ) [ ‖x − µθ ( z ) ‖ 2 2 ] , and it becomes zero only if σ 2 z = 0 . In the second case , J̃σ2x reduces to Ep̃data ( x ) DKL ( qφ , σ2z ( z|x ) ‖ p ( z ) ) , and σ 2 z becomes 1 at the minimum point , from DKL ( qφ , σ2z ( z|x ) ‖ p ( z ) ) = dz 2 ( σ 2 z − log σ2z − 1 ) + ‖µφ ( x ) ‖22 . This shows that a small σ2x makes σ2z converge to a value around 0 , while a large σ2x makes σ 2 z converge to a value around 1 . If σ2z is sufficiently small , as the training progresses to a certain extent , the perturbed decoding process µθ ( z + z ) around z = µφ ( x ) with z ∼ N ( z|0 , σ2zI ) can be approximated as a linear function . The ELBO can be approximated as follows by using the linear approximation of µθ ( · ) and omitting terms independent of θ and φ : J̃σ2x ( θ , φ , σ 2 z ) ≈ 1 2σ2x Ep̃data ( x ) [ ‖x−µθ ( µφ ( x ) ) ‖22 +σ 2 z ‖∇µθ ( µφ ( x ) ) ‖ 2 F +2σ 2 x ‖µφ ( x ) ‖ 2 2 ] . ( 4 ) In the approximation above , ‖ · ‖F is the Frobenius norm and σ2z is treated as a function parameter . The derivation of the above approximation can be found in Appendix B . Equation ( 4 ) decomposes the objective function into three terms : a reconstruction error term , gradient penalty term and L2 regularization term . As one can see from ( 4 ) , σ2z regularizes the smoothness of the decoder by penalizing its gradient norm in training . Although the linear approximation above is derived for the simplified VAE parameterization , we also provide the linear approximation of the ELBO for the standard VAE parameterization ( 2 ) in Appendix B , where the second term in ( 4 ) becomes a weighted gradient penalty . Summarizing the above observations shows that σ2x affects the smoothness via σ 2 z , while σ 2 z directly regularizes the smoothness . This means that if σ2x is excessively large , it will cause overregularization of the decoder and suppress I ( z , x′ ) ( ≥ I ( x , x′ ) ) , which finally leads to posterior collapse . To avoid such over-regularization , σ2x and σ 2 z should be determined appropriately . In addition , an experiment shows that posterior collapse can be triggered by directly manipulating σ2z , as discussed in Appendix E .

Summary: The paper considers Gaussian VAEs and their tendency to suffer from posterior collapse. In particular, the authors analyse the impact of the usually fixed covariance $\sigma_x$ of the decoder Gaussian on the learned encoder variance. They show that the former can be seen as a regulariser for the latter and therefore impacts the "smoothness" of the encoder. The authors hypothesize that a large value of $\sigma_x$ causes posterior collapse as a consequence.