NCP-VAE: Variational Autoencoders with Noise Contrastive Priors

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 the base prior closer to the aggregate posterior q ( z ) . Variational autoencoders ( VAEs ) ( Kingma & Welling , 2014 ; Rezende et al. , 2014 ) are one of the powerful likelihood-based generative models that have applications in image generation ( Brock et al. , 2018 ; Karras et al. , 2019 ; Razavi et al. , 2019 ) , music synthesis ( Dhariwal et al. , 2020 ) , speech generation ( Oord et al. , 2016 ; Ping et al. , 2020 ) , image captioning ( Aneja et al. , 2019 ; Deshpande et al. , 2019 ; Aneja et al. , 2018 ) , semi-supervised learning ( Kingma et al. , 2014 ; Izmailov et al. , 2020 ) , and representation learning ( Van Den Oord et al. , 2017 ; Fortuin et al. , 2018 ) . 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 ( Rosca et al. , 2018 ; Hoffman & Johnson , 2016 ) . 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 ( Klushyn et al. , 2019 ) , autoregressive models ( Gulrajani et al. , 2016 ) , a mixture of approximate posteriors ( Tomczak & Welling , 2018 ) , normalizing flows ( Xu et al. , 2019 ; Chen et al. , 2016 ) , resampled priors ( Bauer & Mnih , 2019 ) , and energy-based models ( Pang et al. , 2020 ; Vahdat et al. , 2018b ; a ; 2020 ) . Among them , energy-based models ( EBMs ) ( Du & Mordatch , 2019 ; Pang et al. , 2020 ) have shown promising results in learning expressive priors . However , they require running iterative MCMC steps during training which is computationally expensive , especially 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 . Our key insight in this paper 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 simply 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 expensive MCMC sampling , we use noise contrastive estimation ( NCE ) ( Gutmann & Hyvärinen , 2010 ) for training the EBM prior . We show that NCE naturally trains the reweighting factor in our prior by learning a binary classifier to distinguish samples from a target distribution ( i.e. , samples from the approximate posterior ) vs. samples from a noise distribution ( i.e. , the base trainable prior ) . However , since NCE ’ s success depends on how close the noise distribution is 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 learned as a post-training mechanism to replace the original prior with a more flexible prior , which can improve the generative performance of VAEs with any structure . 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 VAEs by a large margin across datasets . 2 BACKGROUND . We first review VAEs , their extension to hierarchical VAEs , and 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 on the log-likelihood log p ( x ) : log p ( x ) ≥ Ez∼q ( z|x ) [ log p ( x|z ) ] − KL ( q ( z|x ) ||p ( z ) ) : = LVAE ( x ) , ( 1 ) where q ( z|x ) is an approximate posterior and KL is the Kullback–Leibler divergence . The final training objective is formulated by Epd ( x ) [ LVAE ( x ) ] where pd ( x ) is the data distribution ( Kingma & Welling , 2014 ) . Hierarchical VAEs ( HVAEs ) : To increase the expressivity of both prior and approximate posterior , earlier work adapted a hierarchical latent variable structure ( Vahdat & Kautz , 2020 ; Kingma et al. , 2016 ; Sønderby et al. , 2016 ; Gregor et al. , 2016 ) . 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 group1 . The 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 ) ] with respect to 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 ) . 1For k = 1 , the expectation inside the summation is simplified to KL ( q ( z1|x ) ||p ( z1 ) ) . 3 NOISE CONTRASTIVE PRIORS ( NCPS ) . One of the main causes of the prior hole problem is the limited expressivity of the prior that prevents it from matching the aggregate posterior . Recently , energy-based models have shown promising results in representing complex distributions . Motivated by their success , we introduce the noise contrastive prior ( NCP ) pNCP ( z ) = 1Z r ( z ) p ( z ) , where p ( z ) is a base prior distribution , e.g. , a Normal , r ( z ) is a reweighting factor , and Z = ∫ r ( z ) p ( z ) dz is the normalization constant . The function r : Rn → R+ maps the latent variable z ∈ Rn to a positive scalar , and can be implemented using neural nets . The reweighting factor r ( z ) can be trained using MCMC sampling as discussed in Appendix A . However , MCMC requires expensive sampling iterations that scale poorly to hierarchical VAEs . To address this , we describe a noise contrastive estimation based approach to train pNCP ( z ) without MCMC sampling . 3.1 LEARNING THE REWEIGHTING FACTOR WITH NOISE CONTRASTIVE ESTIMATION . Recall that training VAEs closes the gap between the prior and the aggregate posterior by minimizing KL ( q ( z ) ||p ( z ) ) with respect to prior . Assuming the base prior p ( z ) to be fixed , KL ( q ( z ) ||pNCP ( z ) ) is zero when r ( z ) = q ( z ) /p ( z ) . However , since we do not have the density function for q ( z ) , we can not compute the ratio explicitly . Instead , in this paper , we propose to estimate r ( z ) using noise contrastive estimation ( Gutmann & Hyvärinen , 2010 ) , also known as the likelihood ratio trick that has been popularized in machine learning by predictive coding ( Oord et al. , 2018 ) and generative adversarial networks ( GANs ) ( Goodfellow et al. , 2014 ) . Since , we can generate samples from both p ( z ) and q ( z ) 2 , we train a binary classifier to distinguish samples from q ( z ) and samples from the base prior p ( z ) by minimizing the binary cross-entropy loss : min D − Ez∼q ( z ) [ logD ( z ) ] − Ez∼p ( z ) [ log ( 1−D ( z ) ) ] . ( 3 ) Here , D : Rn → ( 0 , 1 ) is a binary classifier that generates the classification prediction probabilities . Eq . ( 3 ) is minimized when D ( z ) = q ( z ) q ( z ) +p ( z ) . Denoting the classifier at optimality by D ∗ ( z ) , we estimate the reweighting factor r ( z ) = q ( z ) p ( z ) ≈ D∗ ( z ) 1−D∗ ( z ) . The appealing advantage of this estimator is that it is obtained by simply training a binary classifier rather than using expensive MCMC sampling . 3.2 TWO-STAGE TRAINING FOR NOISE CONTRASTIVE PRIORS . To properly learn the reweighting factor , NCE training requires the base prior distribution to be close to the target distribution . Intuitively , if p ( z ) is very close to q ( z ) ( i.e. , p ( z ) ≈ q ( z ) ) , the optimal classifier will have a large loss value in Eq . ( 3 ) , and we will have r ( z ) ≈ 1 . If p ( z ) is instead far from q ( z ) , the binary classifier will easily learn to distinguish samples from the two distributions and it will not learn the likelihood ratios correctly . If p ( z ) is roughly close to q ( z ) , then the binary classifier can learn the ratios . To ensure that the base prior distribution is close to the target aggregate posterior distribution , we propose a two-stage training algorithm . In the first stage , we train the VAE with only the base prior p ( z ) . From Appendix B.1 , we know that at the end of training , p ( z ) is as close as possible to q ( z ) . In the second stage , we freeze the VAE model including the approximate posterior q ( z|x ) , the base prior p ( z ) , and the likelihood p ( x|z ) , and we only train the reweighting factor r ( z ) using Eq . ( 3 ) . The second stage can be thought of as replacing the base distribution p ( z ) with a more expressive distribution of the form pNCP ( z ) ∝ r ( z ) p ( z ) . Hence , NCP matches the prior to the aggregate posterior q ( z ) by using the reweighting factors . Note that our proposed method is generic as it only assumes that we can draw samples from q ( z ) and p ( z ) , which applies to any VAE . Our training is illustrated in Fig . 2 . 2We generate samples from the aggregate posterior q ( z ) = Epd ( x ) [ q ( z|x ) ] via ancestral sampling : draw data from the training set ( x ∼ pd ( x ) ) and then sample from z ∼ q ( z|x ) .

Summary: The authors highlight an important problem in VAE - the prior-hole problem - which is that the approximate posterior and the simple gaussian prior do not match in spite of the KL term in the ELBO which makes sampling an issue - leading to the prior putting probability mass on latents that are not decoded to high probability mass regions in data manifold. Prior approaches have overcome this problem by increasing the expressivity of the prior through autoregressive models, and/or using hierarchical latents, EBMs with MCMC sampling. This paper proposes a very simple two stage method - (1) train a regular VAE, (2) train a binary classifier in NCE style to distinguish samples from prior and approx. posterior; use the re-weighting term from the NCE score to sample from a better re-weighted prior - either through langevin dynamics or re-sampling. The authors combine this approach with the use of hierarchical latents and produce really good performing generative models on a host of benchmarks with good looking samples.