A Contrastive Learning Approach for Training Variational Autoencoder Priors
Authors: Jyoti Aneja
Alex Schwing
Jan Kautz
Arash Vahdat
Venue: NeurIPS 2021 Poster (2021)
Keywords: Variational Autoencoders
Noise Contrastive Estimation
Sampling

Abstract: 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.

A Contrastive Learning Approach for Training Variational Autoencoder Priors
Jyoti Aneja', Alexander G. Schwing's .aan Kautz², Arash Vahdat²
UUniversity of Illinois at Urbana-Champaign, NNVIDI (janeja,, aschwingjeillinoss. "aik, "Jkkutzz.waadatteenidia...c
Abstract
Variational autoencoders (VAEs) are one of the powerful likelihood-basee gen- erative models with applications in many domains. However, they struggle 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 decoded to corrupted images. To tackle this issue, we propose an energy-based prior defined by the product 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 hierarchical VAEs with many latent variable groups Our experiments confirm that the proposed noise contrastive priors improve the generative performance of state-of-the of-the-art 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 improve the expressivity their prior distribution.
base prior distribution and a
Introduction
Variational autoencoders (VAEs)
33
64
are one of the powerful
q(x)
likelihood-based generative models that have applications in image gener-
ation 66 35 62) music synthesis [12, speech generation [54] 66). image
captioning 2
semi-supervised learning [33 40. and representa-
p(a)
tion learning 15 79
Although there has been tremendous progress in improving the expressiv- ity of the approximate posterior, several studies have observed that VAE priors fail to match the aggregate (approximate) posterior 30 66 This
We propose an EBM prior using product of base prior p(z) and reweighting factor r(z). designed to bring p(z) closer to the aggregate posterior (zz).
Figure 1:
phenomenon is sometimes described as holes in the prior, referring to regions the latent space that are not decoded to data-like like samples. Such regions often have a high density under the prior but have 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). 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 have shown promising results. However, they require running iterative MCMC during training which is computationally expensive when the energy function represented by 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 (NcurIPS 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. [ 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 [3], we use noise contrastive estimation (NCE) [2] 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 target distribution (i.e., approximate posterior) vs. samples from a noise distribution (i.e., the base trainable prior). However, since NCE's '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 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
2 Prior to the modern deep learning era, most attempts for building generative models using
EBMs were centered around Boltzmann machines [26] 28 and their "'eep" 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 [33. Pang 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 [0) combines the VAE decoder with an EBM defined
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 cannot be applied to discrete data such text or graphs. In contrast, NCP- 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-VAL VAE are complementary. To avoid MCMC sampling. NCE (22) has recently been used for training a normalizing flow on data distributions
[66. Moreover, Han al. [23] 24) use divergence triangulation to sidesteps MCMC sampling. In contrast, we use NCE to train an EBM prior where a noise distribution easily available through a pre-trained VAE.
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 recent works likelihood ratio trick to the latent space of VAFS. The main difference: the base prior and approximate posterior are trained with the VAE objective rather than the adversarial loss. Adversarial loss has
Adversarial Training: Similar to NCE, generative adversarial networks (GANs) 18
73 that link GANs, defined in the pixels space, to EBMs. We apply the
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.
Among prior works on this problem, VampPrior 72 uses a mixture of encoders to represent the prior. However, this requires storing training data pseudo-data to generate samples at test time. Takahashi et al. [70] use the likelihood ratio estimator to train a simple prior distribution. However test time, the aggregate posterior used for sampling in the latent space.
Prior Hole Problem:
Reweighted Priors: Bauer& & Mnih (5) propose reweighting factor similar to ours, but it is trained via truncated rejection sampling. Lawson et al. [45] introduce energy inngiyddd models (EIMs) that

---

define distributions induced by the sampling processes used by Bauer & Mnih [S] as well as our sampling-importtee importance-resamppaa
(SIR) sampling (called SNIS by Lawson et al. 455). 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 [9] first trains an autoencoder and then fits an autoregressive PixeICNN 778 prior to the latent variables. Albeit impressive results, autoregressive models can be very slow to sample from. Two-stage VAE VAE) [IO] trains a VAE on the data, and then, trains another VAE in the latent space. Regularized autoencoders (RAE) 17 train an autoencoder, and subsequently 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.
Background
We first review VAEs, their extension hierarchical VAEs before discussing the prior hole problem.
Variational Autoencoders: VAES learn a generative distribution p(x. z) = p(z)p(x(zx) where p(z) is prior distribution over the latent variable and p(x[z) is likelihood function that generates the data x given z. VAEs are trained by maximizing a variational lower bound LvAE(x) on the log-likeldhood log p(x) (vAE(x) where
LvvE(x): == (x(t)) [lg - KLL(g())))lx<2))
(1)
Here, q(z|x) is an approximate posterior and KI. is the Kullback-Leibbl divergence. The final training objective is (4VAE)))VEE(x) where (x(x) is the data distribution 39
Hierarchical VAFs (HVAEs): To increase the expressivity of both prior and approximate posterior,
earlier work adapted hierarchical latent variable structure 3 19 68 (74) In HVAEs, the latent
variable is divided into K separate groups,
{z,, Zk} J. The approximate posterior and the
prior distributions are then defined by ((z)x) = I q k x) and p(z)
I p(zz) (zz))
Using these, the training objective becomes
CHVAE(x)= (x)( [log logp(((()]]
Ex(nee) (LL( 1zck [Z<+))]
where q(z<k/x) I Q(zz/zzz x) is the approximate posterior up the (A- 1)th group
Prior Hole Problem: Let q(z) A Epp((>) [x]] denote the aggregate (approximate) posterior. In Appendix B.I we show that maximizing Bpa(x) (XVFF(X)) w.r.t. the prior parameters corresponds to bringing the prior as close as possible to the aggregate posterior by minimizing (L(g(z)]lpaa w.r.t.
p(z). Formally, the prior hole problem refers to the phenomenon that p(z) fails to match q(z).
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, EBMs have shown promising results in representing complex distributions. Motivated by their success, introduce the noise contrastive prior (NCP) pzcp(z) = fr(z)p(z) where p(z) is a base prior distribution, e.g., a Normal, r(z) is reweighting factor, and 2 = f r(z)((() (Z))((z)de dz is the normalization constant. The function r R - IR+ maps the latent variable E IR" to a positive scalar, and can be implemented using neural nets.
The reweighting factor r(z) can be trained using MCMC 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.
For & = 1. the expectation inside the summation is simplitied to (L(g(z) (00022/(x)]pp0 (x)])(())).
3

---

Stage (VAE Training)
Stage (NCE Training)
Test Time Sampling
strla)
((:)
D(4)
pece (+) x(a)r(s)
Palls(sia) lia(all
D(J
p(a)
Encoder
Deconler
Decoder
Binary Clanstter
SRIDD
plan s)
x~
phcce (+)
Reconstructio
Figure 2: NCP- VEE- is trained in two stages. In the first stage, we train a VAE using the original VAE objective. In the second stage, train the reweighting factor r(z) using noise contrastive estimation (NCE). NCE trains a classifier to distinguish samples from the prior and samples from the aggregate posterior. Our noise contrastive prior (NCP) is then constructed by the product of the base prior and the reweighting factor, formed via the classifier. At test time, we sample from NCP using (SIR) or Langevin dynamics (LD). These samples are then passed to the decoder to generate output samples.
4.1 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. To this end, we propose two-stage training algorithm. In the first stage, train the VAE with only the base prior p(z). From Appendix [.I, we know that at the end of training, p(z) is close as possible q(z). In the second stage, we freeze the VAE model including the approximate posterior g( xx), the base prior p(z), and the likelihood p(x[z), and we only train the reweighting factor r(z). This second stage with more expressive distribution of the form pNCP(z) Z) X r(z)p(z). Note that our proposed method is generic it only assumes that we can draw samples from (z) and p(z), which applies any VAE. This proposed training is illustrated in Fig. 2 Next, we present our approach for training r(z).
be thought of as replacing the base distribution p(z)
4.2 Learning The Reweighting Factor with Noise Contrastive Estimation
Recall that maximizing the variational bound in Eq. | with respect to the prior's ss parameters corresponds to closing the gap between the prior and the aggregate posterior by minimizing <L(q(Z)]PPNc((2
with respect to the prior (ncr(z). Assuming that the base p(z) in PNcp(z) fixed after the first stage, (L({(z) PNer (z)) is zero when r(z) = g(z)/p(z). However, since we do not have the density function for qzz,, we cannot compute the ratio explicitly. Instead, in this paper, we propose to estimate rzz) using noise contrastive estimation [22. also known as the likelihood ratio trick, that has been popularized in machine learning by predictive coding (55) and generative adversarial networks (GANs) [8]. Since, we can generate samples from both p(z) and q(2) P we train a binary classifier to distinguish samples from g(z) and samples from the base prior p(z) by minimizing the binary cross-entropy loss
min Ex)](()) [log D(z) Ex) [log(1 D(z))].
(3)
Here, IR" + (0, 1) is binary classifier that generates the classification prediction probabilities. Eq. is minimized when D(z)
(())+p(z) Denoting the classifier at optimality by D* (z), we
9(2)
estimate the reweighting factor r(z) g(z) (z)
p(*) ~ (=) The appealing advantage of this estimator is that it is obtained by simply training a binary classifier rather than using expensive MCMC sampling.
Intuitively. if p(z) is very close to q(z) (i.c., p(z) = q())), the optimal classifier will have a large loss value in Eq. (3). and we will have r(z) R 1. If p(z) is instead far from g(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 qz)), then the binary classifier can learn the ratios successfully.
We generate samples from the aggregate posterior g(z) = (x)/()/(((()))))))))) via ancestral sampling: draw data from the training set (x pa(x)) and then sample from g(2(x).

---

4.3
Test Time Sampling
To sample from VAE with an NCP, we first generate samples from the NCP and pass them to the decoder to generate output samples (Fig. 2). We propose two methods for sampling from NCPs.
(SIR): We first generate M samples from the base prior dis- tribution [z(m²) J m==1 ~ p(z). We then resample one of the M proposed samples using importance weights proportional to mm(m) = pNcp(z) (m)) / //p(z((m) = r(2(m)). The benefit of this technique: both the samples
Sampling-mmportace-Resaaa
proposal generation and the evaluation of
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 sample zo by drawing from p(z) and update the sample iteratively using: Zt++1 2, 0.5 XV 2 E(z) + Vλe, where 6, NN(,, 1) and is the step size. LDD is run for a finite number of iterations, and in contrast to SIR, it is slower given its sequential form.
4.4 Generalization to Hierarchical VAEs
The state-of-the-art VAEs 99| 74] hierarchical g(zix) and p(z). Here p(z) is chosen to be
B.2 shows that training a HVAE encourages the prior to minimize ]=x]] for each conditional, where g(z<x) _ Epats))((z (g(*= xx] is the aggregate posterior up to the 1)th group, and (<<//Z<e)) - Exp(x) [g( (z( Z x)] is the aggregate conditional for the hthh group. Given this observation, extend NCPs to hierarchical models to match each conditional in the prior with g(z z<x). Formally, we define hierarchical NCPs by
Gaussian distribution. Appendix B.2 Ey(xes) (2)(((=))=))lllll
P2cr(z) = 2 I r(zx|zee )p(z)) (Zck) where each factor is an EBM. PNcr(z) resembles EBMs with autoregressive structure among groups (53)
In the first stage, we train the HVAE with prior IT p(zzlzek), For the second stage, we use K binary classifiers, each for a hierarchical group. Following Appendix we train each classifier via
min Epa(x)((acx)))) E flog Dz c(ZZR))) - Ep Da(zs, cl De
(4)
where the outer expectation samples from groups up to the (A: 1)th group, and the inner expectations sample from approximate posterior and base prior for the hthh group, conditioned on the same Zzk-- The discriminator Dr classifies samples 2₃ while conditioning its prediction on Zzz using shared context feature c(zzk).
The NCE training in Eq. (4) is minimized when Da (R(ZE, c(zzz)) p Denot-
giz, kct)
ing the classifier at optimality by (2 c(z<k)), we obtain the reweighting factor r(zz| R ([(,).((<<)) I- (( 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(zzk) extracts a representation from Zęk- Instead of learning a new representa- tion stage two, simply the representation that is extracted from Z.c in the hierarchical prior, trained the first stage. Note that the binary classifiers are trained in parallel for all groups
5
Experiments
In this section, we situate NCP against prior work several commonly used single group VAE models in Sec. 3.I In Sec. 5.2, we present our main results where we apply NCP to hicrarchical NVAE 74 to demonstrate that approach can be applied to large scale models successfully.
In most of our experiments, we measure the sample quality using the Fréchet Inception Distance (FID) score 25 with 50,000 samples, as computing the log-likelihood requires estimating the intractable normalization constant. For generating samples from the model, we use SIR with 5K proposal samples. To report log-likelihood results, we train models with small latent spaces on the dynamically binarized MNIST [46] dataset. We intentionally limit the latent space to ensure that we can estimate the normalization constant correctly.
5

---

Table Comparison with two-stage VAEs on CelebA-64 64 with RAE [17 networks. Results reported by Ghosh et al. 17
Table 2: Likelihood results on MNIST on single latent group model with ar- chitecture from LARS ( & SNIS [4 (results in nats). We closely follow the training hyperparameters used by Law- son al. [45]
Model
FID
VAE w/ Gaussian prior
48.12
2s-VAR (T0
49.70
WAE 71
42.73
Model
NLL
RAE 17
40.95t
VAE w/ Gaussian prior
84.82
NCP w/ Gaussian prior as base
41.28
VAE w/ LARS prior 5
83.03
NCP w/ GMM prior base
39.00
VAE w/ SNIS prior (45)
82.52
Base VAE-Recon
36.01
NCP-VAE
82.82
5.1 Comparison to Prior Work
In this section, we apply NCP to several commonly used small VAE models. Our goal, here, is to situate our proposed model against (i) two-stage VAE models that train (variational) autoencoder first, and then, fit a prior distribution (Sec. 5.1.1), and (ii) VAEs with reweighted priors (Sec. (51.2).
To make sure that these comparisons are fair, follow exact training setup and network architectures from prior work as discussed below.
5.1.1 Comparison against Two-Stage VAES
Here, we show the generative performance of our approach applied to the VAE architecture in RAE [17] on the CelebA-64 dataset 48 We borrow the exact training setup from 117 and implement our method using their publicly available code Note that this VAE architecture has only one latent variable group. The same base architecture
used in the implementation of 2s- VAE (T0) and WAE [71. In order to compare our method to these models, we use the reported results from RAE 17.. We apply our NCP-VAE on top of both vanilla VAE with a Gaussian prior and a 10- component Gaussian mixture model (GMM) prior that was proposed in RAES. As we can see in Tab. D our NCP-VAE improves the performance of 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. We also report the FID score for reconstructed images using samples from the aggregate posterior qz)) instead of the prior. Note that this value represents the best FID score that can be obtained by perfectly matching the prior to the aggregate posterior in the second stage. The high FID score of 36.01 indicates that the small VAEs cannot reconstruct data samples well due to the small network architecture and latent space. Thus, even with expressive priors, FID for two-stage VAEs are lower bounded by 36.01 in the 2nd stage.
5.1.2 Comparison against Reweighted Priors
LARS S and SNIS 45 train reweighted priors similar to our EBM prior. To compare NCP-VAE VAE against these methods, we implement our method using the VAE and energy-function networks from [45]. We closely follow the training hyperparameters used in [45] as well as their approach for obtaining a lower bound on the log likelihood (i.e., the SNIS objective in 45) provides lower bound on data likelihood). As shown in Tab. NCP-VAE obtains the negative log-likelihood (NLL)
of 02.82, comparable to Lawson et 45), while outperforming LARS (5). Although NCP-VAR VAE slightly inferior to SNIS on MNIST, it has several advantages as discussed in Sec.
5.1.3 Training using Normalizing Flows
Chen et al. 8 (Sec. 3.2) show that a normalizing flow in the approximate posterior is equivalent to having its inverse in the prior. The base NVAE uses normalizing flows in the encoder. As a part of VAE training, prior and aggregate posterior are brought close, i.e., normalizing flows are implicitly used. We argue that normalizing flows provide limited gains to address the prior-hole problem (see Fig. by Kingma et al. [1 Yet, our model further improves the base VAE equipped with normalizing flow.

---

4
7
7
0
6
Y
6
/
5
8
3
2
6
0
6 8 5 9 8
6
7
7
/
7
8
9 4 8 7 6
2
5
4
1
0
7
4
7
2
2
0
2
8
9
4
8
3
0
4
0
7
6
7
3
7 /
2
y
0
4
6
6
(a)
MNIST (L
= 1.0)
(b) CIFAR-10 (t 0.5)
(c) CelebA 64x64 (t 0.7)
(d) CelebA HQ 256 0.7)
Figure Randomly sampled images from NCP- with the temperature t for the prior.
Table 3: Generative performance on
Table 4: Generative performance on CIFAR- 110
CelebA-64.
Model
FIDL
Model
FID|
NCP- -APE (ours)
24.08
NCP-VAE (ours)
5.25
VAEBM 80
12.96
VAEBM 80
5.31
NVAE 74
51.71
NVAE 774
13.48
RAE 17
74.16
RAE [17]
40.95
2s-VAE 10
72.9
2s- VAE 10
44.4
Perceptial AE 82
51.51
WAE 71
35
EBM [3
40.58
Perceptial 82
13.8
Latent EBM
57
70.15
Latent EBM [57]
37.87
Style-GANv2 36
3.26
COCO-GAN 47
4.0
DDPM 20
3.17
QA-GAN 58
6.42
Score SDE 69
3.20
NVAE-Rccooo 74
1.03
NVAE-R Recon 74]
2.67
5.2
Quantitative Results Hierarchical Models
In this section, we apply NCP to the hicrarchical VAE model proposed in NVAE 77 We examine NCP-VAE on four datasets including dynamically binarized MÑIST [46), CIFAR-10 43 CelebA-64
and CelebA- HH---HQ--666 -56 34 For CIFAR-10 and CelebA-64 the model has 30 groups, and for CelebA-HQ it has 20 groups For MNIST, we train an NVAE model with small latent space on 4 latent variables. The small latent space allows us estimate the partition function confidently (std. of log Z estimation < 0.23). The quantitative results are reported Tab. 3 Tab. 4 Tab. 5 and Tab. 6 On all four datasets, our model improves upon NVAE, and 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 improves the NVAE FID of 51.71 to 24.08. On MNIST, although our latent space is much smaller, model outperforms previous VAEs. NVAE has reported 78.01 nats on this dataset with a larger latent space.
48
MNIST with 10 groups
On CIFAR-10 and CelebA-HO H00-AA-HQ-26. recently proposed VAEBM 80 outperforms our NCP. VAE, However, we should note that (i) NCP-VAE and VAEBM are complementary to each other, NCP. VAE targets the latent space while VAEBM forms an EBM on the data space. We expect improvements by combining these two models. (ii) VAEBM assumes that the data lies on continuous space whereas NCP- VAE does not make any such assumption and it can be applied to discrete data (like binarized MNIST Tab. (6. graphs, and text. (iii) NCP.VAE is much simpler to setup as it involves training a binary classifier whereas VAEBM requires MCMC for both training and test.
7

---

Table 5: Generative results on CelebA-HQ-256.
Table 6: Likelihood results on MNIST in nats.
Model
FID
Model
NLL,
NCP- VAE (ours)
24.79
NCP-VAE (ours)
78.10
VAEBM 80
20.38
NVAE-small 74
78.67
NVAE 74
40.26
BIVA [50]
78.41
GLOW 38
68,93
DAVE++ 76
78.49
Advers. LAE 59
19.21
IAF-VAE
41
79.10
PGGAN [3]
8.03
VampPrior AR dec.
72
78.45
WVAE-Recon
0.45
DVAE
65
80.15
5.3 Qualitative Results
We visualize samples generated by NCP-VAI with the NVAE backbone in Fig. without any manual intervention. We adopt the common practice of reducing the temperature of the base prior p(z) by scaling down the standard-deviate observe that re-adjusting the nonch-nomalliztiion (BN), given a temperature applied the prior, improves the generative quality. Similarly, we achieve diverse, high-quality images by re-adjusting the BN statistics as described by (74). Additional qualitative results are shown in Appendix G
of the conditional Normal distributions (38) 66 74] also
Nearest Neighbors from the Training Dataset:
To highlight that hierarchical NCP generates unseen samples at test time rather than memorizing the training dataset, Figures 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 pair of images, we downsample to 64 x 64, center crop 40 x 40 and compute the Fuclidean distance. The KD-tree algorithm used 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
Figure 4: Query images (left) and their nearest neighbors from the CelebA-HQ -Q-226 training dataset.
5.4 Additional Ablation Studies
We perform additional experiments study i) how hierarchical NCPs perform as the number of latent groups increases, ii) the impact of SIR and LD byperparameters, and iii) what the classification loss in NCE training conveys about p(z) and qz)) All experiments are performed CelebA-64.
Number of latent variable groups: Tab. 7 shows the generative performance of hierarchical NCP with different amounts of latent variable groups As we increase the number of groups, the FID
"Lowering the temperature only used to obtain qualitative samples, not for the quantitative results in Sec.

---

Binary Classifier for CelebA
Effective Sample Loss
70
Effective Sample.
40
30
0.6
30
0.4
20
0.2
0.0
20
60 Training Epoch
40
80 100
0.1
0.2
0.3
0.5
0.7
Binary Classifier (b)
(a)
Figure 5: (a) Classification loss for binary classifiers on latent variable groups. A larger final loss upon training indicates that q(2) and p(z) are more similar. (b) The effective sample size the final loss value at the end of training. Higher effective sample size implies similarity of two distributions.
score of both NVAE and model improves. This shows the efficacy of NCPs, with expressive hierarchical priors in the presence of many groups
SIR and LD parameters: The computational complexity of SIR is similar to LD if we set the number of proposal samples in SIR equal to the number LD iterations. In Tab. 8 we observe that increasing both the number of proposal samples STR and the LD iterations leads to noticeable improvement in FID score. For SIR, the proposal generation and the evaluation of r(z) ) are parallelizable. Hence, as shown in Tab. 8 image generation is faster with SIR than with LD. However, GPU
Table 7: # groups & genera- tive performance in FID.,
# groups NVAE NCP-VAE
6
33.18
18.68
15
14.96
5.96
30
13.48
5.25
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 (z)). Denoting the classification loss in Eq. (3) at optimality by L. Goodfellow et al. 188 show that (p(())](q(())) = log 2 0.5 x L where JSD denotes the Jensen-Shannon divergence between two distributions. Fig. ((a) plots the classification loss (Eq. 4 for each classifier for 15-group NCP trained on the CelebA-64 dataset. Assume that the classifier loss the end of training is good approximation of L We observe that 8 out of 15 groups have L 2 0.4, indicating a good overlap between p(z) and g(z) for those groups. To further assess the impact of the distribution match SIR sampling, in Fig. ()b), size (ESS
visualize the effecttve sample F in SIR vs. L for the same group. We observe strong correlation between L and the effective sample size. SIR is more reliable the same 8 groups that have high classification loss. These groups are at the top of the NVAE hierarchy which have been shown to control the global structure of generated samples (see B.6 in (74)
Table 8: Effect of SIR sample size and LD iterations. Time-N is the time used to generate a batch of N images.
# SIR proposal samples
Time-1 (sec)
Time- 10 (sec)
Memory (GB)
LD iterations
Time- I (sec)
Time- 10
Memory (GB)
FID
FID
(sec)
5
11.75
0.34
0.42
1.96
5
14.44
3.08
3.07
1.94
50
8.58
0.40
1.21
4.30
50
12.76
27.85
28.55
1.94
500
6.76
1.25
9.43
20.53
500
8.12
276.13
260.35
1.94
5000
5.25
10.11
95.67
23.43
1000
6.98
552
561.44
1.94
Analysis of the re-weighting technique: To show that samples from NCP (PxcP(2) (z)) are closer to the aggregate posterior q(z) compared to the samples from the base prior p(z). we take 5k samples from g(z), p(z), and PNCP(2) (z) at different hierarchy/group levels. Samples are projected to lower dimension (d= 500) using PCA and populations are compared via Maximum Mean Discrepancy (MMD), Consistent with Fig. 5(a),
Table 9: MMD comparison.
# group
(g. p) (g. PNCP)
5
0.002
0.002
10
0.08
0.06
0.08
0.07
$ESS measures reliability of SIR via 1/ Σ,, (m)( where (m)) = ((z(m))) /mm r()(mm)) [56].

---

Tab. shows that groups with lower classification loss had mismatch between and q. and NCP is able to reduce the dissimilarity by re-weighting.
Conclusions
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 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 's expressivity. We also showed how NCP training scales to large hierarchical VAEs, as it can be done in parallel simultancously for all the groups. Finally, we demonstrated that NCPs improve the generative performance of small single group VAEs and ofate-of oftee-----h-rrr NVAEs by a large margin.
Impact Statement
The main contributions of this paper towards tackling fundamental issue with training VAE models - the prior hole problem. The proposed method increases the expressivity of the distribution used sample latent codes for test-time image generation, thereby increasing the quality (sharpness) and diversity of the generated solutions. Therefore, ideas from this paper could find applications in VAE-based content generation domains, such computer graphics, biomedical imaging, computation fluid dynamics, among others. More generally, we expect the improved data generation to be beneficial for data augmentation and representation learning techniques.
When generating new content from a trained VAE model, one must carefully assess if the sampled distribution bears semblance to the real data distribution used for training, in terms of capturing the different modes of the real data, as well as the long tail. A model that fails to achieve real data distribution result should be considered biased and corrective steps should be taken to proactively address it. Methods such as NCP- VAE, that increase the prior expressivity, hold the promise to reduce the bias in image VAEs. Even so, factors such as the VAE architecture. training hyper-paan pareeeerrreeeers, and temperature for test-time generation, could impact the potential for bias and ought to be given due consideration. We recommend incorporating the active research into bias correction for generative modeling 200 into any potential applications that use this work.
8 Acknowledgements
The authors would like thank Zhisheng Xiao for helpful discussions. They also would like to extend their sincere gratitude to the NGC team at NVIDIA for their compute support.
Funding Transparency Statement
This work mostly funded by NVIDIA during internship. It also partially supported by NSF under Grant #1718221, 2008387, 2045586, 2106825, MRI #1725729, and NIFA award 2020-67021-32799.
References
[1] Ackley, D. H., Hinton, G. E., and Sejnowski, T. J. Cognitive science, 1985.
[2] Aneja, J., Deshpande, A., and Schwing, A. G. Convolutional image captioning. In Proceedings of the IEEE conference on computer vision and pattern recognition (CVPR), 2018.
[3] Aneja, J., Agrawal, H., Batra, D., and Schwing, A. Sequential latent spaces for modeling the intention during diverse image captioning. In Proceedings of the IEEE International Conference on Computer Vision (ICCV), 2019.
[4] Azadi, S., Olsson, C., Darrell, T., Goodfellow, I., and Odena, A. Discriminator rejection sampling. In International Conference Learning Representations (ICLR), 2018.
10

---

[5] Bauer, M. and Mnih, A. Resampled priors for variational autoencoders. In The 22nd Interna- tional Conference on Artificial Intelligence and Statistics (AISTATS), 2019.
[6] Brock, A., Donahue, J., and Simonyan, K. Large scale gan training for high fidelity natural image synthesis. arXiv:1809, 1006, 2018.
[7] Che, T., Zhang, R., Sohl-Dickstein, J., Larochelle, H., Paull, L., Cao. Y., and Bengio, Y. Your GAN is secretly an energy- based model and you should use discriminator driven latent sampling. arXiv:2003.06060 2020.
[8] Chen, X., Kingma, D. P., Salimans, T., Duan, Y., Dhariwal, P.. Schulman, J., Sutskever, I., and Abbeel, P. Variational lossy autoencoder. In International Conference Learning Representations (ICLR), 2017.
[9] Child, R. Very deep (vac}s generalize autoregressive models and outperform them images. In International Conference on Learning Representations (ICLR), 2021.
[10] Dai, B. and Wipf, D. Diagnosing and enhancing models. In International Conference Learning Representations (ICLR), 2018.
[11] Deshpande, A., Aneja, J., Wang, I.., Schwing, A. G., and Forsyth, D. Fast, diverse and accurate image captioning guided by part-of-speech. In Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition (CVPR), 2019.
[12] Dhariwal, P. Jun, H., Payne, C., Kim, J. W. Radford, A., and Sutskever, I. Jukebox: A generative model for music. arXiv:2005, 00041, 2020.
[13] Du, Y. and Mordatch, I. Implicit generation and modeling with energy based models. In Advances Neural Information Processing Systems, 2019.
[14] Engel, J., Hoffman, M., and Roberts, A. Latent constraints: Learning to generate conditionally from unconditional generative models. In International Conference on Learning Representations (ICLR), 2018.
[15] Fortuin, V., Hüser, M., Locatello, F., Strathmann, H., and Rätsch, G. Som-vae: Interpretable discrete representation learning time series. In International Conference Learning Representations (ICLR), 2018.
[16] Gao, R., Nijkamp, E., Kingma, D. P., Xu, Z., Dai, A. M., and Wu, Y. N. Flow contrastive estimation of energy-based based models. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR), 2020.
[17] Ghosh, P., Sajjadi, M. S. M., Vergari, A., Black, M., and Scholkopf, B. From variational to deterministic autoencoders. In International Conference Learning Representations (ICLR), 2020.
[18] Goodfellow, I. J., Pouget-Abadie, Mirza, M., Xu, B., Warde-Farley, D., Ozair, S., Courville, A., and Bengio, Y. Generative Adversarial Networks. 1rXiv::406.222 2661, 2014.
[19] Gregor, K., Besse, F., Rezende, D. J., Danihelka, and Wierstra, D. Towards conceptual compression. In Advances In Neural Information Processing Systems, 2016.
[20] Grover, A., Song, J., Kapoor, A., Tran, K., Agarwal, A., Horvitz, E. J., and Ermon, S. Bias correction of learned generative models using likelihood-frie importance weighting. In Advances in Neural Information Processing Systems, 2019.
[21] Gulrajani, I., Kumar, K., Ahmed, F., Taiga, A. A., Visin, F., Vazquez, D., and Courville, A. PixeIVAE: A latent variable model for natural images. arXiv: 1.05013, 2016.
[22] Gutmann, M. and Hyvärinen, A. Noise-contrastive estimation: A new estimation principle for unnormalized statistical models. In International Conference on Artificial Intelligence and Statistics (AISTATS), 2010.

---

[23] Han, T., Nijkamp, E., Fang, X., Hill, M.. Zhu. S.-C., and Wu, Y. N. Divergence triangle for joint training of generator model, energy-banc based model, and inferential model. In Proceedings of the HEEE Conference on Computer Vision and Pattern Recognition (CVPR), 2019.
[24] Han, T., Nijkamp, E.,, Zhou, L.,, Pang, B., Zhu, S.-C., and Wu, N. Joint training of variational auto-encoder and latent energy-based model. In Proceedings of the IEEE/CVE Conference on Computer Vision and Pattern Recognition (CVPR), 2020.
[25] Heusel, M., Ramsauer, H., Unterthiner, T., Nessler, B., and Hochreiter, Gans trained by a two time-scale update rule converge to local nash equilibrium. In Advances in neural information processing systems, 2017.
[26] Hinton, G. E. Training products of experts by minimizing contrastive divergence. Neural compulation, 2002.
[27] Hinton, G. E., Sejnowski, T. J., et Learning and releaming in boltzmann machines. Parallel distributed processing: Explorations in the microstructure of cognition, 1986.
[28] Hinton, G. E., Osindero, S., and Teh, Y.-W. A fast learning algorithm for deep belief nets. Neural computation, 2006.
[29] Ho, J., Jain, A., and Abbeel, P. Denoising diffusion probabilistic models. arxiv:2006. 2020.
[30] Hoffman, M. D. and Johnson, M. J. Elbo surgery: yet another way to carve up the variational evidence lower bound. In Workshop in Advances in Approximate Bayesian Inference, Neural Information Processing Systems, 2016.
[31] Hu, J., Shen, L., and Sun, G. Squeeze-ande ande--ndd-eee andeexciatiion networks. In Proceedings of the IEEE conference on computer vision and pattern recognition (CVPR), 2018.
[32] loffe, S. and Szegedy, C. Batch normalization: Accelerating deep network training by reducing internal covariate shift. In International Conference Machine Learning (ICML), 2015.
[33] Izmailov, P., Kirichenko, Finzi, M., and Wilson, A. G. Semi-supervised learning with normalizing flows. In ICML, 2020.
[34] Karras, T., Aila, T., Laine, S., and Lehtinen, J. Progressive growing of gans for improved quality, stability, and variation. X7vi 10196, 2017.
[35] Karras, T., Laine, S., and Aila, T. A stylc-based generator architecture for generative adversarial networks. In Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition (CVPR), 2019.
[36] Karras, T., Aittala, M., Hellsten, J., Laine, S., Lehtinen, J.. and Aila, T. Training generative adversarial networks with limited data. arXiv:2006.06676 2020.
[37] Kingma, D. and Ba, J. Adam: A method for stochastic optimization. arXiv:1412.6980 2014.
[38] Kingma, D. P. and Dhariwal, P. Glow: Generative flow with invertible 1x1 convolutions. In Advances Neural Information Processing Systems 31, 2018.
[39] Kingma, D. P. and Welling, M. Auto-encoding variational bayes. In The International Confer- ence on Learning Representations (ICLR), 2014.
[40] Kingma, D. P., Mohamed, S., Rezende, J., and Welling, M. Semi-supervised learning with deep generative models. In Advances in Neural Information Processing Systems, 2014.
[41] Kingma, D. P., Salimans, T., Jozefowicz, R., Chen, X., Sutskever, I., and Welling, M. Improved variational inference with inverse autoregressive flow. In Neural Information Processing Systems, 2016.
[42] Klushyn, A., Chen, N., Kurle, R., Cseke, B., and van der Smagt, Learning hierarchical priors in vacs. In Advances in Neural Information Processing Systems, 2019.
12

---

[43] Krizhevsky, A., Hinton, G., et al. Learning multiple layers of features from tiny images, 2009.
[44]
Larochelle, H. and Bengio, Y. Classification using discriminative restricted boltzmann machines. In Proceedings of the 25th international conference on Machine learning, pp. 536-543, 2008.
[45] Lawson, J., Tucker, G., Dai, B., and Ranganath, R. Energy-inspired models: Learning with sampler-induced distributions. In NeurIPS, 2019.
[46] LeCun, Y. The mnist database of handwritten digits. http://ynnn. lecun, com/exib/mnist/, 1998.
[47] Lin, C. H., Chang, C.-C., Chen, Y.S.. Juan, D.-C., Wei, W., and Chen, H.-T. Coco-gan: generation by parts via conditional coordinating. In Proceedings of the IEEE International Conference Computer Vision (ICCV), 2019.
[48] Liu, Z., Luo, P., Wang, X., and Tang. X. Deep learning face attributes in the wild. In Proceedings of the IEEE international conference on computer vision (ICCV), 2015.
[49] Loshchilov, and Hutter, F. Sgdr:
Stochastic gradient descent with warm restarts.
arXiv:1608.03983, 2016.
[50] Maaløe, L., Fraccaro, M., Liévin, V., and Winther, O. BIVA: A very deep hierarchy of latent variables for generative modeling. In Advances in neural information processing systems, 2019.
[51] Makhzani, A., Shlens, J., Jaitly, N., Goodfellow, I., and Frey, B. Adversarial autoencoders. arXiv:/511.05644, 2015.
[52] Mescheder, L.. Nowozin, S., and Geiger, A. Adversarial variational bayes: Unifying variational autoencoders and generative adversarial networks. International Conference on Machine Learning (ICML), 2017.
[53] Nash, C. and Durkan, C. Autoregressive energy machines. arXiv:1904.055: 2019.
[54] Oord, A. d., Dieleman, S., Zen, H., Simonyan, K., Vinyals, O., Graves, A., Kalchbrenner, N., Senior, A., and Kavukcuoglu, K. Wavenet: A generative model for raw audio. arXiv:1609.033999 2016.
[55] Oord, A. d., Li, Y., and Vinyals, O. Representation learning with contrastive predictive coding. TrXivv1807.083 2018.
[56] Owen, A. B. Monte Carlo theory, methods and examples. 2013.
[57] Pang, B., Han, T., Nijkamp, E., Zhu, S.-C., and Wu, Y. N. Learning latent space energy-based prior model. arXiv:2006.08200 2020.
[58] Parimala, and Channappayya. Quality aware generative adversarial networks. In Advances in Neural Information Processing Systems, 2019.
[59] Pidhorskyi. S., Adjeroh, D. A., and Doretto, G. Adversarial latent autoencoders. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR), 2020.
[60] Ping. W., Peng, K., Zhao, K., and Song, Z. Waveflow: A compact flow-based model for raw audio. ICML, 2020.
[61] Ramachandran, P., Zoph, B., and Le, Q. V. Searching for activation functions. arXiv: 2017.
[62] Razavi, A., van den Oord, A., and Vinyals, Generating diverse high-fidelity images with vq-vae-2. In Advances in Neural Information Processing Systems, 2019.
[63] Rezende, D. J. and Viola, F. Taming vaes. arXiv: 180000097 2018.
[64] Rezende, D. J., Mohamed, S., and Wierstra, D. Stochastic backpropagation and approximate inference in deep generative models. In International Conference on Machine Learning (ICLR), 2014.
[65] Rolfe, J. T. Discrete variational autoencoders, arXhv:1609.02200.0 2016.
13

---

[66] Rosca, M., Lakshminarayanan, B., and Mohamed, S. Distribution matching in variational inference. arX1v.1802.06847, 2018.
[67] Salakhutdinov, R. and Hinton, G. Deep boltzmann machines. In Artificial intelligence and stutistics, 2009.
[68] Sønderby, C. K., Raiko, T., Maaløe, L.., Sønderby, S. K., and Winther, O. Ladder variational autoencoders. In Advances in neural information processing systems, 2016.
[69] Song, Y., Sohl-Dickstein, J., Kingma, D. P., Kumar, A., Ermon, S., and Poole, B. Score-based generative modeling through stochastic differential equations. In International Conference on Learning Representations (ICLR), 2021.
[70] Takahashi, H., Iwata, T., Yamanaka, Y., Yamada, M., and Yagi, S. Variational autoencoder with implicit optimal priors. Proceedings of the AAAI Conference on Artificial Intelligence, 2019.
[71] Tolstikhin, I., Bousquet, O., Gelly, S., and Schölkopf. B. Wasserstein auto-encoders. In International Conference Learning Representations (ICLR 2018). OpenReview. net, 2018.
[72] Tomezak, J. and Welling, M. Vae with vampprior. In International Conference Artificial Intelligence and Statistics, pp. 1214-121 1223, 2018.
[73] Turner, R., Hung, J., Frank, E., Saatchi, Y., and Yosinski, J. Metropolis-hastine generative adversarial networks. In International Conference on Machine Learning. PMLR, 2019.
[74] Vahdat, A. and Kautz, J. NVAE: deep hierarchical variational autoencoder. In NeurIPS, 2020.
[75] Vahdat, A., Andriyash, E., and Macready, W. G. DVAE#: Discrete variational autoencoders with relaxed Boltzmann priors. In Neural Information Processing Systems, 2018.
[76] Vahdat, A., Macready, W. G., Bian, Z., Khoshaman, A., and Andríyash, E. DVAE++ Discrete variational autoencoders with overlapping transformations. International Conference on Machine Learning (ICML), 2018.
[77] Vahdat, A., Andriyash, E., and Macready, W. G. Undirected graphical models as approximate posteriors. In International Conference on Machine Learning (ICML), 2020.
[78] Van Den Oord, A., Kalchbrenner, N., and Kavukcuoglu, K. Pixel recurrent neural networks. In Proceedings of the 33rd International Conference on International Conference on Machine Learning, 2016.
[79] Van Den Oord, A., Vinyals, O., et al. Neural discrete representation learning. In Advances in Neural Information Processing Systems, 2017.
[80] Xiao, Z., Kreis, K., Kautz, J., and Vahdat, A. (VAEBM): A symbiosis between variational au- toencoders and energy-based models. In International Conference Learning Representations (ICLR), 2021.
[81] Xu, H., Chen, W., Lai, J., Li, Z., Zhao, Y., and Pei, D. On the necessity and effectiveness of learning the prior of variational auto-encoder. arXiv: 13452, 2019.
[82] Zhang, Z.,, Zhang, R., Li, Z., Bengio, Y., and Paull, L. Perceptual generative autoencoders. In International Conference on Machine Learning (ICML), 2020.
14


Reviews:
[{"review_id": "javgySy4kGN", "reviewer": "Reviewer_C3kP", "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. ", "questions": "", "limitations": "", "rating": 6, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "Pros:\n- The idea is very simple and easy to implement.\n- The method can be plugged with many different VAE variants.\n- The method seems to improve sample quality as measured by FID across the board.\n- The paper provides quite thorough empirical validations, on many different datasets, with many different architectures.\n- The ablations, while mainly confirming intuitions, are informative.\n- The paper is well written, and most claims are well substantiated.\n\nCons:\n- The method cannot reliably be used to evaluate a bound on the log-likelihood.\n- The only performance metric used is the FID, which arguably favors quality of samples over generalization (as mentionned in (Razavi et al. 2019). As a consequence, I am not sure that the claim that 'diversity' of samples is improved (l.319) is well substantiated. Evaluating performance using a more diverse set of metrics, and notably the metrics used in (Razavi et al. 2019) could help in substantiating this claim.\n\nOther comments:\n- I don't understand the statement l.240. Fig 1. from Kingma et al. 2016 seems to indicate that IAF exactly attend to the prior holes problem.\n- l294. Could you also give the number of proposal samples, so that we can compare it to the effective sample size?\n\nOverall, I think the pros outweigh the cons. This is a simple and interesting contribution. I recommend acceptance.\n\nReferences:\n_Generating Diverse High-Fidelity Images with VQ-VAE-2_ Ali Razavi, Aäron van den Oord, Oriol Vinyals", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "The authors provide a statement assessing the potential societal impact of their work.", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "3"}, "extra_text": {}}, {"review_id": "bcBsBAGO9x", "reviewer": "Reviewer_a91M", "summary": "The authors propose a way to improve the performance of a trained VAE by modifying the learned prior with a reweighting factor. The reweighting factor is computed using a network trained to discriminate between samples from the learned prior and samples from the encoder given real data. To sample from the VAE, the latent variables are sampled from the reweighted prior using MC inference (e.g. sampling-importance-resampling) and fed through the decoder. This leads to better samples (versus not using reweighting) but also greater computational cost due to the MC inference.", "questions": "", "limitations": "", "rating": 7, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "# Novelty\nThe authors situate their work well in relation to prior work. Other (cited) papers use reweighted priors, so the novelty comes mainly from using noise contrastive estimation to train the reweighting factor (via learning to discriminate between samples from the prior and from the encoder). This is a significant improvement as it makes training simpler and scalable to larger VAE architectures.\n\n# Strengths\n- Clearly written and a pleasure to read.\n- Thorough experimental comparisons and ablation studies.\n- Applicable to modern hierarchical VAE architectures, unlike much prior work.\n\n# Concerns\n- Sampling from the VAE with the suggested parameters (e.g. using 5000 proposal samples) is very slow, taking 10s to create a single output (based on Table 8). This might limit the practical applicability of this approach. It would also be interesting to see these times compared to those for the VAEBM baseline, which also involves MC inference.\n- The experimental results are not overly convincing, with the reported FIDs/NLLs often being similar to, or worse than, the baselines. However, given that the authors address this (e.g. pointing out that VAEBM baseline is more complex and not applicable to discrete data), I think the results are good enough for acceptance.", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "Yes.", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "4"}, "extra_text": {}}, {"review_id": "G0jKgkW4Bn", "reviewer": "Reviewer_wa7R", "summary": "Update after Response:\n\nThe response of the authors resolved some of my concerns and I am therefore improving my score by 1 point, provided that the authors add the additional information mentioned in their reply. I think the work is relevant and presents a potential solution to solve some of the problems that plague autoencoder sample generation. \n\nOld:\nThe authors propose a novel method of adapting the prior p(z) of VAE to the experimentally observed approximate posterior q_{\\theta}(z) over the training samples. The modification of the prior is performed as a separate step, after model training. This improves the visual quality of generated samples drawn from the modified prior. The article includes an extensive evaluation and the appendix includes further results with important studies.", "questions": "", "limitations": "", "rating": 7, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "Strengths:\n- very clearly written method\n- simple, straightforward method\n- detailed evaluation on multiple datasets and with multiple baselines\n- High visual quality of qualitative samples shown.\n- topic highly relevant to the neurips community\n\nWeaknesses:\n- implementation not accessible\n- On initial reading, I feared that by adapting the prior to the aggregate approximate posterior of the training dataset, we are teaching the model memorization of the training set. The nearest-neighbor evaluation given in Appendix F has convinced me of the opposite. Sad, that I had to reach to the appendix for this very relevant experimental result.\n- evaluation by FID: What has been used as the comparison data distribution? If the training set is used, I'm not surprised that this method achieves high FID scores, as the NCE training step is focused on aligning the prior to the approximate posterior of the training dataset. How do the FID metric on the validation/test set compare?\n- In the experiments, it is unclear how the log partition function is computed. It is well known that for many estimators, a small standard deviation is not a reliable estimate for correctness of the mean since the distribution of estimates can be highly skewed.\n- I feel the use of the t-value during image generation, e.g. Figure 3, undermines the message of the paper. If it is beneficial to accurately represent the marginal posterior, why would we then sample from a different distribution than the marginal posterior?  Also, it is unclear how the temperature was selected.  \n\nCorrectness:\nThe experimental evaluation is mostly appropriate to determine the correctness of the work. However, details are unclear.\n\n\nClarity: Is the paper well written?\n- The introduction and method sections are very clearly written\n- The Experimental section could explain the total model size, used compute, and evaluation metrics used in more detail/higher clarity.\n- Consider moving required background knowledge for the method from the related work section (after the method) to the background section (before the method).\n\n\nAdditional feedback:\n\nThe authors phrase the problem as the inability of the prior to model the posterior correctly. But wouldn't it rather be the inability of the posterior to correctly model the prior or our inability to train a posterior that accurately models the prior? It is well known that GANs provide very good results when sampling from a standard normal prior, thus it should be possible to train a VAE that uses the same prior. I think the authors should provide more evidence to the claim that the prior is at fault here and not the posterior/ELBO or rephrase the relevant parts of the paper. \n\nThe checklist includes several answers that are debatable\n- 1b) I do not see section 4 discussing limitations of the presented work but only limitations of the presented work. I would expect those to be discussed in the discussion/conclusion of experiments and results.\n- 3a) asks for code to which the authors reply yes, without making the code accessible. A description of the method is not equivalent to making the code public, e.g. to check for correctness or reproduction of the experiments. \n- 3b) I think number of training iterations/stopping criteria are missing.\n- 3d) asks for total estimated compute time used for the experiments. The paper only includes partial measurements and it is unclear how costly a full reproduction of the results would be.\n\n\nTypos:\nline 196, prior art -> prior work\nline 199, In most our experiments -> In most of our experiments\n\n", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "The authors do not discuss limitations of their work. An important limitation of their work is that their proposed SIR sampling can fail. By proposing samples using the prior it can be difficult to sample regions which have high probability under the marginal posterior but low probability under the prior. While the authors explore Langevin sampling as alternative, the study does not compare how well the samples drawn by the samples reflect the marginal posterior.\n\nWhile the impact contains an impact statement, it is only of limited value since it is mostly used to praise the work presented. The societal impact of generating images of high fidelity could include for example the dangers of creating deep fakes. I do appreciate however the reference to bias in image generation.", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "5h"}, "extra_text": {}}, {"review_id": "-a1lMFKgTY", "reviewer": "Reviewer_DUyg", "summary": "In trained VAEs, mismatch between the prior and the variational posterior can lead to low-quality generated samples.\n\nTo address this issue, the authors suggest the following procedure:\n1) given a trained VAE, train a classifier to distinguish images generated from the variational posterior vs images generated from the prior;\n2) 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.\n\nMain contribution:\na practical method to improve the generative quality of several forms of VAE by reweighing the prior to diminish the mismatch from the variational posterior", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "Originality:\nOverall, the combination of ideas/methods to improve image generation in VAEs is original.\n\nQuality:\nThe submission is technically sound.\n\nClarity:\nGenerally well written. I have two minor comments:\n- I think the manuscript would read better if the authors first discuss the main steps of the procedure (train classifier given a trained VAE, then inference with sampling reweighting) as done in Section 3.2 and then describe the specifics of contrastive estimation in Section 3.1;\n- the authors state that prior hole problem is when \"the prior distribution fails to match the aggregate approximate posterior\". I would say that the contrary happens: after training, the aggregate approximate posterior may not match the prior sufficiently well (this may due to limited expressivity of the posterior, trainability, low sample size, etc);\n\nSignificance:\nThis is my only major concern. While I see the practicality of this procedure, the technical contribution is narrow and I wonder why one would prefer this procedure over applying state-of-the-art VAE methods such as VDVAE and VQVAE. Could the authors provide some guidance on the use cases where their method should be preferred?\n\nAfter revision:\nAfter reading authors' reply and other reviewers' comments, I decided to change my score to a 6 (+2 difference from previous score). I think my original assessment on significance was incomplete. As Reviewer C3kP notices (i) \"the idea is very simple and easy to implement, and \"The method can be plugged with many different VAE variants\". After more consideration, these pros address my original concern on the significance of this work. The authors may consider making this more explicit in their manuscript, should it be accepted.", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "NA", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "4"}, "extra_text": {}}]

Author rebuttal:

Meta-review: The paper proposes a new class of priors for Variational Auto-Encoders (VAEs). The main idea is to use Noise Contrastive Estimation and a two-stage training method. The reviewers highlighted that:
- The paper is clearly written.
- The idea is simple (in the positive sense!) and easy to implement.
- The proposed prior could be plugged in any VAE.

However, the reviewers raised some concerns:
- Sampling from the VAE with the suggested parameters is very slow.
- The experimental results are not overly convincing.
- The organization of the paper could be improved.
- The method cannot reliably be used to evaluate a bound on the log-likelihood.
- The only performance metric used is the FID, which arguably favors the quality of samples overgeneralization.

Moreover, I find it rather surprising that the authors do not compare their approach to at least one of the following two methods:
- Tomczak, J., & Welling, M. (2018). VAE with a VampPrior. In International Conference on Artificial Intelligence and Statistics (pp. 1214-1223). PMLR.
- Norouzi, S., Fleet, D. J., & Norouzi, M. (2020). Exemplar VAE: Linking Generative Models, Nearest Neighbor Retrieval, and Data Augmentation. arXiv preprint arXiv:2004.04795.

It seems rather natural to focus on a thorough comparison against other priors for VAE. Instead, the authors have a limited comparison in Table 2. Moreover, the authors decided to compare their method to various deep generative models. I totally agree that it is also important, however, without a proper comparison against various priors, it is hard to properly evaluate their idea.

The authors provided a thorough rebuttal and they promised to improve the paper. Overall, the paper is solid and the idea is neat, therefore, I tend to accept the paper.

Decision: Accept (Poster)