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 goal of the paper is to model the marginal over latents in VAEs in such a way to minimize the mismatch with the aggregated posterior. The paper proposes a new class of marginal distributions over the latent space that is a product of two experts: the first expert is a non-trainable probability distribution, and the second expert is an unnormalized probability distribution parameterized using neural networks. Since training a product of experts requires to apply an approximate inference (e.g., MCMC sampling), the authors propose to use the likelihood ratio trick. Eventually, a VAE is trained in two stages. First, they assume the marginal over z's to be simply the non-trainable distribution, and the VAE is trained. At the second stage, they propose to train the second expert (i.e., the binary classifier that distinguishes z ~ q(z) and z ~ p(z)) in order to obtain the final NCP that better matches the aggregated posterior. Further, the idea is extended to hierarchical VAEs, and a separate binary classifier is trained per each stochastic level.