summaries: 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. Authors approach the "hole problem" of variational autoencoders where the aggregate posterior fails to match the prior, causing some areas of the prior distribution to be left out. Consequently, the decoder is not trained properly to operate in such regions, and the whole generate models is then subject to suboptimal performance. To attack this problem authors introduce two changes: 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.
abstractText: Variational autoencoders (VAEs) are one of the powerful likelihood-based generative models with applications in various 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.
references.authors: Jyoti Aneja Harsh Agrawal Dhruv Batra Alexander Schwing
references.title: Sequential latent spaces for modeling the intention during diverse image captioning
references.venue: In Proceedings of the IEEE International Conference on Computer Vision,
references.authors: Samaneh Azadi Catherine Olsson Trevor Darrell Ian Goodfellow Augustus Odena
references.title: Discriminator rejection sampling
references.venue: In International Conference on Learning Representations,
references.authors: Matthias Bauer Andriy Mnih
references.title: Resampled priors for variational autoencoders
references.venue: In The 22nd International Conference on Artificial Intelligence and Statistics,
references.authors: Andrew Brock Jeff Donahue Karen Simonyan
references.title: Large scale gan training for high fidelity natural image synthesis
references.venue: arXiv preprint arXiv:1809.11096,
references.authors: Tong Che Ruixiang Zhang Jascha Sohl-Dickstein Hugo Larochelle Liam Paull Yuan Cao Yoshua Bengio
references.title: Your GAN is secretly an energy-based model and you should use discriminator driven latent sampling
references.venue: arXiv preprint arXiv:2003.06060,
references.authors: Xi Chen Diederik P Kingma Tim Salimans Yan Duan Prafulla Dhariwal John Schulman Ilya Sutskever Pieter Abbeel
references.title: Variational lossy autoencoder
references.venue: arXiv preprint arXiv:1611.02731,
references.authors: Bin Dai David Wipf
references.title: Diagnosing and enhancing vae models
references.venue: In International Conference on Learning Representations,
references.authors: Aditya Deshpande Jyoti Aneja Liwei Wang Alexander G Schwing David Forsyth
references.title: Fast, diverse and accurate image captioning guided by part-of-speech
references.venue: In Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition,
references.authors: Prafulla Dhariwal Heewoo Jun Christine Payne Jong Wook Kim Alec Radford Ilya Sutskever
references.title: Jukebox: A generative model for music
references.venue: arXiv preprint arXiv:2005.00341,
references.authors: Yilun Du Igor Mordatch
references.title: Implicit generation and modeling with energy based models
references.venue: In Advances in Neural Information Processing Systems,
references.authors: Jesse Engel Matthew Hoffman Adam Roberts
references.title: Latent constraints: Learning to generate conditionally from unconditional generative models
references.venue: In International Conference on Learning Representations,
references.authors: Vincent Fortuin Matthias Hüser Francesco Locatello Heiko Strathmann Gunnar Rätsch
references.title: Som-vae: Interpretable discrete representation learning on time series
references.venue: In International Conference on Learning Representations,
references.authors: Ruiqi Gao Erik Nijkamp Diederik P Kingma Zhen Xu Andrew M Dai Ying Nian Wu
references.title: Flow contrastive estimation of energy-based models
references.venue: In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition,
references.authors: Partha Ghosh Mehdi S.M. Sajjadi Antonio Vergari Michael Black Bernhard
references.title: Scholkopf. From variational to deterministic autoencoders
references.venue: In International Conference on Learning Representations,
references.authors: Karol Gregor Frederic Besse Danilo Jimenez Rezende Ivo Danihelka Daan Wierstra
references.title: Towards conceptual compression
references.venue: In Advances In Neural Information Processing Systems,
references.authors: Ishaan Gulrajani Kundan Kumar Faruk Ahmed Adrien Ali Taiga Francesco Visin David Vazquez Aaron Courville
references.title: PixelVAE: A latent variable model for natural images
references.venue: arXiv preprint arXiv:1611.05013,
references.authors: Michael Gutmann Aapo Hyvärinen
references.title: Noise-contrastive estimation: A new estimation principle for unnormalized statistical models
references.venue: In Proceedings of the Thirteenth International Conference on Artificial Intelligence and Statistics,
references.authors: Tian Han Erik Nijkamp Xiaolin Fang Mitch Hill Song-Chun Zhu Ying Nian Wu
references.title: Divergence triangle for joint training of generator model, energy-based model, and inferential model
references.venue: In Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition,
references.authors: Tian Han Erik Nijkamp Linqi Zhou Bo Pang Song-Chun Zhu Ying Nian Wu
references.title: Joint training of variational auto-encoder and latent energy-based model
references.venue: In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition,
references.authors: Martin Heusel Hubert Ramsauer Thomas Unterthiner Bernhard Nessler Sepp Hochreiter
references.title: Gans trained by a two time-scale update rule converge to a local nash equilibrium
references.venue: In Advances in neural information processing systems,
references.authors: Geoffrey E Hinton
references.title: Training products of experts by minimizing contrastive divergence
references.venue: Neural computation,
references.authors: Geoffrey E Hinton Terrence J Sejnowski
references.title: Learning and relearning in boltzmann machines. Parallel distributed processing: Explorations in the microstructure of cognition
references.authors: Geoffrey E Hinton Simon Osindero Yee-Whye Teh
references.title: A fast learning algorithm for deep belief nets
references.venue: Neural computation,
references.authors: Jonathan Ho Ajay Jain Pieter Abbeel
references.title: Denoising diffusion probabilistic models
references.venue: arXiv preprint arxiv:2006.11239,
references.authors: Matthew D Hoffman Matthew J Johnson
references.title: Elbo surgery: yet another way to carve up the variational evidence lower bound
references.venue: In Workshop in Advances in Approximate Bayesian Inference,
references.authors: Jie Hu Li Shen Gang Sun
references.title: Squeeze-and-excitation networks
references.venue: In Proceedings of the IEEE conference on computer vision and pattern recognition,
references.authors: Sergey Ioffe Christian Szegedy
references.title: Batch normalization: Accelerating deep network training by reducing internal covariate shift
references.venue: In International Conference on Machine Learning,
references.authors: Pavel Izmailov Polina Kirichenko Marc Finzi Andrew Gordon Wilson
references.title: Semi-supervised learning with normalizing flows
references.venue: In ICML,
references.authors: Tero Karras Timo Aila Samuli Laine Jaakko Lehtinen
references.title: Progressive growing of gans for improved quality, stability, and variation
references.venue: arXiv preprint arXiv:1710.10196,
references.authors: Tero Karras Samuli Laine Timo Aila
references.title: A style-based generator architecture for generative adversarial networks
references.venue: In Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition,
references.authors: Tero Karras Miika Aittala Janne Hellsten Samuli Laine Jaakko Lehtinen Timo Aila
references.title: Training generative adversarial networks with limited data
references.venue: arXiv preprint arXiv:2006.06676,
references.authors: Diederik Kingma Jimmy Ba
references.title: Adam: A method for stochastic optimization
references.venue: arXiv preprint arXiv:1412.6980,
references.authors: Diederik P Kingma Max Welling
references.title: Auto-encoding variational bayes
references.venue: In The International Conference on Learning Representations (ICLR),
references.authors: Diederik P Kingma Shakir Mohamed Danilo Jimenez Rezende Max Welling
references.title: Semi-supervised learning with deep generative models
references.venue: In Advances in Neural Information Processing Systems,
references.authors: Diederik P Kingma Tim Salimans Rafal Jozefowicz Xi Chen Ilya Sutskever Max Welling
references.title: Improved variational inference with inverse autoregressive flow
references.venue: In Advances in Neural Information Processing Systems,
references.authors: Durk P Kingma Prafulla Dhariwal
references.title: Glow: Generative flow with invertible 1x1 convolutions
references.venue: Advances in Neural Information Processing Systems
references.authors: Alexej Klushyn Nutan Chen Richard Kurle Botond Cseke Patrick van der Smagt
references.title: Learning hierarchical priors in vaes
references.venue: In Advances in Neural Information Processing Systems,
references.authors: Alex Krizhevsky Geoffrey Hinton
references.title: Learning multiple layers of features from tiny
references.authors: Hugo Larochelle Yoshua Bengio
references.title: Classification using discriminative restricted boltzmann machines
references.venue: In Proceedings of the 25th international conference on Machine learning,
references.authors: John Lawson George Tucker Bo Dai Rajesh Ranganath
references.title: Energy-inspired models: Learning with sampler-induced distributions
references.venue: In Advances in Neural Information Processing Systems,
references.authors: Yann LeCun
references.title: The mnist database of handwritten digits. http://yann
references.venue: lecun. com/exdb/mnist/,
references.authors: Chieh Hubert Lin Chia-Che Chang Yu-Sheng Chen Da-Cheng Juan Wei Wei Hwann-Tzong Chen
references.title: Coco-gan: generation by parts via conditional coordinating
references.venue: In Proceedings of the IEEE International Conference on Computer Vision,
references.authors: Ziwei Liu Ping Luo Xiaogang Wang Xiaoou Tang
references.title: Deep learning face attributes in the wild
references.venue: In Proceedings of the IEEE international conference on computer vision,
references.authors: Ilya Loshchilov Frank Hutter
references.title: Sgdr: Stochastic gradient descent with warm restarts
references.venue: arXiv preprint arXiv:1608.03983,
references.authors: Lars Maaløe Marco Fraccaro Valentin Liévin Ole Winther
references.title: BIVA: A very deep hierarchy of latent variables for generative modeling
references.venue: In Advances in neural information processing systems,
references.authors: Lars Mescheder S Nowozin Andreas Geiger
references.title: Adversarial variational bayes: Unifying variational autoencoders and generative adversarial networks
references.venue: In 34th International Conference on Machine Learning (ICML),
references.authors: Charlie Nash Conor Durkan
references.title: Autoregressive energy machines
references.venue: arXiv preprint arXiv:1904.05626,
references.authors: Aaron van den Oord Sander Dieleman Heiga Zen Karen Simonyan Oriol Vinyals Alex Graves Nal Kalchbrenner Andrew Senior Koray Kavukcuoglu
references.title: Wavenet: A generative model for raw audio
references.venue: arXiv preprint arXiv:1609.03499,
references.authors: Aaron van den Oord Yazhe Li Oriol Vinyals
references.title: Representation learning with contrastive predictive coding
references.venue: arXiv preprint arXiv:1807.03748,
references.authors: Art B. Owen
references.title: Monte Carlo theory, methods and examples
references.authors: Bo Pang Tian Han Erik Nijkamp Song-Chun Zhu Ying Nian Wu
references.title: Learning latent space energy-based prior model
references.venue: arXiv preprint arXiv:2006.08205,
references.authors: Kancharla Parimala Sumohana Channappayya
references.title: Quality aware generative adversarial networks
references.venue: In Advances in Neural Information Processing Systems,
references.authors: Stanislav Pidhorskyi Donald A Adjeroh Gianfranco Doretto
references.title: Adversarial latent autoencoders
references.venue: In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition,
references.authors: Wei Ping Kainan Peng Kexin Zhao Zhao Song
references.title: Waveflow: A compact flow-based model for raw audio
references.authors: Prajit Ramachandran Barret Zoph Quoc V Le
references.title: Searching for activation functions
references.venue: arXiv preprint arXiv:1710.05941,
references.authors: Ali Razavi Aaron van den Oord Oriol Vinyals
references.title: Generating diverse high-fidelity images with vq-vae-2
references.venue: In Advances in Neural Information Processing Systems,
references.authors: Danilo Jimenez Rezende Shakir Mohamed Daan Wierstra
references.title: Stochastic backpropagation and approximate inference in deep generative models
references.venue: In International Conference on Machine Learning,
references.authors: Jason Tyler Rolfe
references.title: Discrete variational autoencoders
references.venue: arXiv preprint arXiv:1609.02200,
references.authors: Mihaela Rosca Balaji Lakshminarayanan Shakir Mohamed
references.title: Distribution matching in variational inference
references.venue: arXiv preprint arXiv:1802.06847,
references.authors: Ruslan Salakhutdinov Geoffrey Hinton
references.title: Deep boltzmann machines
references.venue: In Artificial intelligence and statistics,
references.authors: Casper Kaae Sønderby Tapani Raiko Lars Maaløe Søren Kaae Sønderby Ole Winther
references.title: Ladder variational autoencoders
references.venue: In Advances in neural information processing systems,
references.authors: Hiroshi Takahashi Tomoharu Iwata Yuki Yamanaka Masanori Yamada Satoshi Yagi
references.title: Variational autoencoder with implicit optimal priors
references.venue: In Proceedings of the AAAI Conference on Artificial Intelligence,
references.authors: I Tolstikhin O Bousquet S Gelly B Schölkopf
references.title: Wasserstein auto-encoders
references.venue: In International Conference on Learning Representations (ICLR 2018). OpenReview
references.authors: Jakub Tomczak Max Welling
references.title: Vae with a vampprior
references.venue: In International Conference on Artificial Intelligence and Statistics,
references.authors: Ryan Turner Jane Hung Eric Frank Yunus Saatchi Jason Yosinski
references.title: Metropolis-hastings generative adversarial networks
references.venue: In International Conference on Machine Learning,
references.authors: Arash Vahdat Jan Kautz
references.title: NVAE: A deep hierarchical variational autoencoder
references.venue: In Neural Information Processing Systems (NeurIPS),
references.authors: Arash Vahdat Evgeny Andriyash William G Macready
references.title: DVAE#: Discrete variational autoencoders with relaxed Boltzmann priors
references.venue: In Neural Information Processing Systems,
references.authors: Arash Vahdat William G. Macready Zhengbing Bian Amir Khoshaman Evgeny Andriyash
references.title: DVAE++: Discrete variational autoencoders with overlapping transformations
references.venue: In International Conference on Machine Learning (ICML),
references.authors: Arash Vahdat Evgeny Andriyash William G Macready
references.title: Undirected graphical models as approximate posteriors
references.venue: In International Conference on Machine Learning (ICML),
references.authors: Aäron Van Den Oord Nal Kalchbrenner Koray Kavukcuoglu
references.title: Pixel recurrent neural networks
references.venue: In Proceedings of the 33rd International Conference on International Conference on Machine Learning,
references.authors: Aaron Van Den Oord Oriol Vinyals
references.title: Neural discrete representation learning
references.venue: In Advances in Neural Information Processing Systems,
references.authors: Haowen Xu Wenxiao Chen Jinlin Lai Zhihan Li Youjian Zhao Dan Pei
references.title: On the necessity and effectiveness of learning the prior of variational auto-encoder
references.authors: Zijun Zhang Ruixiang Zhang Zongpeng Li Yoshua Bengio Liam Paull
references.title: Perceptual generative autoencoders
references.venue: In International Conference on Machine Learning,
references.authors: Adam (Kingma Ba
references.title: Learning Rate Initialize at 1e-3, CosineAnnealing (Loshchilov & Hutter, 2016
sections.text: 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.
sections.heading: 2 BACKGROUND
sections.text: 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)).
sections.heading: 3 NOISE CONTRASTIVE PRIORS (NCPS)
sections.text: 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.
sections.heading: 3.1 LEARNING THE REWEIGHTING FACTOR WITH NOISE CONTRASTIVE ESTIMATION
sections.text: 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 cannot 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.
sections.heading: 3.2 TWO-STAGE TRAINING FOR NOISE CONTRASTIVE PRIORS
sections.text: 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).
sections.heading: 3.3 TEST TIME SAMPLING
sections.text: To sample from a VAE with an NCP, we first generate samples from the NCP before passing them to the decoder to generate output samples (shown in Fig. 2). We propose two methods for sampling from NCPs.
  Sampling-Importance-Resampling (SIR): We first generate M samples from the base prior distribution {z(m)}Mm=1 ∼ p(z). We then resample one of the M proposed samples using importance weights proportional to w(m) = pNCP(z(m))/p(z(m)) = r(z(m)). The benefit of this technique: both proposal generation and the evaluation of r on the samples are done in parallel.
  Langevin Dynamics (LD): Since our NCP is an EBM, we can use LD for sampling. Denoting the energy function by E(z) = − log r(z) − log p(z), we initialize a sample z0 by drawing from p(z) and update the sample iteratively using: zt+1 = zt − 0.5 λ∇zE(z) + √ λ t where t ∼ N (0, 1) and λ is the step size. LD is run for a finite number of iterations, and in contrast to SIR, it can be slow given its sequential form.
sections.heading: 3.4 GENERALIZATION TO HIERARCHICAL VAES
sections.text: The state-of-the-art VAE (Vahdat & Kautz, 2020) uses a hierarchical q(z|x) and p(z). Here p(z) is chosen to be a Gaussian distribution. Appendix B.2 shows that training a HVAE encourages the prior to minimize Eq(z<k) [KL(q(zk|z<k)||p(zk|z<k))] for each conditional, where q(z<k) , Epd(x)[q(z<K |x)] is the aggregate posterior up to the (k − 1)th group, and q(zk|z<k) , Epd(x)[q(zk|z<k,x)] is the aggregate conditional for the kth group. Given this observation, we extend NCPs to hierarchical models to match each conditional in the prior with q(zk|z<k). Formally, we define hierarchical NCPs by pNCP(z) = 1 Z ∏K k=1 r(zk|z<k)p(zk|z<k) where each factor is an EBM. pNCP(z) resembles energy machines with an autoregressive structure among groups (Nash & Durkan, 2019).
  In the first stage, we train the HVAE with prior ∏K
  k=1 p(zk|z<k). For the second stage, we use K binary classifiers, each for a hierarchical group. Following Appendix C, we train each classifier via:
  min Dk
  Epd(x)q(z<k|x) [ − Eq(zk|z<k,x)[logDk(zk, c(z<k))]− Ep(zk|z<k)[log(1−Dk(zk, c(z<k)))] ] , (4)
  where the outer expectation samples from groups up to the (k−1)th group, and the inner expectations sample from approximate posterior and base prior for the kth group, conditioned on the same z<k. The discriminator Dk classifies samples zk while conditioning its prediction on z<k using a shared context feature c(z<k).
  The NCE training in Eq. (4) is minimized when Dk(zk, c(z<k)) = q(zk|z<k)
  q(zk|z<k)+p(zk|z<k) . Denoting the
  classifier at optimality by D∗k(z, c(z<k)), we obtain the reweighting factor r(zk|z<k) ≈ D∗k(zk,c(z<k))
  1−D∗k(zk,c(z<k)) in the second stage. Given our hierarchical NCP, we use ancestral sampling to sample from the prior. For sampling from each group, we can use SIR or LD as discussed before.
  The context feature c(z<k) extracts a representation from z<k. Instead of learning a new representation at stage two, we simply use the representation that is extracted from z<k in the hierarchical prior, trained in the first stage. Note that the binary classifiers in the second stage are trained in parallel for all groups.
sections.heading: 4 RELATED WORK
sections.text: In this section, we review prior works related to the proposed method.
  Energy-based Models (EBMs): Early work on EBMs for generative learning goes back to 1980s (Ackley et al., 1985; Hinton et al., 1986). Prior to the modern deep learning era, most attempts for building generative models using EBMs were centered around Boltzmann machines (Hinton, 2002; Hinton et al., 2006) and their “deep” extensions (Salakhutdinov & Hinton, 2009; Larochelle & Bengio, 2008). 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 (Rolfe, 2016; Vahdat et al., 2018a;b; 2020). Recently, EBMs with neural energy functions have gained popularity for representing complex data distribution (Du & Mordatch, 2019). Pang et al. (2020) 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. To eliminate MCMC sampling, NCE (Gutmann & Hyvärinen, 2010) has recently been used for training a normalizing flow on data distributions (Gao et al., 2020). Moreover, Han et al. (2019; 2020) have used 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) (Goodfellow et al., 2014) also 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 (Azadi et al., 2018; Turner et al., 2019; Che et al., 2020) 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 (Makhzani et al., 2015; Mescheder et al., 2017; Engel et al., 2018). But, they have not been linked to energy-based priors as we do explicitly.
  Prior Hole Problem: Among prior works on this problem, VampPrior (Tomczak & Welling, 2018) 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. (2019) 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. Bauer & Mnih (2019) propose a reweighting factor similar to ours, but it is trained via importance sampling. Recently, Lawson et al. (2019) introduced energy-inspired models (EIMs) that define distributions induced by the sampling processes used by (Bauer & Mnih, 2019) as well as our SIR sampling. Although, EIMs have the advantage of end-to-end training and can be used as either prior or decoder in VAEs, 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 often very memory intensive and are trained with a few training samples per GPU.
  Two-stage VAEs: VQ-VAE (Van Den Oord et al., 2017; Razavi et al., 2019) first trains an autoencoder and then fits an autoregressive PixelCNN (Van Den Oord et al., 2016) prior to the latent variables which is slow to sample from. Two-stage VAE (2s-VAE) (Dai & Wipf, 2018) trains a VAE on the data, and then, trains another VAE in the latent space. Regularized autoencoders (RAE) (Ghosh et al., 2020) 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.
sections.heading: 5 EXPERIMENTS
sections.text: Our implementation of NCP-VAE builds upon NVAE (Vahdat & Kautz, 2020), the state-of-the-art hierarchical VAE. We examine NCP-VAE on four datasets including dynamically binarized MNIST (LeCun, 1998), CIFAR-10 (Krizhevsky et al., 2009), CelebA-64 (Liu et al., 2015) and CelebA-HQ-256 (Karras et al., 2017). For CIFAR-10 and CelebA-64, the model has 30 groups, and for CelebA-HQ-256 it has 20 groups. We sample from NCP-VAE using SIR with 5K proposal samples at the test (image generation) time. On these datasets, we measure the sample quality using the Fréchet Inception Distance (FID) score (Heusel et al., 2017) with 50,000 samples, as computing the log-likelihood requires estimating the intractable normalization constant. To report log-likelihood results, we train an NVAE model with a small latent space on MNIST with 10 groups of 4 × 4 latent variables. We estimate log normalization constant in NCPs using 1000 importance weighted samples. We intentionally limit the latent space to ensure that we can estimate the normalization constant correctly (standard deviation of logZ estimation ≤ 0.23). Thus, on this dataset, we report negative log-likelihood (NLL). Implementation details are provided in Appendix E.
sections.heading: 5.1 QUANTITATIVE RESULTS
sections.text: The quantitative results are reported in Tab. 1, Tab. 2, Tab. 3, and Tab. 4. On all four datasets, our model improves upon state-of-the-art NVAE, and it reduces the gap with GANs by a large margin. On CelebA 64, we improve NVAE from an FID of 13.48 to 5.25, comparable to GANs. On CIFAR-10, NCP-VAE
  Table 1: Generative performance on CelebA-64
  Model FID↓ NCP-VAE (ours) 5.25 NVAE (Vahdat & Kautz, 2020) 13.48 RAE (Ghosh et al., 2020) 40.95 2s-VAE (Dai & Wipf, 2018) 44.4 WAE (Tolstikhin et al., 2018) 35 Perceptial AE(Zhang et al., 2020) 13.8 Latent EBM (Pang et al., 2020) 37.87 COCO-GAN (Lin et al., 2019) 4.0 QA-GAN (Parimala & Channappayya, 2019) 6.42 NVAE-Recon (Vahdat & Kautz, 2020) 1.03
  Table 3: Generative results on CelebA-HQ-256
  Model FID↓ NCP-VAE (ours) 24.79 NVAE (Vahdat & Kautz, 2020) 40.26 GLOW (Kingma & Dhariwal, 2018) 68.93 Advers. LAE (Pidhorskyi et al., 2020) 19.21 PGGAN (Karras et al., 2017) 8.03 NVAE-Recon (Vahdat & Kautz, 2020) 0.45
  Table 4: Likelihood results on MNIST in nats
  Model NLL↓ NCP-VAE (ours) 78.10 NVAE-small (Vahdat & Kautz, 2020) 78.67 BIVA (Maaløe et al., 2019) 78.41 DAVE++ (Vahdat et al., 2018b) 78.49 IAF-VAE (Kingma et al., 2016) 79.10 VampPrior AR dec. (Tomczak & Welling) 78.45
  improves the NVAE FID of 51.71 to 24.08. On MNIST, although our latent space is much smaller, our model outperforms previous VAEs. NVAE has reported 78.01 nats on this dataset with a larger latent space.
  Reconstruction: In the last row of Tab. 1, Tab. 2, and Tab. 3, we report the FID score for the reconstructed images of the NVAE baseline. Note how reconstruction FID is much lower than our FID, indicating that our model is far from memorizing the training data. In Appendix G, we also provide nearest neighbours from training data for generated samples. In Appendix F, we present our NCP applied to vanilla VAEs.
sections.heading: 5.2 QUALITATIVE RESULTS
sections.text: We visualize samples generated by NCP-VAE in Fig. 3 without any manual intervention. We adopt the common practice of reducing the temperature of the base prior p(z) by scaling down the standard-deviation of the conditional Normal distributions (Kingma & Dhariwal, 2018)3. Brock et al. (2018); Vahdat & Kautz (2020) also observe that re-adjusting the batch-normalization (BN), given a temperature applied to the prior, improves the generative quality. Similarly, we achieve diverse, high-quality images by re-adjusting the BN statistics as described by Vahdat & Kautz (2020). Additional qualitative results are shown in Appendix H.
sections.heading: 5.3 ADDITIONAL STUDIES
sections.text: We perform additional experiments to study i) how hierarchical NCPs perform as the number of latent groups increases, ii) the impact of SIR and LD hyperparameters, and iii) what the classification loss in NCE training conveys about p(z) and q(z). All experiments in this section are performed on CelebA-64 data.
  3Lowering of the temperature is only used to obtain qualitative samples. It’s not used when computing any of the quantitative results in Sec. 5.1.
  Number of latent variable groups: Tab. 5 shows the generative performance of hierarchical NCP with different amounts of latent variable groups. As we increase the number of groups, the FID score of both NVAE and our model improves. This demonstrates the efficacy of our NCPs, even with expressive hierarchical priors and in the presence of many groups.
  SIR and LD parameters: The computationally complexity of SIR is similar to LD if we set the number of proposal samples in SIR equal to the number LD iterations. Tab. 6 reports the impact of these parameters. We observe that increasing both the number of proposal samples in SIR and the LD iterations leads to a noticeable improvement in FID score. For SIR, the proposal generation and the evaluation of r(z) are parallelizable. Hence, as shown in Tab. 6, image generation is faster with SIR than with LD (LD iterations are sequential). However, GPU memory usage scales with the number of SIR proposals, but not with the number of LD iterations. Interestingly, SIR, albeit simple, performs better than LD when using about the same compute.
  Classification loss in NCE: We can draw a direct connection between the classification loss in Eq. 3 and the similarity of p(z) and q(z). Denoting the classification loss in Eq. 3 at optimality by L∗, Goodfellow et al. (2014) show that JSD(p(z)||q(z)) = log 2 − 0.5 × L∗ where JSD denotes the Jensen–Shannon divergence between two distributions. Fig. 4(a) plots the classification loss (Eq. 4) for each classifier for a 15-group NCP trained on the CelebA-64 dataset. Assume that the classifier loss at the end of training is a good approximation of L∗. We observe that 8 out of 15 groups have L∗ ≥ 0.4, indicating a good overlap between p(z) and q(z) for those groups. To further assess the impact of the distribution match on SIR sampling, in Fig. 4(b), we visualize the effective sample size (ESS)4 in SIR vs. L∗ for the same group. We observe a strong correlation between L∗ and the effective sample size. SIR is more reliable on the same 8 groups that have high classification loss. These groups are mostly at the top of the NVAE hierarchy which have been shown to control the global structure of generated samples (see B.6 in Vahdat & Kautz (2020)).
  4ESS measures reliability of SIR via 1/ ∑
  m(ŵ (m))2, where ŵ(m) = r(z(m))/ ∑ m′ r(z (m′)) (Owen, 2013).
sections.heading: 5.4 ADDITIONAL COMPARISON TO PRIOR ART
sections.text: In our previous experiments, we rely on NVAE as the VAE backbone. Although this shows that our approach can be applied to large scale models successfully, this may make comparison against the prior art unfair as they are often applied to smaller models. To provide a fair comparison, in this section, we apply NCP to several commonly used small VAE models.
  Comparison against RAE, 2s-VAE, and WAE: In Tab. 7 we show the generative performance of our approach applied to the VAE architecture in RAE (Ghosh et al., 2020). Note that this VAE architecture has only one latent variable group. The same base architecture was used in the implementation of 2s-VAE (Dai & Wipf, 2018) and WAE (Tolstikhin et al., 2018). We borrow the training setup from Ghosh et al. (2020) on the CelebA-64 dataset, and compare to the baselines reported in this work. We apply our NCP-VAE on top of vanilla VAE with a Gaussian prior as well as a 10-component Gaussian mixture model (GMM) prior that was proposed in RAEs. Our NCP-VAE improves the performance on the base VAE, improving the FID score to 41.28 from 48.12. Additionally, when NCP is applied to the VAE with GMM prior (the RAE model), it improves its performance from 40.95 to the FID score of 39.00.
  Comparison against LARS and SNIS: In Tab. 8, we compare NCP-VAE to LARS (Bauer & Mnih, 2019) and SNIS (Lawson et al., 2019) priors on MNIST. We implement our NCP-VAE using the VAE and energyfunction architectures that were used by Lawson et al. (2019). We closely follow the training hyperparameters used by Lawson et al. (2019) as well as their approach for obtaining a lower bound on the log likelihood. NCP-VAE obtains the negative log-likelihood (NLL) of 82.82, comparable to Lawson et al. (2019), while outperforming LARS (Bauer & Mnih, 2019).
sections.heading: 6 CONCLUSIONS
sections.text: The prior hole problem is one of the main reasons for VAEs’ poor generative quality. In this paper, we tackled this problem by introducing the noise contrastive prior (NCP), defined by the product of a reweighting factor and a base prior. We showed how the reweighting factor is trained by contrasting samples from the aggregate posterior with samples from the base prior. Our proposal is simple and can be applied to any VAE to increase its prior’s expressivity. We also showed how NCP training scales to large hierarchical VAEs, as it can be done in parallel simultaneously for all the groups. Finally, we demonstrated that NCPs improve the generative performance of state-of-the-art NVAEs by a large margin, closing the gap to GANs.
sections.heading: A TRAINING ENERGY-BASED PRIORS USING MCMC
sections.text: In this section, we show how a VAE with energy-based model in its prior can be trained. Assuming that the prior is in the form pEBM(z) = 1Z r(z)p(z), the variational bound is of the form:
  Epd(x)[LVAE] = Epd(x) [ Eq(z|x)[log p(x|z)]− KL(q(z|x)||pEBM(z)) ] = Epd(x) [ Eq(z|x)[log p(x|z)− log q(z|x) + log r(z) + log p(z)] ] − logZ,
  where the expectation term, similar to VAEs, can be trained using the reparameterization trick. The only problematic term is the log-normalization constant logZ, which captures the gradient with respect to the parameters of the prior pEBM(z). Denoting these parameters by θ, the gradient of logZ is obtained by:
  ∂
  ∂θ logZ =
  1
  Z
  ∫ ∂(r(z)p(z))
  ∂θ dz =
  ∫ r(z)p(z)
  Z
  ∂ log(r(z)p(z))
  ∂θ dz = EPEBM (z)[
  ∂ log(r(z)p(z))
  ∂θ ], (5)
  where the expectation can be estimated using MCMC sampling from the EBM prior.
sections.heading: B MAXIMIZING THE VARIATIONAL BOUND FROM THE PRIOR’S PERSPECTIVE
sections.text: In this section, we discuss how maximizing the variational bound in VAEs from the prior’s perspective corresponds to minimizing a KL divergence from the aggregate posterior to the prior. Note that this relation has been explored by Hoffman & Johnson (2016); Rezende & Viola (2018); Tomczak & Welling (2018) and we include it here for completeness.
  B.1 VAE WITH A SINGLE GROUP OF LATENT VARIABLES
  Denote the aggregate (approximate) posterior by q(z) , Epd(x)[q(z|x)]. Here, we show that maximizing the Epd(x)[LVAE(x)] with respect to the prior parameters corresponds to learning the prior by minimizing KL(q(z)||p(z)). To see this, note that the prior p(z) only participates in the KL term in LVAE (Eq. 1). We hence have:
  argmax p(z) Epd(x)[LVAE(x)] = argmin p(z) Epd(x)[KL(q(z|x)||p(z))]
  = argmin p(z) −Epd(x)[H(q(z|x))]− Eq(z)[log p(z)]
  = argmin p(z)
  −H(q(z))− Eq(z)[log p(z)]
  = argmin p(z)
  KL(q(z)||p(z)),
  where H(.) denotes the entropy. Above, we replaced the expected entropy Epd(x)[H(q(z|x))] with H(q(z)) as the minimization is with respect to the parameters of the prior p(z).
  B.2 HIERARCHICAL VAES Denote hierarchical approximate posterior and prior distributions by: q(z|x) = ∏K
  k=1 q(zk|z<k,x) and p(z) = ∏K k=1 p(zk|z<k). The hierarchical VAE 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))] , (6)
  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. Denote the aggregate posterior up to the (K − 1)th group by q(z<k) , Epd(x)[q(z<K |x)] and the aggregate conditional for the kth group given the previous groups q(zk|z<k) , Epd(x)[q(zk|z<k,x)].
  Here, we show that maximizing Epd(x)[LHVAE(x)] with respect to the prior corresponds to learning the prior by minimizing Eq(z<k) [KL(q(zk|z<k)||p(zk|z<k))] for each conditional:
  argmax p(zk|z<k) Epd(x)[LHVAE(x)] = argmin p(zk|z<k)
  Epd(x) [ Eq(z<k|x) [KL(q(zk|z<k,x)||p(zk|z<k))] ] = argmin
  p(zk|z<k) −Epd(x)q(z<k|x)q(zk|z<k,x) [log p(zk|z<k)]
  = argmin p(zk|z<k) −Eq(zk,z<k) [log p(zk|z<k)]
  = argmin p(zk|z<k)
  −Eq(z<k) [ Eq(zk|z<k) [log p(zk|z<k)] ] = argmin
  p(zk|z<k) Eq(z<k)
  [ −H(q(zk|z<k))− Eq(zk|z<k) [log p(zk|z<k)] ] = argmin
  p(zk|z<k) Eq(z<k) [KL(q(zk|z<k)||p(zk|z<k))] . (7)
sections.heading: C CONDITIONAL NCE FOR HIERARCHICAL VAES
sections.text: In this section, we describe how we derive the NCE training objective for hierarchical VAEs given in Eq. (4). Our goal is to learn the likelihood ratio between the aggregate conditional q(zk|z<k) and the prior p(zk|z<k). We can define the NCE objective to train the discriminator Dk(zk, z<k) that classifies zk given samples from the previous groups z<k using:
  min Dk − Eq(zk|z<k)[logDk(zk, z<k)]− Ep(zk|z<k)[log(1−Dk(zk, z<k))] ∀z<k. (8)
  Since z<k is in a high dimensional space, we cannot apply the minimization ∀z<k. Instead, we sample from z<k using the aggregate approximate posterior q(z<k) as done for the KL in a hierarchical model (Eq. (7)):
  min Dk
  Eq(z<k) [ − Eq(zk|z<k)[logDk(zk, z<k)]− Ep(zk|z<k)[log(1−Dk(zk, z<k))] ] . (9)
  Since q(z<k)q(zk|z<k) = q(zk, z<k) = Epd(x)[q(z<k|x)q(zk|z<k,x)], we have:
  min Dk
  Epd(x)q(z<k|x) [ − Eq(zk|z<k,x)[logDk(zk, z<k)]− Ep(zk|z<k)[log(1−Dk(zk, z<k))] ] . (10)
  Finally, instead of passing all the samples from the previous latent variables groups to D, we can pass the context feature c(z<k) that extracts a representation from all the previous groups:
  min Dk
  Epd(x)q(z<k|x) [ − Eq(zk|z<k,x)[logDk(zk, c(z<k))]− Ep(zk|z<k)[log(1−Dk(zk, c(z<k)))] ] . (11)
sections.heading: D NVAE BASED MODEL AND CONTEXT FEATURE
sections.text: Context Feature: The base model NVAE (Vahdat & Kautz, 2020) is hierarchical. To encode the information from the lower levels of the hierarchy to the higher levels, during training of the binary classifiers, we concatenate the context feature c(z<k) to the samples from both p(z) and q(z). The context feature for each group is the output of the residual cell of the top-down model and encodes a representation from z<k.
  Image Decoder p(x|z): The base NVAE (Vahdat & Kautz, 2020) uses a mixture of discretized logistic distributions for all the datasets but MNIST, for which it uses a Bernoulli distribution. In our model, we observe that replacing this with a Normal distribution for the RGB image datasets leads to significant improvements in the base model performance. This is also reflected in the gains of our approach.
  E IMPLEMENTATION DETAILS
  The binary classifier is composed of two types of residual blocks as in Fig. 5. The residual blocks use batchnormalization (Ioffe & Szegedy, 2015), the Swish activation function (Ramachandran et al., 2017), and the Squeeze-and-Excitation (SE) block (Hu et al., 2018). SE performs a squeeze operation (e.g., mean) to obtain a single value for each channel. An excitation operation (non-linear transformation) is applied to these values to get per-channel weights. The Residual-Block-B differs from Residual-Block-A in that it doubles the number of channels (C → 2C), while down-sampling the other spatial dimensions. It therefore also includes a factorized reduction with 1×1 convolutions along the skip-connection. The complete architecture of the classifier is:
sections.heading: F PERFORMANCE ON SINGLE GROUP VAE
sections.text: To demonstrate the efficacy of our approach on any off-the-shelf VAE, we, apply our NCE based approach to the VAE in (Ghosh et al., 2020). Note we use the vanilla VAE with Normal prior provided by the authors. The FID for CelebA 64 improves from 48.12 to 41.28. Note that the FID for reconstruction is reported as 39.12. Single group VAEs are known to perform poorly on the task of image geeration, which is reflected in the high FID value of reconstruction.
sections.heading: G NEAREST NEIGHBORS FROM THE TRAINING DATASET
sections.text: To highlight that hierarchical NCP generates unseen samples at test time rather than memorizing the training dataset, Figures 6-7 visualize samples from the model along with a few training images that are most similar
  to them (nearest neighbors). To get the similarity score for a pair of images, we downsample to 64 × 64, center crop to 40 × 40 and compute the Euclidean distance. The KD-tree algorithm is used to fetch the nearest neighbors. We note that the generated samples are quite distinct from the training images.
  Query Image Nearest neighbors from the training dataset .
  Query Image Nearest neighbors from the training dataset .
sections.heading: H ADDITIONAL QUALITATIVE EXAMPLES
sections.heading: I EXPERIMENT ON SYNTHETIC DATA
sections.text: In Fig. 11 we demonstrate the efficacy of our approach on the 25-Gaussians dataset, that is generated by a mixture of 25 two-dimensional Gaussian distributions that are arranged on a grid. The encoder and decoder of the VAE have 4 fully connected layers with 256 hidden units, with 20 dimensional latent variables. The discriminator has 4 fully connected layers with 256 hidden units. Note that the samples decoded from prior p(z) Fig. 11(b)) without the NCP approach generates many points from the the low density regions in the data distribution. These are removed by using our NCP approach (Fig. 11(c)).
  We use 50k samples from the true distribution to estimate the log-likelihood. Our NCP approach obtains an average log-likelihood of −0.954 nats compared to the log-likelihood obtained by sampling from the Gaussian prior, −2.753 nats. We use 20k Monte Carlo samples to estimate the log partition function for the calculation of log-likelihood.
sections.heading: J ADDITIONAL QUALITATIVE EXAMPLES
sections.text: In Fig. 12, we show additional examples of images generated by NVAE (Vahdat & Kautz, 2020) and our NCP-VAE. We use temperature=0.7 for both. Visually corrupt images are highlighted with a red square.