Source-linked AI summary
Generating Diverse High-Fidelity Images with VQ-VAE-2
Ali Razavi, Aaron van den Oord, Oriol Vinyals
TL;DR
The paper addresses high-resolution image generation while balancing fidelity, diversity, sampling speed, and generalization. It combines hierarchical VQ-VAE representations with powerful autoregressive priors over discrete latent codes. The resulting samples have competitive fidelity and broader diversity than GAN comparisons in several settings, while latent-space sampling is substantially faster than pixel-space sampling.
Problem
High-resolution image generation must balance sample quality, diversity, speed, and objective evaluation, while GANs can lack diversity and pixel-space likelihood has quality-measurement limitations.
Method
The paper uses a hierarchical VQ-VAE with simple feed-forward encoder and decoder networks and autoregressive priors over compressed discrete latent codes.
Results
VQ-VAE produces fidelity competitive with state-of-the-art GANs and higher diversity in several ImageNet comparisons, with FID improving from roughly ∼30 to ∼10 after rejection sampling.
Takeaways & Limitations
Autoregressive modeling in a hierarchical latent space is presented as a simple and effective approach for diverse high-resolution image generation.
Abstract
from arXiv · showhide
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.
1 Introduction
Generative models offer different trade-offs in quality, diversity, speed, and evaluation. This paper uses lossy compression and discrete latent-space priors to generate coherent, high-resolution images efficiently while addressing limitations associated with pixel-space likelihood models and GANs.
- Generative models include likelihood-based approaches and GANs, with trade-offs among sample quality, diversity, speed, and evaluation.
- GANs can generate high-quality, high-resolution images but may fail to capture the full diversity of the true distribution and lack a satisfactory generalization measure.
- Likelihood-based models support objective comparison and generalization measurement, but pixel-space likelihood is not always a reliable measure of sample quality.
- Lossy compression removes negligible information by encoding images into discrete latent representations that remain reconstructable with little distortion.The representations are over 30x smaller than the original image.
- A self-attention PixelCNN prior models the discrete representations, producing decoded samples with reconstruction-level quality and coherence.
- Training and sampling in the discrete latent space are 30x faster than applying the autoregressive model directly to pixels, enabling higher-resolution training.
2 Background
VQ-VAE represents observations with discrete codebook indices and reconstructs them through a decoder. Its training combines reconstruction with codebook and commitment objectives, while autoregressive models factorize distributions into conditional probabilities.
- VQ-VAE encodes an observation into a vector, replaces it with the nearest codebook prototype, and transmits the resulting discrete index to the decoder.The quantization process can be lossy.
- The decoder maps codebook indices back to vectors and reconstructs the input, with reconstruction gradients passed through the decoder and encoder using a straight-through estimator.
- The codebook loss aligns selected codebook vectors with encoder outputs, while the commitment loss keeps encoder outputs near their chosen vectors.
- The commitment mechanism discourages encoder outputs from switching too frequently between codebook vectors.
- Exponential moving average updates replace the codebook loss, using a decay parameter γ set to 0.99 in all experiments.
- Deep autoregressive models factorize a joint distribution into conditional distributions for each dimension given preceding dimensions.
3 Method
The method encodes images into hierarchical discrete latent maps, then learns autoregressive priors over those codes for sampling and decoding. Global and local information are modeled at different scales, while classifier-based rejection sampling trades off sample diversity and quality.
- 3 Method: The proposed method first trains a hierarchical VQ-VAE, then fits PixelCNN priors over the induced discrete latent spaces.Stage 1 learns the encoder and decoder representation; stage 2 learns priors over the latent codes.
- 3 Method: The hierarchy separates global information in a top latent code from local information in a bottom latent code conditioned on the top latent.This allows each prior to model correlations specific to its level.
- 3 Method: For 256 × 256 images, the encoder produces 64 × 64 bottom and 32 × 32 top latent maps, while a feed-forward decoder combines all quantized levels.The encoder downsamples by factors of four and two before quantization.
- 3 Method: The learned latent prior makes test-time samples closer to the decoder’s training distribution, supporting more coherent decoded outputs.The paper interprets prior fitting as improving the approximation to the latent variables’ true distribution.
- 3 Method: The top prior uses PixelCNN with multi-headed self-attention for long-range spatial correlations, while the bottom prior models higher-resolution local latents conditionally.The top-level network operates on 32 × 32 latent variables and the bottom-level prior on 64 × 64 latents.
- 3 Method: Classifier-based rejection sampling provides an automated way to trade off diversity and quality by favoring samples more likely to receive correct class labels.The method is motivated by the idea that samples closer to the true data manifold are more likely to be classified correctly.
4 Related Works
The work builds on VQ-VAE and related hierarchical or likelihood-based generators while distinguishing its latent-space autoregressive priors and feed-forward decoding approach. It also investigates self-attention, compute scaling, and rejection sampling as relevant design choices.
- 4 Related Works: The prior network is based on Gated PixelCNN augmented with self-attention, building on the VQ-VAE framework.The paper identifies VQ-VAE as the foundation of its approach.
- 4 Related Works: BigGAN motivates comparisons involving self-attention, compute scale, and the trade-off between sample diversity and sample quality.The paper investigates whether these elements also improve VQ-VAE sample quality.
- 4 Related Works: Subscale Pixel Networks generate high-resolution images by partitioning spatial dimensions while trading sampling speed against density-estimation performance and sample quality.Unlike the cited parallel multi-scale model, SPN does not impose the corresponding independence assumptions.
- 4 Related Works: Compared with related hierarchical VQ-VAE work, this method uses complementary information across levels, feed-forward decoders, and autoregressive priors only in compressed latent space.The paper contrasts this design with autoregressive pixel-space decoders and hierarchy-collapse concerns in earlier work.
- 4 Related Works: Rejection sampling for improving sample quality has prior precedents in both GANs and VAEs.The cited VAE approach combines a learned rejection-sampling proposal with the prior.
5 Experiments
Experiments evaluate sample quality, diversity, and generalization across ImageNet and FFHQ. VQ-VAE achieves competitive fidelity with broader diversity than BigGAN in several comparisons, while train-validation results suggest limited overfitting.
- ImageNet sample quality and diversity: VQ-VAE produces samples with comparable fidelity and higher diversity than BigGAN-deep in side-by-side ImageNet comparisons.The comparison covers representative ImageNet classes and reports broader sample variation for VQ-VAE.
- FFHQ high-resolution faces: The three-level hierarchical model generates realistic 1024 × 1024 FFHQ faces while preserving long-range facial dependencies and covering lower-density modes.Examples include matching eye colour, symmetric facial features, and green hair.
- Evaluation metrics: NLL provides an objective generalization measure, whereas FID, Inception Score, Precision-Recall, and CAS can overlook memorization.This motivates evaluating both likelihood-based generalization and quality-diversity metrics.
- Negative log-likelihood and reconstruction error: Close train-validation NLL values for the top and bottom priors indicate that neither prior network overfits.Table 1 also reports train-validation reconstruction errors, with small differences suggesting the VQ-VAE does not overfit.
- Precision-recall evaluation: VQ-VAE attains slightly lower precision but higher recall than BigGAN across ImageNet classes under precision-recall evaluation.The comparison uses classifier-based rejection sampling for VQ-VAE and multiple truncation levels for BigGAN-deep.
- FID and Inception Score: Classifier-based rejection sampling improves both VQ-VAE IS and FID, with FID decreasing from roughly ∼30 to ∼10.FID is sensitive to slight blurriness and perturbations in VQ-VAE reconstructions, so FID* is also reported against reconstructions.
6 Conclusion
The paper presents VQ-VAE with powerful autoregressive latent priors as a simple approach to diverse high-resolution image generation. Its samples are competitive with state-of-the-art GANs, while evaluation limitations remain.
- Conclusion: The proposed method combines VQ-VAE, powerful autoregressive priors, and hierarchical multi-scale latent maps for high-resolution image generation.The encoder and decoder remain simple and lightweight feed-forward networks.
- Conclusion: Best class-conditional samples have fidelity competitive with state-of-the-art GANs and broader diversity in several classes.The conclusion contrasts this result with known GAN limitations.
- Conclusion: Visual inspection remains necessary because concrete measures of sample quality and diversity are still in their infancy.
A.1 PixelCNN Prior Networks
The appendix lists hyperparameters for autoregressive prior networks used in the ImageNet-256 and FFHQ-1024 experiments.
- PixelCNN Prior Networks: Table 3 specifies autoregressive prior hyperparameters for the ImageNet-256 experiments.
- PixelCNN Prior Networks: Table 4 specifies autoregressive prior hyperparameters for the FFHQ-1024 experiments.
A.2 VQ-VAE Encoder and Decoder
The appendix provides the hyperparameters used for the VQ-VAE encoder and decoder in the ImageNet-256 and FFHQ-1024 experiments.
- VQ-VAE Encoder and Decoder: Table 5 reports VQ-VAE encoder and decoder hyperparameters for ImageNet-256 and FFHQ-1024.
B Additional Samples
The paper directs readers to a full, losslessly rendered version containing additional samples.
- The full version of the paper is available through a linked resource.
- The linked version is rendered without lossy compression.
- The full version includes additional samples.