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: The paper proposes to learn an energy based model using noise contrastive estimation on the latent of a VAE to mitigate the prior hole problem. To do so, the paper proposes a very simple method that consists in first training a standard VAE, then, given the frozen VAE, training a discriminator in latent space to discriminate between samples from the prior, and samples from encodings of real images. The discriminator can then be used to reweigh the prior, and obtain a better, but non algebraically computable, prior distribution. The implicit distribution can still be sampled using importance resampling or langevin dynamics. The paper proceeds to show that this technique improves the performance of the model in term of FID on various datasets.