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: This paper analyses the "posterior collapse" phenomenon observed in training latent variable models (in particular Variational Auto-Encoders), and propose a new training objective to remedy the problem. The theoretical analysis of the authors suggests that the posterior collapse is induced by an inappropriate choice of variance in the decoder distribution. The new objective they propose, the AR-ELBO, jointly optimises this variance along with the usual network parameters. The authors demonstrate that their objective yields relatively good results on image modelling, compared to other standard VAE methods.