[Paper Review] NVAE: A Deep Hierarchical Variational Autoencoder

Arash Vahdat, Jan Kautz|arXiv (Cornell University)|Jul 8, 2020

Generative Adversarial Networks and Image Synthesis79 references378 citations

TL;DR

NVAE designs a deep hierarchical VAE with depthwise separable convolutions and residual posterior parameterization, achieving state-of-the-art non-autoregressive likelihoods on several image datasets while enabling large-scale image generation.

ABSTRACT

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$ imes$256 pixels. The source code is available at https://github.com/NVlabs/NVAE .

Motivation & Objective

• Design a deep hierarchical VAE architecture tailored for high-quality image generation.
• Stabilize training of very deep VAEs with many latent groups.
• Improve memory efficiency and sampling speed for large images.

Proposed method

• Introduce depthwise separable convolutions in the generative model to expand receptive fields efficiently.
• Use a residual parameterization of the approximate posterior relative to the prior to stabilize KL terms.
• Apply spectral regularization to bound the Lipschitz constant and stabilize training.
• Incorporate batch normalization with tuned momentum and BN-activation pairing to improve training stability.
• Employ mixed-precision training and gradient checkpointing to reduce memory usage.
• Optionally apply lightweight normalizing flows in the encoder to increase posterior expressiveness.

Experimental results

Research questions

• RQ1 Can a carefully designed deep hierarchical VAE outperform existing non-autoregressive likelihood-based models on standard image datasets?
• RQ2 What architectural choices (convolutions, normalization, activation, residual parameterization) most improve VAE training stability for large images?
• RQ3 How do memory- and compute-saving techniques affect training and sampling efficiency for deep VAEs at 256×256 resolutions?
• RQ4 Does adding normalizing flows to the encoder meaningfully improve held-out log-likelihood without sacrificing stability?

Key findings

• NVAE achieves state-of-the-art results among non-autoregressive likelihood-based models on MNIST, CIFAR-10, CelebA 64, and CelebA HQ-256, and is a strong baseline on FFHQ-256.
• On CIFAR-10, NVAE improves from 2.98 to 2.91 bits per dimension.
• NVAE can generate high-quality 256×256 images and is among the first VAEs to do so without changing the standard VAE objective.
• Sampling is fast due to the unconditional decoder, with 56 ms per image (batch size 36) on a Titan V GPU.
• Ablations show that BN with Swish activation and SE, depthwise separable generative cells, SR, and residual posterior parameterization all contribute to performance and stability.
• Memory reduction techniques (mixed-precision and gradient checkpointing) roughly double training throughput.