PaperlayerView exact revision on arXiv

Source-linked AI summary

NVAE: A Deep Hierarchical Variational Autoencoder

Arash Vahdat, Jan Kautz

arXiv:2007.03898v3stat.MLcs.CVcs.LG

TL;DR

VAEs offer efficient sampling and accessible encoders but have been outperformed by competing likelihood-based models, while VAE research has emphasized statistical rather than architectural challenges. NVAE addresses this gap with a deep hierarchical architecture using specialized convolutions, batch normalization, residual Normal posterior parameterization, and spectral regularization. It achieves state-of-the-art results among non-autoregressive likelihood-based models on several image datasets and supports high-quality generation at 256×256 resolution.

  • Problem

    VAEs provide fast, tractable sampling and accessible encoders but have been outperformed by other likelihood-based models, while architectural design for VAEs remains comparatively underexplored.

  • Method

    NVAE is a deep hierarchical VAE using depthwise separable convolutions, batch normalization, residual parameterization of Normal posteriors, and spectral regularization.

  • Results

    NVAE achieves state-of-the-art results among non-autoregressive likelihood-based models on MNIST, CIFAR-10, CelebA 64, and CelebA HQ, while providing a strong FFHQ baseline.

  • Takeaways & Limitations

    NVAE demonstrates that carefully designed deep hierarchical VAE architectures can produce high-quality large images without changing the VAE objective.

  • Takeaways & Limitations

    The paper notes that batch-normalization effects require further study and that instance normalization did not reproduce the same quantitative and qualitative results.

Abstract

from arXiv · show

Normalizing flows, autoregressive models, variational autoencoders (VAEs), and deep energy-based models are among competing likelihood-based frameworks for deep generative learning. Among them, VAEs have the advantage of fast and tractable sampling and easy-to-access encoding networks. However, they are currently outperformed by other models such as normalizing flows and autoregressive models. While the majority of the research in VAEs is focused on the statistical challenges, we explore the orthogonal direction of carefully designing neural architectures for hierarchical VAEs. We propose Nouveau VAE (NVAE), a deep hierarchical VAE built for image generation using depth-wise separable convolutions and batch normalization. NVAE is equipped with a residual parameterization of Normal distributions and its training is stabilized by spectral regularization. We show that NVAE achieves state-of-the-art results among non-autoregressive likelihood-based models on the MNIST, CIFAR-10, CelebA 64, and CelebA HQ datasets and it provides a strong baseline on FFHQ. For example, on CIFAR-10, NVAE pushes the state-of-the-art from 2.98 to 2.91 bits per dimension, and it produces high-quality images on CelebA HQ. To the best of our knowledge, NVAE is the first successful VAE applied to natural images as large as 256$\times$256 pixels. The source code is available at https://github.com/NVlabs/NVAE .

1 Introduction

NVAE addresses the underexplored architectural needs of VAEs by introducing a deep hierarchical design for high-quality image generation. It combines specialized convolutions, batch normalization, residual posterior parameterization, and spectral regularization, achieving strong results across image-generation settings.

  • Motivation: VAEs have fast, tractable sampling and accessible encoders, but their neural architectures have received less attention than statistical improvements.Prior work often borrows architectures from classification tasks despite VAEs’ different information-retention requirements.
  • Method: NVAE is a deep hierarchical VAE designed around depthwise convolutions that rapidly expand receptive fields without dramatically increasing parameter counts.The model targets high-quality image generation while retaining the VAE framework.
  • Method: Batch normalization is identified as important for deep VAE success, while residual posterior parameterization and spectral regularization address KL optimization and training instability.These components are presented as remedies for instability when increasing the number of hierarchical groups.
  • Contributions: NVAE introduces practical memory-reduction solutions and demonstrates that deep hierarchical VAEs can achieve state-of-the-art results on several image datasets.The paper reports high-quality samples from the original VAE objective and identifies NVAE as the first successful VAE application to images as large as 256×256 pixels.
  • Related Work: NVAE differs from VQ-VAE-2 by optimizing the VAE objective directly and using an unconditional data-space decoder rather than a slow PixelCNN latent prior.Compared with IAF-VAEs, it changes the neural networks, posterior parameterization, and scaling to large images.

2 Background

VAEs define a latent-variable generative model whose intractable posterior is approximated by an encoder. Hierarchical VAEs partition latents into groups and use neural networks to implement conditional priors, posteriors, and the variational objective.

  • VAE Foundations: A VAE models data with a prior p(z) and likelihood p(x|z), while an encoder q(z|x) approximates the generally intractable posterior.Training uses the approximate posterior to optimize the generative model.
  • Hierarchical Extension: Hierarchical VAEs partition latent variables into disjoint groups to increase the expressiveness of the approximate posterior and prior.The number of latent groups is denoted by L.
  • Hierarchical Extension: Each hierarchical prior and approximate-posterior conditional is represented by a factorial Normal distribution conditioned on preceding latent groups and, for the posterior, the input.The resulting variational lower bound is optimized using the reparameterization trick.
  • Neural Implementation: A top-down network generates conditional parameters in the generative model, whereas a bottom-up deterministic network extracts representations for inferring latent variables.Samples from latent groups are combined with deterministic feature maps before being passed to subsequent groups.

3 Method

NVAE is a deep hierarchical VAE designed for large-image generation through architecture modules, hierarchical multi-scale modeling, and stabilized posterior optimization. Its design combines depthwise convolutions, normalization and activation choices, residual posterior parameterization, and spectral regularization.

  • Method overview: NVAE targets expressive neural networks and stable scaling for deep hierarchical VAEs applied to large images.The method addresses architecture design alongside the challenges of increasing hierarchical groups and image sizes.
  • Hierarchical modeling: The hierarchical multi-scale generator starts from small spatial latent variables and progressively doubles spatial dimensions.This assigns global long-range correlations to higher hierarchy groups and local fine-grained dependencies to lower groups.
  • Residual cells: Depthwise convolutions increase receptive fields with substantially fewer parameters and lower computational complexity than regular convolutions.NVAE uses them in the generative model, while regular residual cells are retained for the bottom-up encoder because depthwise cells do not improve its performance.
  • Normalization and activation: BN combined with Swish outperforms weight normalization with ELU in the examined VAE cells.The paper identifies batch normalization as important for deep VAE success and reports that its negative impact occurs during evaluation rather than training.
  • Residual cells: Squeeze-and-excitation provides a channel-wise gating mechanism that improves VAEs.The residual generative cell adds BN layers, Swish activation, and SE alongside depthwise convolutions.
  • Posterior parameterization and training stability: Residual Normal distributions parameterize the approximate posterior relative to the prior, while spectral regularization constrains encoder sensitivity to stabilize KL minimization.The residual parameterization makes the posterior move with the prior, and spectral regularization targets the unbounded KL optimization challenge.

4 Experiments

NVAE is evaluated across image datasets using likelihood comparisons, qualitative sampling and reconstruction, and ablations of normalization, residual distributions, and spectral regularization. It achieves strong non-autoregressive likelihood results, high-quality samples, fast sampling, and identifiable benefits from several architectural components.

  • 4.1 Main Quantitative Results: NVAE outperforms non-autoregressive flow and VAE models on all evaluated datasets except ImageNet, where it ranks second after Flow++.On CIFAR-10, it improves the state of the art from 2.98 to 2.91 bpd.
  • 4.2 Qualitative Results: NVAE produces diverse, high-quality CelebA HQ samples even with small prior temperatures, with improved hair detail and diversity in visual comparison.Lowering temperature commonly improves sample quality but reduces diversity; NVAE obtains both properties with its batch-normalization sampling procedure.
  • 4.3 Ablation Studies: Batch normalization improves ELU training, especially at L = 40, and performs better with Swish than the compared weight-normalization configuration.The normalization and activation comparison is reported in Table 2.
  • 4.3 Ablation Studies: Depthwise-convolution residual cells improve the top-down generative model but do not change bottom-up encoder performance, so NVAE uses them only in the generative model.Regular encoder cells require less memory and train faster.
  • 4.3 Ablation Studies: Residual Normal distributions improve the training KL term by 0.04 bpd and final test log-likelihood by 0.03 bpd without materially changing active latent variables or reconstruction loss.The comparison uses the 40-group model with and without residual distributions.
  • 4.3 Ablation Studies: Removing spectral regularization or squeeze-and-excitation hurts performance, while spectral regularization also slightly improves generative performance beyond stabilizing training.These effects are measured in the ablation study summarized in Table 5.
  • 4.1 Main Quantitative Results: 2.03 seconds for 36 images at 256×256 px, compared with MaCow’s 434.2 ms/image, making NVAE approximately 8× faster in batched sampling.The measurements use a 12-GB Titan V GPU.
  • 4.2 Qualitative Results: NVAE reconstructions on CelebA HQ are reported as indistinguishable from training images, while CIFAR-10’s KL values are distributed similarly across most latent groups.Only one of the 30 visualized groups is turned off, with KL balancing applied during the first approximately 25,000 iterations.

5 Conclusions

NVAE is a deep hierarchical VAE whose architecture and training techniques enable high-quality image generation and state-of-the-art results across several image datasets.

  • NVAE uses depthwise separable convolutions in the generative model and regular convolutions in the encoder.
  • Residual parameterization of Normal distributions and spectral regularization stabilize training of very deep models.
  • NVAE’s practical memory remedies speed up training by ∼2×.
  • NVAE achieves state-of-the-art results on MNIST, CIFAR-10, CelebA 64, and CelebA HQ-256, while providing a strong baseline on FFHQ-256.
  • NVAE produces large high-quality images without changing the VAE objective function.

Impact Statement

The paper’s impact centers on expressive architectures for image VAEs, with potential applications in content generation and related tasks. However, generated outputs can reproduce biases present in training data and acquire additional bias during design, training, or sampling.

  • The work examines expressive neural architectures for image VAEs on commonly used public datasets.
  • Potential applications include content generation, computer graphics, data augmentation, semi-supervised learning, and representation learning.
  • VAEs avoid mode collapse and are described as representing data distributions more faithfully than commonly used GANs.
  • Bias in data collection can be reproduced by VAEs, while additional bias may arise during model design, training, or low-temperature sampling.

A Additional Implementation Details

The implementation details describe optimization, KL handling, normalization, decoder choices, evaluation, architecture scaling, datasets, and memory-conscious training configurations for NVAE.

  • The KL term is warmed up at the beginning of training, with β annealed from 0 to 1 during the first 30% of training.
  • KL balancing coefficients encourage comparable information usage across latent groups during warm-up and are reset to 1 afterward.
  • NVAE uses spectral-related smoothness regularization with stronger initial values when training CelebA HQ and FFHQ is unstable.
  • Training uses AdaMax with cosine learning-rate decay, including a reduced FFHQ learning rate of 0.008.
  • The decoder uses discretized Logistic mixtures except on MNIST, where it uses a Bernoulli distribution, and remains unconditional across image spatial locations.
  • Test log-likelihood is estimated with importance-weighted sampling using 1000 encoder-based samples.
  • A smaller model can run on 8 GPUs while increasing negative log-likelihood by only 0.01 bpd on the reported datasets.
  • Experiments cover MNIST, CIFAR-10, ImageNet 32 × 32, CelebA 64 × 64, CelebA HQ, and FFHQ 256×256.

B Additional Experiments and Visualizations

This section provides additional insights into NVAE beyond the main presentation.

  • The section provides additional insights into NVAE.
  • These insights concern aspects of NVAE discussed beyond the preceding material.
  • The section extends the paper’s discussion of NVAE with further observations.

B.1 Is NVAE Memorizing the Training Set?

NVAE’s experiments find no observed overfitting as model capacity increases, and generated samples are not present in the training set. The evaluation compares generated images with similar training examples and considers sampling settings.

  • NVAE shows no observed overfitting as depth and width increase, especially on datasets with large images.The authors report stopping model growth because of compute and training-time considerations rather than observed overfitting.
  • Generated NVAE samples are not present in the training set when compared with the most similar training images using downsampled central-crop L2 distance.The comparison uses aligned images and focuses on facial features including the eyes, nose, and mouth.
  • Figure 7 examines prior temperatures together with default and readjusted batch-normalization statistics when sampling from VAEs.Readjusted statistics improve diversity and quality in the reported CelebA HQ experiment.

B.3 Additional Generated Samples

Figures 8 and 9 show additional NVAE samples generated from a CelebA HQ-trained model. The samples use higher temperatures and were manually selected.

  • Figures 8 and 9 visualize additional samples generated by NVAE.
  • The NVAE model used for these samples was trained on CelebA HQ.
  • Sampling used temperatures from 0.6 to 0.9, with samples manually selected.

B.4 More on the Impact of Residual Normal Distributions

Residual Normal distributions keep more latent channels active during the early KL warm-up phase and achieve a better final KL value. The section also shows sampling results under different temperatures and batch-normalization settings.

  • Residual Normal distributions keep more latent variables active during the initial KL warm-up phase than distributions predicting absolute posterior parameters.The comparison covers the total number of active channels across all latent variables.
  • Readjusting batch-normalization running averages during sampling improves the diversity and quality of CelebA HQ samples.The readjustment uses 500 samples at the given prior temperature.
  • Additional CelebA HQ samples are shown at higher prior temperatures t ∈{0.6, 0.7, 0.8, 0.9}, with manual selection.
  • The residual-distribution model achieves a better KL value at the end of training.This comparison corresponds to the experiment reported in Table 4.

B.5 Stabilizing the Training with Spectral Regularization

Spectral regularization stabilizes NVAE training when the KL term causes instability, including on FFHQ. The hierarchical model uses multiple scales, with global structure concentrated at the top and local variation at lower groups.

  • B.5 Stabilizing the Training with Spectral Regularization: Increasing the spectral regularization coefficient λ from 0.1 to 1.0 stabilizes unstable FFHQ training caused by the KL term.Gradient clipping and restricting Normal-distribution parameters did not stabilize training without negatively affecting performance.
  • B.6 Long-Range Correlations: NVAE’s CelebA HQ 256×256 generative model uses five latent-variable scales, beginning at 8×8 and doubling spatial dimensions up to 128×128.The hierarchy samples groups sequentially while increasing spatial resolution.
  • B.6 Long-Range Correlations: Global long-range correlations are captured mostly at the hierarchy’s top, while local variations are recorded in lower groups.The figure studies this behavior by fixing samples at selected scales and sampling the remaining hierarchy.
Loading 2007.03898v3…