PaperlayerView exact revision on arXiv

Source-linked AI summary

A Contrastive Learning Approach for Training Variational Autoencoder Priors

Jyoti Aneja, Alexander Schwing, Jan Kautz, Arash Vahdat

arXiv:2010.02917v3cs.LGcs.CVstat.ML

TL;DR

VAEs can have prior regions that do not match the aggregate approximate posterior, leaving latent areas that decode to non-data-like samples. The paper reweights a base prior with a factor trained by noise contrastive estimation, and reports improved generative performance across datasets.

  • Problem

    The prior may fail to match the aggregate approximate posterior, leaving high-density latent regions that do not correspond to data-like samples.

    Source: §1 · p. 1, §3 · p. 3

  • Method

    The noise contrastive prior multiplies a base prior by a learned reweighting factor trained by contrasting aggregate-posterior and base-prior samples.

    Source: §1 · p. 2

  • Results

    NCP improves generative performance across four datasets; on CelebA-64 it improves NVAE from an FID of 13.48 to 5.25.

    Source: §5.2 · p. 7, §6 · p. 10

  • Takeaways & Limitations

    NCP can increase prior expressivity across VAE types and its hierarchical training can be parallelized across latent groups.

    Source: §1 · p. 2, §6 · p. 10

    On CIFAR-10 and CelebA-HQ-256, VAEBM outperforms NCP-VAE.

    Source: §5.2 · p. 7

Abstract

from arXiv · show

Variational autoencoders (VAEs) are one of the powerful likelihood-based generative models with applications in many domains. However, they struggle to generate high-quality images, especially when samples are obtained from the prior without any tempering. One explanation for VAEs' poor generative quality is the prior hole problem: the prior distribution fails to match the aggregate approximate posterior. Due to this mismatch, there exist areas in the latent space with high density under the prior that do not correspond to any encoded image. Samples from those areas are decoded to corrupted images. To tackle this issue, we propose an energy-based prior defined by the product of a base prior distribution and a reweighting factor, designed to bring the base closer to the aggregate posterior. We train the reweighting factor by noise contrastive estimation, and we generalize it to hierarchical VAEs with many latent variable groups. Our experiments confirm that the proposed noise contrastive priors improve the generative performance of state-of-the-art VAEs by a large margin on the MNIST, CIFAR-10, CelebA 64, and CelebA HQ 256 datasets. Our method is simple and can be applied to a wide variety of VAEs to improve the expressivity of their prior distribution.

1 Introduction

VAEs can have latent regions that are dense under the prior but poorly represented by the aggregate posterior, contributing to the prior-hole problem. The paper proposes a reweighted energy-based prior trained with noise contrastive estimation, including a scalable approach for hierarchical VAEs.

1 Introduction

  • The prior-hole problem occurs when high-prior-density latent regions have low aggregate-posterior density and decode to samples that are not data-like.The paper notes that this mismatch can persist despite progress in making approximate posteriors more expressive.Source: §1 · p. 1
  • Existing neural energy-based priors require costly iterative MCMC during training and scale poorly when applied to each group in hierarchical models.The paper identifies these requirements as obstacles to using EBMs for hierarchical VAEs.Source: §1 · p. 1
  • The proposed energy-based prior multiplies a trainable base prior by a learned reweighting factor to bring it closer to the aggregate posterior.The reweighting factor is represented by neural networks and the base prior by Normal distributions.Source: Figure 1 · p. 1, §1 · p. 2
  • Noise contrastive estimation trains the reweighting factor by distinguishing aggregate-posterior samples from base-prior samples, after the base prior is trained to approach the aggregate posterior.The authors first train the VAE with the base prior, then train the energy-based prior using NCE.Source: §1 · p. 2
  • The paper introduces noise contrastive priors (NCPs), extends their training to hierarchical VAEs by training latent-group factors in parallel, and reports large generative-quality improvements across datasets.The contributions include application to several VAE forms and hierarchical models with many latent-variable groups.Source: §1 · p. 2

2 Related Work

Related work addresses prior expressivity through energy-based models, alternative priors, and staged training, while NCP uses latent-space reweighting trained without MCMC. The paper distinguishes its approach by its compatibility with hierarchical VAEs and non-continuous data.

Energy-based Models (EBMs)

  • Neural energy-based priors can represent expressive distributions, but prior approaches use MCMC and have been limited to one latent group; VAEBM instead places an EBM in pixel space.VAEBM assumes continuous data and therefore cannot be applied to discrete data such as text or graphs, whereas NCP-VAE forms its energy function in latent space.Source: §2 · p. 2

Adversarial Training

  • Unlike GANs [18], whose discriminator updates the generator, NCE uses a fixed noise generator; this paper applies the likelihood-ratio approach in VAE latent space.The base prior and approximate posterior are trained with the VAE objective rather than an adversarial loss.Source: §2 · p. 2

Prior Hole Problem

  • VampPrior requires training or pseudo-data at sampling time, while Takahashi et al. use the aggregate posterior for latent sampling at test time.Source: §2 · p. 2

Reweighted Priors

  • Bauer & Mnih train a similar reweighting factor with truncated rejection sampling, while EIMs require many training samples that challenge memory-intensive hierarchical VAEs.NCP trains each latent-group reweighting factor in parallel, supporting model parallelism for hierarchical models.Source: §2 · pp. 2–3

Two-stage VAEs

  • VQ-VAE, 2s-VAE, and RAE fit additional latent-space models after autoencoder training, whereas this work improves prior expressivity with an EBM.Source: §2 · p. 3

3 Background

The background defines the VAE and hierarchical latent-variable setup, then characterizes prior holes as a mismatch between the prior and the aggregate approximate posterior.

Variational Autoencoders

  • A VAE models p(x, z) as p(z)p(x|z), with a latent prior and a likelihood that generates data given the latent variable.Its training objective maximizes a variational lower bound on log-likelihood, averaged over the data distribution.Source: §3 · p. 3

Hierarchical VAEs (HVAEs)

  • A hierarchical VAE divides z into K latent groups and defines its prior and approximate posterior over those groups conditionally.The approximate posterior up to group k−1 is denoted q(z<k|x).Source: §3 · p. 3

Prior Hole Problem

  • The aggregate approximate posterior q(z) averages q(z|x) over the data distribution, and the prior-hole problem is the failure of p(z) to match q(z).The paper relates optimizing the VAE bound with respect to prior parameters to minimizing KL(q(z)||p(z)).Source: §3 · p. 3

4 Noise Contrastive Priors (NCPs)

The section defines NCP as a reweighted base prior and develops a two-stage NCE procedure to learn it. It also describes test-time sampling and a conditional extension for hierarchical VAEs.

4 Noise Contrastive Priors (NCPs)

  • The NCP reweights a base prior by positive r(z), forming an energy-based prior intended to address limited prior expressivity.The prior is defined by the product of the base prior and the reweighting factor.Source: §4 · p. 3

4.1 Two-stage Training for Noise Contrastive Priors

  • Stage one trains the VAE with the base prior; stage two freezes the VAE and trains only r(z).The two-stage procedure first brings the base prior close to the aggregate posterior, then replaces it with the more expressive NCP.Source: §4.1 · p. 4, Figure 2 · p. 4

4.2 Learning The Reweighting Factor with Noise Contrastive Estimation

  • NCE trains a binary classifier on aggregate-posterior and base-prior samples to estimate r(z), working best when these distributions are reasonably close.The method uses classifier predictions because the aggregate-posterior density ratio cannot be computed explicitly.Source: §4.2 · p. 4

4.3 Test Time Sampling

  • SIR resamples base-prior proposals with weights r(z), while Langevin dynamics iteratively follows the NCP energy but is slower.SIR proposal generation and weight evaluation can be performed in parallel.Source: §4.3 · p. 5

4.4 Generalization to Hierarchical VAEs

  • Hierarchical NCPs reweight each conditional prior to match its aggregate conditional, using group-specific classifiers trained in parallel.The method samples hierarchically, using ancestral sampling and SIR or Langevin dynamics for each group.Source: §4.4 · p. 5

5 Experiments

Experiments compare NCP with prior-work approaches and apply it to hierarchical NVAE models. Across four datasets, NCP-VAE improves on NVAE; additional analyses examine sample quality, latent-group scaling, and sampling and matching behavior.

5.1.1 Comparison against Two-Stage VAEs

  • On CelebA-64, NCP changes vanilla-VAE FID from 48.12 to 41.28 and GMM-prior FID from 40.95 to 39.00; aggregate-posterior reconstruction FID is 36.01.The 36.01 reconstruction score indicates a lower bound for the two-stage VAE results given the small architecture and latent space.Source: Table 1 · p. 6, §5.1.1 · p. 6

5.1.2 Comparison against Reweighted Priors

  • On MNIST, NCP-VAE reaches NLL 82.82 nats, comparable to SNIS and better than LARS, but slightly inferior to SNIS.Source: Table 2 · p. 6, §5.1.2 · p. 6

5.1.3 Training using Normalizing Flows

  • Normalizing flows provide limited gains against prior holes, while NCP further improves the base VAE equipped with flows.Source: §5.1.3 · p. 6

5.2 Quantitative Results on Hierarchical Models

  • NCP-VAE improves on NVAE across four datasets and reduces its gap with GANs; CelebA-64 FID falls from 13.48 to 5.25 and CIFAR-10 FID from 51.71 to 24.08.The reported tables cover CelebA-64, CIFAR-10, CelebA-HQ-256, and MNIST.Source: Table 3 · p. 7, Table 4 · p. 7, §5.2 · p. 7

5.3 Qualitative Results

  • Qualitative samples use a temperature-adjusted prior and re-adjusted batch-normalization statistics; nearest-neighbor comparisons report generated samples distinct from training images.Nearest neighbors are found after downsampling, center-cropping, and measuring Euclidean distance.Source: Figure 3 · p. 7, §5.3 · p. 8, Figure 4 · p. 8

5.4 Additional Ablation Studies

  • As the number of latent groups increases, FID improves for both NVAE and NCP-VAE, showing NCP efficacy with expressive hierarchical priors.Source: §5.4 · pp. 8–9, Table 7 · p. 9
  • Increasing SIR proposals or LD iterations improves FID; SIR generates faster than LD at similar compute, but its GPU memory use scales with proposal count.Source: §5.4 · p. 9, Table 8 · p. 9
  • Across 15 CelebA-64 groups, higher NCE classification loss tracked greater p(z)-q(z) overlap and SIR effective sample size, while reweighting reduced mismatch in low-loss groups.Eight groups had final classification loss at least 0.4; SIR was more reliable on those groups.Source: §5.4 · pp. 9–10, Figure 5 · p. 9, Table 9 · p. 9

6 Conclusions

The conclusion presents the noise contrastive prior (NCP) as a way to address prior holes by reweighting a base prior, with training that extends to hierarchical VAEs.

6 Conclusions

  • NCP defines the prior as a reweighting factor multiplied by a base prior, trained by contrasting aggregate-posterior samples with base-prior samples.The paper presents it as a post-training mechanism to increase prior expressivity.Source: §6 · p. 10
  • Hierarchical NCP training scales to large VAEs by training all latent-variable groups in parallel.Source: §6 · p. 10
  • The paper reports that NCPs improve generative performance by a large margin for small single-group VAEs and state-of-the-art NVAEs.Source: §6 · p. 10

7 Impact Statement

The impact statement links more expressive VAE priors with improved generated-image quality and diversity, while emphasizing that bias assessment and correction remain important.

7 Impact Statement

  • The proposed method increases the expressivity of the latent-code sampling distribution, thereby increasing the quality and diversity of generated solutions.Source: §7 · p. 10
  • The authors expect improved generation to benefit data augmentation and representation learning, among other applications.They also identify computer graphics, biomedical imaging, and computational fluid dynamics as potential application domains.Source: §7 · p. 10
  • Potential bias reduction is not assured: VAE architecture, training hyperparameters, and test-time temperature can affect bias, so generation should be assessed against real-data modes and long-tail coverage.The statement recommends considering active research on generative-model bias correction in applications.Source: §7 · p. 10

9 Funding Transparency Statement

The authors report funding from NVIDIA and partial support from NSF and NIFA.

9 Funding Transparency Statement

  • The work was mostly funded by NVIDIA during an internship and was also partially supported by NSF and NIFA awards.The statement lists NSF grant numbers 1718221, 2008387, 2045586, 2106825, MRI 1725729, and NIFA award 2020-67021-32799.Source: §9 · p. 10

A Training Energy-based Priors using MCMC

These sections describe energy-based-prior training, the relation between VAE prior fitting and aggregate-posterior matching, and conditional NCE for hierarchical VAEs. They also outline NVAE context features and image-decoder choices.

A Training Energy-based Priors using MCMC

  • For an energy-based prior proportional to a reweighting factor times a base prior, the variational bound includes a log-normalization term whose parameter gradient is estimated with MCMC.The expectation term can be trained with the reparameterization trick.Source: A Training Energy-based Priors using… · p. 15

B Maximizing the Variational Bound from the Prior’s Perspective

B.1 VAE with a Single Group of Latent Variables

  • For a single latent group, maximizing the VAE objective with respect to prior parameters is equivalent to learning the prior by minimizing KL(q(z)||p(z)).The derivation replaces expected conditional entropy with aggregate-posterior entropy when optimizing the prior.Source: §B.1 · p. 15

B.2 Hierarchical VAEs

  • For hierarchical VAEs, prior optimization minimizes the expected conditional KL between aggregate posterior and prior for each latent group.The passage defines aggregate conditional distributions for each group given preceding groups.Source: §B.2 · pp. 15–16

C Conditional NCE for Hierarchical VAEs

  • Hierarchical conditional NCE trains a discriminator to distinguish aggregate-posterior samples from prior samples given preceding latent groups, sampling those contexts from the aggregate posterior.A context feature can represent the preceding groups instead of passing all their samples to the discriminator.Source: C Conditional NCE for Hierarchical VAEs · p. 16

D NVAE Based Model and Context Feature

  • NVAE classifiers receive a context feature from the top-down model’s residual cell; for RGB datasets, replacing the base decoder’s discretized logistic mixture with a Normal distribution improved base-model performance.MNIST uses a Bernoulli decoder in the base NVAE.Source: D NVAE Based Model and Context Feature · p. 16

E Implementation Details

The implementation details describe the binary classifier’s residual-block design, while the appendix provides nearest-neighbor examples and additional generated samples across CelebA variants and temperatures.

E Implementation Details

  • The binary classifier uses residual blocks with batch normalization, Swish, and squeeze-and-excitation; one block doubles channels and downsamples spatial dimensions.Block B also uses a factorized reduction with 1 × 1 convolutions on the skip connection; Table 10 lists classifier training hyperparameters.Source: E Implementation Details · pp. 16–17, Figure 6 · p. 17, Table 10 · p. 17

F Additional Examples - Nearest Neighbors from the Training Dataset

G Additional Qualitative Examples

H Additional Qualitative Examples

I Additional Qualitative Examples

J Experiment on Synthetic Data

On the 25-Gaussians synthetic dataset, the NCP approach removes prior samples in low-density regions and achieves a higher reported average log-likelihood than a vanilla VAE.

J Experiment on Synthetic Data

  • NCP removes many prior-decoded points in low-density regions of the 25-Gaussians data distribution.The dataset is a grid-arranged mixture of 25 two-dimensional Gaussians; Figure 13 compares prior samples without NCP and with NCP.Source: J Experiment on Synthetic Data · p. 23, Figure 13 · p. 23
  • −0.954 nats is the NCP-VAE average log-likelihood, compared with −2.753 nats for vanilla VAE.The estimate uses 50k samples from the true distribution and 20k Monte Carlo samples for the log partition function.Source: J Experiment on Synthetic Data · p. 23
Loading 2010.02917v3…