A Contrastive Learning Approach for Training Variational Autoencoder Priors

Variational autoencoders ( VAEs ) are one of the powerful likelihood-based generative models with applications in many domains . However , they struggle to generate high-quality images , especially when samples are obtained from the prior without any tempering . One explanation for VAEs ’ poor generative quality is the prior hole problem : the prior distribution fails to match the aggregate approximate posterior . Due to this mismatch , there exist areas in the latent space with high density under the prior that do not correspond to any encoded image . Samples from those areas are decoded to corrupted images . To tackle this issue , we propose an energy-based prior defined by the product of a base prior distribution and a reweighting factor , designed to bring the base closer to the aggregate posterior . We train the reweighting factor by noise contrastive estimation , and we generalize it to hierarchical VAEs with many latent variable groups . Our experiments confirm that the proposed noise contrastive priors improve the generative performance of state-of-the-art VAEs by a large margin on the MNIST , CIFAR-10 , CelebA 64 , and CelebA HQ 256 datasets . Our method is simple and can be applied to a wide variety of VAEs to improve the expressivity of their prior distribution . 1 Introduction q ( z ) p ( z ) r ( z ) Figure 1 : We propose an EBM prior using the product of a base prior p ( z ) and a reweighting factor r ( z ) , designed to bring p ( z ) closer to the aggregate posterior q ( z ) . Variational autoencoders ( VAEs ) [ 39 , 64 ] are one of the powerful likelihood-based generative models that have applications in image generation [ 6 , 35 , 62 ] , music synthesis [ 12 ] , speech generation [ 54 , 60 ] , image captioning [ 2 , 3 , 11 ] , semi-supervised learning [ 33 , 40 ] , and representation learning [ 15 , 79 ] . Although there has been tremendous progress in improving the expressivity of the approximate posterior , several studies have observed that VAE priors fail to match the aggregate ( approximate ) posterior [ 30 , 66 ] . This phenomenon is sometimes described as holes in the prior , referring to regions in the latent space that are not decoded to data-like samples . Such regions often have a high density under the prior but have a low density under the aggregate approximate posterior . The prior hole problem is commonly tackled by increasing the flexibility of the prior via hierarchical priors [ 42 ] , autoregressive models [ 21 ] , a mixture of encoders [ 72 ] , normalizing flows [ 8 , 81 ] , resampled priors [ 5 ] , and energy-based models [ 57 , 75–77 ] . Among them , energy-based models ( EBMs ) [ 13 , 57 ] have shown promising results . However , they require running iterative MCMC during training which is computationally expensive when the energy function is represented by a neural network . Moreover , they scale poorly to hierarchical models where an EBM is defined on each group of latent variables . 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . Our key insight in this work is that a trainable prior is brought as close as possible to the aggregate posterior as a result of training a VAE . The mismatch between the prior and the aggregate posterior can be reduced by reweighting the prior to re-adjust its likelihood in the area of mismatch with the aggregate posterior . To represent this reweighting mechanism , we formulate the prior using an EBM that is defined by the product of a reweighting factor and a base trainable prior as shown in Fig . 1 . We represent the reweighting factor using neural networks and the base prior using Normal distributions . Instead of computationally expensive MCMC sampling , notorious for being slow and often sensitive to the choice of parameters [ 13 ] , we use noise contrastive estimation ( NCE ) [ 22 ] for training the EBM prior . We show that NCE trains the reweighting factor in our prior by learning a binary classifier to distinguish samples from a target distribution ( i.e. , approximate posterior ) vs. samples from a noise distribution ( i.e. , the base trainable prior ) . However , since NCE ’ s success depends on closeness of the noise distribution to the target distribution , we first train the VAE with the base prior to bring it close to the aggregate posterior . And then , we train the EBM prior using NCE . In this paper , we make the following contributions : i ) We propose an EBM prior termed noise contrastive prior ( NCP ) which is trained by contrasting samples from the aggregate posterior to samples from a base prior . NCPs are simple and can be learned as a post-training mechanism to improve the expressivity of the prior . ii ) We also show how NCPs are trained on hierarchical VAEs with many latent variable groups . We show that training hierarchical NCPs scales easily to many groups , as they are trained for each latent variable group in parallel . iii ) Finally , we demonstrate that NCPs improve the generative quality of several forms of VAEs by a large margin across datasets . 2 Related Work . In this section , we review related prior works . Energy-based Models ( EBMs ) : Early work on EBMs for generative learning goes back to the 1980s [ 1 , 27 ] . Prior to the modern deep learning era , most attempts for building generative models using EBMs were centered around Boltzmann machines [ 26 , 28 ] and their “ deep ” extensions [ 44 , 67 ] . Although the energy function in these models is restricted to simple bilinear functions , they have been proven effective for representing the prior in discrete VAEs [ 65 , 75–77 ] . Recently , EBMs with neural energy functions have gained popularity for representing complex data distributions [ 13 ] . Pang et al . [ 57 ] have shown that neural EBMs can represent expressive prior distributions . However , in this case , the prior is trained using MCMC sampling , and it has been limited to a single group of latent variables . VAEBM [ 80 ] combines the VAE decoder with an EBM defined on the pixel space and trains the model using MCMC . Additionally , VAEBM assumes that data lies in a continuous space and applies the energy function in that space . Hence , it can not be applied to discrete data such as text or graphs . In contrast , NCP-VAE forms the energy function in the latent space and can be applied to non-continuous data . For continuous data , our model can be used along with VAEBM . We believe VAEBM and NCP-VAE are complementary . To avoid MCMC sampling , NCE [ 22 ] has recently been used for training a normalizing flow on data distributions [ 16 ] . Moreover , Han et al . [ 23 , 24 ] use divergence triangulation to sidesteps MCMC sampling . In contrast , we use NCE to train an EBM prior where a noise distribution is easily available through a pre-trained VAE . Adversarial Training : Similar to NCE , generative adversarial networks ( GANs ) [ 18 ] rely on a discriminator to learn the likelihood ratio between noise and real images . However , GANs use the discriminator to update the generator , whereas in NCE , the noise generator is fixed . In spirit similar are recent works [ 4 , 7 , 73 ] that link GANs , defined in the pixels space , to EBMs . We apply the likelihood ratio trick to the latent space of VAEs . The main difference : the base prior and approximate posterior are trained with the VAE objective rather than the adversarial loss . Adversarial loss has been used for training implicit encoders in VAEs [ 14 , 51 , 52 ] . But , they have not been linked to energy-based priors as we do explicitly . Prior Hole Problem : Among prior works on this problem , VampPrior [ 72 ] uses a mixture of encoders to represent the prior . However , this requires storing training data or pseudo-data to generate samples at test time . Takahashi et al . [ 70 ] use the likelihood ratio estimator to train a simple prior distribution . However at test time , the aggregate posterior is used for sampling in the latent space . Reweighted Priors : Bauer & Mnih [ 5 ] propose a reweighting factor similar to ours , but it is trained via truncated rejection sampling . Lawson et al . [ 45 ] introduce energy-inspired models ( EIMs ) that define distributions induced by the sampling processes used by Bauer & Mnih [ 5 ] as well as our sampling-importance-resampling ( SIR ) sampling ( called SNIS by Lawson et al . [ 45 ] ) . Although , EIMs have the advantage of end-to-end training , they require multiple samples during training ( up to 1K ) . This can make application of EIMs to deep hierarchical models such as NVAEs very challenging as these models are memory intensive and are trained with a few training samples per GPU . Moreover , our NCP scales easily to hierarchical models where the reweighting factor for each group is trained in parallel with other groups ( i.e. , NCP enables model parallelism ) . We view our proposed training method as a simple alternative approach that allows us to scale up EBM priors to large VAEs . Two-stage VAEs : VQ-VAE [ 62 , 79 ] first trains an autoencoder and then fits an autoregressive PixelCNN [ 78 ] prior to the latent variables . Albeit impressive results , autoregressive models can be very slow to sample from . Two-stage VAE ( 2s-VAE ) [ 10 ] trains a VAE on the data , and then , trains another VAE in the latent space . Regularized autoencoders ( RAE ) [ 17 ] train an autoencoder , and subsequently a Gaussian mixture model on latent codes . In contrast , we train the model with the original VAE objective in the first stage , and we improve the expressivity of the prior using an EBM . 3 Background . We first review VAEs , their extension to hierarchical VAEs before discussing the prior hole problem . Variational Autoencoders : VAEs learn a generative distribution p ( x , z ) = p ( z ) p ( x|z ) where p ( z ) is a prior distribution over the latent variable z and p ( x|z ) is a likelihood function that generates the data x given z. VAEs are trained by maximizing a variational lower bound LVAE ( x ) on the log-likelihood log p ( x ) ≥ LVAE ( x ) where LVAE ( x ) : = Eq ( z|x ) [ log p ( x|z ) ] − KL ( q ( z|x ) ||p ( z ) ) . ( 1 ) Here , q ( z|x ) is an approximate posterior and KL is the Kullback–Leibler divergence . The final training objective is Epd ( x ) [ LVAE ( x ) ] where pd ( x ) is the data distribution [ 39 ] . Hierarchical VAEs ( HVAEs ) : To increase the expressivity of both prior and approximate posterior , earlier work adapted a hierarchical latent variable structure [ 3 , 9 , 19 , 41 , 68 , 74 ] . In HVAEs , the latent variable z is divided into K separate groups , z = { z1 , . . . , zK } . The approximate posterior and the prior distributions are then defined by q ( z|x ) = ∏K k=1 q ( zk|z < k , x ) and p ( z ) = ∏K k=1 p ( zk|z < k ) . Using these , the training objective becomes LHVAE ( x ) : = Eq ( z|x ) [ log p ( x|z ) ] − K∑ k=1 Eq ( z < k|x ) [ KL ( q ( zk|z < k , x ) ||p ( zk|z < k ) ) ] , ( 2 ) where q ( z < k|x ) = ∏k−1 i=1 q ( zi|z < i , x ) is the approximate posterior up to the ( k − 1 ) th group.1 Prior Hole Problem : Let q ( z ) , Epd ( x ) [ q ( z|x ) ] denote the aggregate ( approximate ) posterior . In Appendix B.1 , we show that maximizing Epd ( x ) [ LVAE ( x ) ] w.r.t . the prior parameters corresponds to bringing the prior as close as possible to the aggregate posterior by minimizing KL ( q ( z ) ||p ( z ) ) w.r.t . p ( z ) . Formally , the prior hole problem refers to the phenomenon that p ( z ) fails to match q ( z ) .

Summary: In trained VAEs, mismatch between the prior and the variational posterior can lead to low-quality generated samples.  To address this issue, the authors suggest the following procedure: 1) given a trained VAE, train a classifier to distinguish images generated from the variational posterior vs images generated from the prior; 2) adjust the prior to minimize the accuracy of the trained classifier from step 1. This is done by reweighting the sampling procedure using the predictions from the classifier.  Main contribution: a practical method to improve the generative quality of several forms of VAE by reweighing the prior to diminish the mismatch from the variational posterior