Simple and Effective VAE Training with Calibrated Decoders

1 INTRODUCTION . Deep density models based on the variational autoencoder ( VAE ) ( Kingma & Welling , 2014 ; Rezende et al. , 2014 ) have found ubiquitous use in probabilistic modeling and representation learning as they are both conceptually simple and are able to scale to very complex distributions and large datasets . These VAE techniques are used for tasks such as future frame prediction ( Castrejon et al. , 2019 ) , image segmentation ( Kohl et al. , 2018 ) , generating speech ( Chung et al. , 2015 ) and music ( Dhariwal et al. , 2020 ) , as well as model-based reinforcement learning ( Hafner et al. , 2019a ) . However , in practice , many of these approaches require careful manual tuning of the balance between two terms that correspond to distortion and rate from information theory ( Alemi et al. , 2017 ) . This balance trades off fidelity of reconstruction and quality of samples from the model : a model with low rate would not contain enough information to reconstruct the data , while allowing the model to have high rate might lead to unrealistic samples from the prior as the KL-divergence constraint becomes weaker ( Alemi et al. , 2017 ; Higgins et al. , 2017 ) . While a proper variational lower bound does not expose any free parameters to control this tradeoff , many prior works heuristically introduce a weight on the prior KL-divergence term , often denoted β . Usually , β needs to be tuned for every dataset and model variant as a hyperparameter , which slows down development and can lead to poor performance as finding the optimal value is often prohibitively computationally expensive . Moreover , using β 6= 1 precludes the appealing interpretation of the VAE objective as a bound on the data likelihood , and is undesirable for applications like density modeling . While many architectures for calibrating decoders have been proposed in the literature ( Kingma & Welling , 2014 ; Kingma et al. , 2016 ; Dai & Wipf , 2019 ) , more applied work typically employs VAEs with uncalibrated decoding distributions , such as Gaussian distributions without a learned variance , where the decoder only outputs the mean parameter ( Castrejon et al. , 2019 ; Denton & Fergus , 2018 ; Lee et al. , 2019 ; Babaeizadeh et al. , 2018 ; Lee et al. , 2018 ; Hafner et al. , 2019b ; Pong et al. , 2019 ; Zhu et al. , 2017 ; Pavlakos et al. , 2019 ) , or uses other ad-hoc modifications to the objective ( Sohn et al. , 2015 ; Henaff et al. , 2019 ) . Indeed , it is well known that attempting to learn the variance in a Gaussian decoder may lead to numerical instability ( Rezende & Viola , 2018 ; Dai & Wipf , 2019 ) , and naı̈ve approaches often lead to poor results . As a result , it remains unclear whether practical empirical performance of VAEs actually benefits from calibrated decoders or not . To rectify this , our first contribution is a comparative analysis of various calibrated decoder architectures and practical recommendations for simple and effective VAE training . We find that , while naı̈ve calibrated decoders often lead to worse results , a careful choice of the decoder distribution can work very well , and removes the need to tune the additional parameter β . Indeed , we note that the entropy of the decoding distribution controls the mutual information I ( x ; z ) . Calibrated decoders allow the model to control I ( x ; z ) automatically , instead of relying on manual tuning . Our second contribution is a simple but novel technique for optimizing the decoder variance analytically , without requiring the decoder network to produce it as an additional output . We call the resulting approach to learning the Gaussian variance the σ-VAE . In our experiments , the σ-VAE outperforms the alternative of learning the variance through gradient descent , while being simpler to implement and extend . We validate our results on several VAE and sequence VAE models and a range of image and video datasets . 2 RELATED WORK . Prior work on variational autoencoders has studied a number of different decoder parameterizations . Kingma & Welling ( 2014 ) ; Rezende et al . ( 2014 ) use the Bernoulli distribution for the binary MNIST data and Kingma & Welling ( 2014 ) use Gaussian distributions with learned variance parameter for grayscale images . However , modeling images with continuous distributions is prone to instability as the variance can converge to zero ( Rezende & Viola , 2018 ; Mattei & Frellsen , 2018 ; Dai & Wipf , 2019 ) . Some work has attempted to rectify this problem by using dequantization ( Gregor et al. , 2016 ) , which is theoretically appealing as it is tightly related to the log-likelihood of the original discrete data ( Theis et al. , 2016 ) , optimizing the variance in a two-stage procedure ( Arvanitidis et al. , 2017 ) , or training a post-hoc prior ( Ghosh et al. , 2019 ) . Takahashi et al . ( 2018 ) ; Barron ( 2019 ) proposed more expressive distributions . Additionally , different choices for representing such variance exist , including diagonal covariance ( Kingma & Welling , 2014 ; Sønderby et al. , 2016 ; Rolfe , 2016 ) , or a single shared parameter ( Kingma et al. , 2016 ; Dai & Wipf , 2019 ; Edwards & Storkey , 2016 ; Rezende & Viola , 2018 ) . We analyze these and notice that learning a single variance parameter shared across images leads to stable training and good performance , without the use of dequantization or even clipping the variance , although these techniques can be used with our decoders ; and further improve the estimation of this variance with an analytic solution . Early work on discrete VAE decoders for color images modeled them with the Bernoulli distribution , treating the color intensities as probabilities ( Gregor et al. , 2015 ) . Further work has explored various parameterizations based on discretized continuous distributions , such as discretized logistic ( Kingma et al. , 2016 ) . More recent work has improved expressivity of the decoder with a mixture of discretized logistics ( Chen et al. , 2016 ; Maaløe et al. , 2019 ) . However , these models also employ powerful autoregressive decoders ( Chen et al. , 2016 ; Gulrajani et al. , 2016 ; Maaløe et al. , 2019 ) , and the latent variables in these models may not represent all of the significant factors of variation in the data , as some factors can instead be modeled internally by the autoregressive decoder ( Alemi et al. , 2017 ) .1 While many calibrated decoders have been proposed , outside the core generative modeling community uncalibrated decoders are ubiquitous . They are used in work on video prediction ( Denton & Fergus , 2018 ; Castrejon et al. , 2019 ; Lee et al. , 2018 ; Babaeizadeh et al. , 2018 ) , image segmentation ( Kohl et al. , 2018 ) , image-to-image translation ( Zhu et al. , 2017 ) , 3D human pose ( Pavlakos et al. , 2019 ) , as well as model-based reinforcement learning ( Henaff et al. , 2019 ; Hafner et al. , 2019b ; a ) , and representation learning ( Lee et al. , 2019 ; Watter et al. , 2015 ; Pong et al. , 2019 ) . Most of these works utilize the heuristic hyperparameter β instead , which is undesirable both as the resulting objective is no longer a bound on the likelihood , and as β usually requires extensive tuning . In this work , we analyze the common pitfalls of using calibrated decoders that may have prevented the practitioners from using them , propose a simple and effective analytic way of learning such calibrated distribution , and provide a comprehensive experimental evaluation of different decoding distributions . Alternative discussions of the hyperparameter β are presented by Zhao et al . ( 2017 ) ; Higgins et al . ( 2017 ) ; Alemi et al . ( 2017 ) ; Achille & Soatto ( 2018 ) , who show that it controls the amount of information in the latent variable , I ( x ; z ) . Peng et al . ( 2018 ) ; Rezende & Viola ( 2018 ) further discuss constrained optimization objectives for VAEs , which also yield a similar hyperparameter . Here , we focus on β-VAEs with Gaussian decoders with constant variance , as commonly used in recent work , and show that the hyperparameter β can be incorporated in the decoding likelihood for these models . 1BIVA ( Maaløe et al. , 2019 ) uses the Mixture of Logistics decoder proposed in ( Salimans et al. , 2017 ) that produces the channels for each pixel autoregressively , see also App D . 3 ANALYSING DECODING DISTRIBUTIONS . The generative model of a VAE ( Kingma & Welling , 2014 ; Rezende et al. , 2014 ) with parameters θ is specified with a prior distribution over the latent variable pθ ( z ) , commonly unit Gaussian , and a decoding distribution pθ ( x|z ) , which for color images is commonly a conditional Gaussian parameterized with a neural network . We would like to fit this generative model to a given dataset by maximizing the evidence lower bound ( ELBO ( Neal & Hinton , 1998 ; Jordan et al. , 1999 ; Kingma & Welling , 2014 ; Rezende et al. , 2014 ) ) , which uses an approximate posterior distribution qφ ( z|x ) , also commonly a conditional Gaussian specified with a neural network . In this work , we focus on the form of the decoding distribution pθ ( x|z ) . To achieve the best results , we want a decoding distribution that represents the required probability p ( x|z ) accurately In this section , we will review and analyze various choices of decoding distributions that enable better decoder calibration , including expressive decoding distributions that can represent both the prediction of the image and the uncertainty about such prediction , or even multimodal predictions . 3.1 GAUSSIAN DECODERS . We first analyse the commonly used Gaussian decoders . We note that the commonly used MSE reconstruction loss between the reconstruction x̂ and ground truth data x is equivalent to the negative log-likelihood objective with a Gaussian decoding distribution with constant variance : − ln p ( x|z ) = 1 2 ||x̂− x||2 +D ln √ 2π = 1 2 ||x̂− x||2 + c = D 2 MSE ( x̂ , x ) + c , where p ( x|z ) ∼ N ( x̂ , I ) , the prediction x̂ is produced with a neural network x̂ = µθ ( z ) , and D is the dimensionality of x . This demonstrates a drawback of methods that rely simply on the MSE loss ( Castrejon et al. , 2019 ; Denton & Fergus , 2018 ; Lee et al. , 2019 ; Hafner et al. , 2019b ; Pong et al. , 2019 ; Zhu et al. , 2017 ; Henaff et al. , 2019 ) , as it is equivalent to assuming a particular , constant variance of the Gaussian decoding distribution . By learning this variance , we can achieve much better performance due to better calibration of the decoder . There are several ways in which we can specify this variance . An expressive way to specify the variance is to specify a diagonal covariance matrix for the image , with one value per pixel ( Kingma & Welling , 2014 ; Sønderby et al. , 2016 ; Rolfe , 2016 ) . This can be done , for example , by letting a neural network σθ output the diagonal entries of the covariance matrix given a latent sample z : pθ ( x|z ) ∼ N ( µθ ( z ) , σθ ( z ) 2 ) . ( 1 ) This parameterization of the decoding distribution outputs one variance value per each pixel and channel . While powerful , we observe in Section 5.3 that this approach attains suboptimal performance , and is moreover prone to numerical instability . Instead , we will find experimentally that a simpler parameterization , in which the covariance matrix is specified with a single shared ( Kingma et al. , 2016 ; Dai & Wipf , 2019 ; Edwards & Storkey , 2016 ; Rezende & Viola , 2018 ) parameter σ as Σ = σI often works better in practice : pθ , σ ( x|z ) ∼ N ( µθ ( z ) , σ 2I ) . ( 2 ) The parameter σ can be optimized together with parameters of the neural network θ with gradient descent . Of particular interest is the interpretation of this parameter . Writing out the expression for the decoding likelihood , we obtain − ln p ( x|z ) = 1 2σ2 ||x̂−x||2+D lnσ √ 2π = 1 2σ2 ||x̂−x||2+D lnσ+c = D lnσ+ D 2σ2 MSE ( x̂ , x ) +c . The full objective of the resulting Gaussian σ-VAE is : Lθ , φ , σ = D lnσ + D 2σ2 MSE ( x̂ , x ) +DKL ( q ( z|x ) ||p ( z ) ) . ( 3 ) Note that σ may be viewed as a weighting parameter between the MSE reconstruction term and the KL-divergence term in the objective . Moreover , this objective explicitly specifies how to select the optimal variance : the variance should be selected to minimize the ( weighted ) MSE loss while also minimizing the logarithm of the variance . Decoder Calibration It is important that the decoder distribution be calibrated in the statistical sense , that is , the predicted probabilities should correspond to the frequencies of seeing a particular value of x given that prediction ( DeGroot & Fienberg , 1983 ; Dawid , 1982 ) . The calibration of a neural network can be usually improved by estimating the uncertainty of that prediction ( Guo et al. , 2017 ) , such as the variance of a Gaussian ( Kendall & Gal , 2017 ) . Since the naive MSE loss assumes a constant variance , it does not effectively represent the uncertainty of the prediction , and is often poorly calibrated . Instead , learning the variance as in Eq . 3 leads to better uncertainty estimation and better calibration . In Sec 5.1 , we show that learning a good estimate of this uncertainty is crucial for the quality of the VAE generations . Connection to β-VAE . The β-VAE objective ( Higgins et al. , 2017 ) for a Gaussian decoder with unit variance is : Lβ = D 2 MSE ( x̂ , x ) + βDKL ( q ( z|x ) ||p ( z ) ) . ( 4 ) We see that it can be interpreted as a particular case of the objective ( 3 ) , where the variance is constant and the term D lnσ can be ignored during optimization . The β-VAE objective is then equivalent to a σ-VAE with a constant variance σ = √ β/2 ( for a particular learning rate setting ) . In recent work ( Zhu et al. , 2017 ; Denton & Fergus , 2018 ; Lee et al. , 2019 ) , β-VAE models are often used in this exact regime . By tuning the β term , practitioners are able to tune the variance of the decoder , manually producing a more calibrated decoder . However , by re-interpreting the β-VAE objective as a special case of the VAE and introducing the missing D lnσ term , we can both obtain a valid evidence lower bound , and remove the need to manually select β . Instead , the variance σ can instead simply be learned end-to-end , reducing the need for hyperparameter tuning . An alternative discussion of this connection in the context of linear VAEs is also presented by Lucas et al . ( 2019 ) . While the β term is not necessary for good performance if the decoder is calibrated , it can still be employed if desired , such as when the aim is to attain better disentanglement ( Higgins et al. , 2017 ) or a particular rate-distortion tradeoff ( Alemi et al. , 2017 ) . However , we found that with calibrated decoders , the best sample quality is obtained when β = 1 . Loss implementation details . For the correct evidence lower bound computation , it is necessary to add the values of the MSE loss and the KL divergence across the dimensions . We observe that common implementations of these losses ( Denton & Fergus , 2018 ; Abadi et al. , 2016 ; Paszke et al. , 2019 ) use averaging instead , which will lead to poor results if the number of image dimensions is significantly different from the number of the latent dimensions . While this can be conveniently ignored in the β-VAE regime , where the balance term is tuned manually anyway , for the σ-VAE it is essential to compute the objective value correctly . Variance implementation details . Since the variance is non-negative , we parameterize it logarithmically as σ2 = e2λ , where λ is the logarithm of the standard deviation . For some models , such as per-pixel variance decoders , we observed that it is necessary to restrict the variance range for numerical stability . We do so by using the soft clipping operations proposed by Chua et al . ( 2018 ) : λ : = λmax − softplus ( λmax − λ ) ; λ : = λmin + softplus ( λ− λmin ) . We observe that setting λmin = −6 to lower bound the standard deviation to be at least half of the distance between allowed color values works well in practice . We also observe that this clipping is unnecessary when learning a shared σ value .

Summary: This paper discusses a well-known problem of VAE training that decoder produces blurry reconstruction with constant variance. While much existing work addressed this problem by introducing independent variance training (as of the original VAE model) or additional hyper-parameters, those approaches usually come with additional training/tuning difficulty and even break the ELBO assumption. This paper proposed a simple $\sigma$-VAE that addresses the above problem by optimizing a single variance variable. This also could be easily connected to the well known $\beta$-VAE works. The experiment results in Tables 2 and 3 show the proposed model obtains a better FID score than the existing works on multiple datasets.