[Paper Review] Very Deep VAEs Generalize Autoregressive Models and Can Outperform Them on Images

Rewon Child|arXiv (Cornell University)|Nov 20, 2020

Advanced Neural Network ApplicationsComputer Science34 references79 citations

TL;DR

The paper introduces a very deep hierarchical VAE that, when extended beyond previous depths, outperforms PixelCNN-based autoregressive models in log-likelihood on CIFAR-10, ImageNet variants, and FFHQ across high-dimensional images while using fewer parameters and enabling faster sampling.

ABSTRACT

We present a hierarchical VAE that, for the first time, generates samples quickly while outperforming the PixelCNN in log-likelihood on all natural image benchmarks. We begin by observing that, in theory, VAEs can actually represent autoregressive models, as well as faster, better models if they exist, when made sufficiently deep. Despite this, autoregressive models have historically outperformed VAEs in log-likelihood. We test if insufficient depth explains why by scaling a VAE to greater stochastic depth than previously explored and evaluating it CIFAR-10, ImageNet, and FFHQ. In comparison to the PixelCNN, these very deep VAEs achieve higher likelihoods, use fewer parameters, generate samples thousands of times faster, and are more easily applied to high-resolution images. Qualitative studies suggest this is because the VAE learns efficient hierarchical visual representations. We release our source code and models at https://github.com/openai/vdvae.

Motivation & Objective

• Motivate whether VAEs can match or exceed autoregressive models for natural images when deeply stacked latent variables are used.
• Propose and validate a very deep hierarchical VAE architecture capable of scaling to many stochastic layers.
• Empirically compare log-likelihoods and efficiency against PixelCNN/PixelCNN++ baselines on CIFAR-10, ImageNet (32/64), FFHQ variants.
• Show that depth enables efficient hierarchical representations and parallelizable generation, leading to faster sampling.

Proposed method

• Formalize N-layer VAE with top-down factorization of prior and posterior (p(z)=p(z0)p(z1|z0)…; q(z|x)=q(z0|x)q(z1|z0,x)…).
• Argue and prove that N-layer VAEs generalize autoregressive models when N equals data dimension and can universally approximate N-dimensional latent densities.
• Introduce a very deep VAE architecture with bottleneck residual blocks, diagonal Gaussian priors/posteriors, and residual scaling (1/√N) to stabilize deep training.
• Use nearest-neighbor upsampling for unpooling to avoid posterior collapse and remove KL warmup/free-bits requirements.
• Stabilize training with a gradient-skipping heuristic to avoid divergent updates during deep training.
• Train and evaluate on CIFAR-10, ImageNet-32, ImageNet-64, FFHQ-256, and FFHQ-1024, comparing log-likelihoods to autoregressive baselines.

Experimental results

Research questions

• RQ1 Can VAEs, when made sufficiently deep, match or outperform autoregressive models on image log-likelihood?
• RQ2 Does increasing stochastic depth (depth of latent layers) improve VAE performance independent of total parameter count?
• RQ3 Do very deep VAEs learn efficient hierarchical representations enabling parallelizable generation and faster sampling?
• RQ4 How does a very deep VAE scale to high-resolution images compared with PixelCNN-based models?

Key findings

• Very deep VAEs achieve higher log-likelihoods than PixelCNN-based models on CIFAR-10, ImageNet-32, ImageNet-64, FFHQ-256, and FFHQ-1024.
• Depth improvements occur independent of model capacity, with deeper models outperforming shallower ones using similar parameter counts.
• The very deep VAE learns efficient hierarchical representations, enabling many latent variables at high resolution to be generated in parallel.
• Compared to PixelCNN-based autoregressive models, the Very Deep VAE uses fewer parameters and enables sampling thousands of times faster, while scaling to larger images.
• The architecture scales well to 1024x1024 images, maintaining fast sampling (single forward pass) versus autoregressive counterparts that require many evaluations.