[Paper Review] Generating Diverse High-Fidelity Images with VQ-VAE-2

Ali Razavi, Aäron van den Oord+1 more|arXiv (Cornell University)|Jun 2, 2019

Generative Adversarial Networks and Image Synthesis35 references106 citations

TL;DR

Two-stage hierarchy: learn discrete latent codes with VQ-VAE and then model them with powerful autoregressive priors (PixelCNN with self-attention) to generate high-fidelity, diverse images at large scales.

ABSTRACT

We explore the use of Vector Quantized Variational AutoEncoder (VQ-VAE) models for large scale image generation. To this end, we scale and enhance the autoregressive priors used in VQ-VAE to generate synthetic samples of much higher coherence and fidelity than possible before. We use simple feed-forward encoder and decoder networks, making our model an attractive candidate for applications where the encoding and/or decoding speed is critical. Additionally, VQ-VAE requires sampling an autoregressive model only in the compressed latent space, which is an order of magnitude faster than sampling in the pixel space, especially for large images. We demonstrate that a multi-scale hierarchical organization of VQ-VAE, augmented with powerful priors over the latent codes, is able to generate samples with quality that rivals that of state of the art Generative Adversarial Networks on multifaceted datasets such as ImageNet, while not suffering from GAN's known shortcomings such as mode collapse and lack of diversity.

Motivation & Objective

• Motivate and demonstrate a scalable, likelihood-based approach for high-resolution image generation that rivals GANs in quality and diversity.
• Leverage a hierarchical VQ-VAE to compress images into discrete latent representations enabling efficient autoregressive modeling in latent space.
• Show that sampling in latent space is significantly faster than pixel-space sampling for large images.

Proposed method

• Train a two-level hierarchical VQ-VAE to encode 256x256 images into discrete latents at 64x64 (bottom) and 32x32 (top).
• Quantize encoder outputs with shared codebooks and optimize with the VQ-VAE loss including codebook and commitment terms; use exponential moving average to update codebooks.
• Fit powerful autoregressive priors over the discrete latents using PixelCNN with self-attention: a top-level prior over 32x32 latents (global structure) and a bottom-level prior over 64x64 latents (local details) conditioned on the top latents.
• Condition priors on class labels; employ multi-headed self-attention at the top level and a large conditioning stack for the bottom level to capture long-range and local dependencies.
• Perform classifier-based rejection sampling to trade off diversity and sample quality (evaluate using a pre-trained ImageNet classifier).
• Evaluate using NLL and reconstruction error, precision/recall trade-offs, CAS, and FID/IS metrics, comparing against BigGAN-deep and reconstructions."],
• research_questions:[

Experimental results

Research questions

• RQ1 Can a hierarchical VQ-VAE with autoregressive priors generate high-fidelity, globally coherent images at large resolutions?
• RQ2 Does modeling in the discrete latent space provide faster sampling and better diversity than pixel-space autoregressive models?
• RQ3 How does the proposed method compare to state-of-the-art GANs in quality and diversity across datasets like ImageNet and FFHQ?
• RQ4 Can classifier-based rejection sampling provide a controllable diversity-quality trade-off without SGD training instability?

Key findings

• Class-conditional samples on ImageNet (256x256) achieve high fidelity with competitive quality to state-of-the-art GANs.
• The hierarchical latent approach yields broad diversity, with qualitative comparisons showing comparable fidelity and higher diversity than BigGAN-deep on several classes.
• Quantitative metrics indicate strong generalization (NLLs on train/validation are close), and precision-recall with rejection sampling improves both quality (precision) and diversity (recall).
• CAS results show VQ-VAE reconstructions plus classifier evaluation yield higher top-1 and top-5 accuracy than BigGAN-deep in their setup (with real data higher).
• On FFHQ-1024, the three-level hierarchy captures long-range dependencies (e.g., matching eye color) and produces realistic faces across high resolution.
• Overall, autoregressive modeling in the latent space provides a simple and effective path to large-scale, high-fidelity image generation with broader diversity than GANs, mitigating mode-collapse concerns.