PaperlayerView exact revision on arXiv

Source-linked AI summary

Autoregressive Image Generation using Residual Quantization

Doyup Lee, Chiheon Kim, Saehoon Kim, Minsu Cho, Wook-Shin Han

arXiv:2203.01941v2cs.CVcs.LG

TL;DR

Previous VQ methods struggle to shorten image-code sequences while preserving reconstruction quality, limiting efficient autoregressive modeling. The paper combines RQ-VAE’s coarse-to-fine residual quantization with RQ-Transformer to generate high-quality images from short code sequences, outperforming previous autoregressive models across multiple benchmarks.

  • Problem

    Previous VQ methods face a rate-distortion trade-off: shortening the spatial code sequence requires exponentially larger codebooks, which can be inefficient and unstable.

  • Method

    RQ-VAE recursively quantizes feature vectors from coarse to fine using stacked codes, while RQ-Transformer predicts the next stack to model images autoregressively.

  • Results

    The framework outperforms previous autoregressive models across image-generation benchmarks including LSUNs, FFHQ, ImageNet, and CC-3M.

  • Takeaways & Limitations

    Using residual quantization, the approach represents images with shorter code sequences while preserving reconstruction quality and reducing autoregressive computational costs.

  • Takeaways & Limitations

    The models do not outperform StyleGAN2 on unconditional generation, especially on small-scale datasets such as FFHQ, where autoregressive overfitting is implicated.

Abstract

from arXiv · show

For autoregressive (AR) modeling of high-resolution images, vector quantization (VQ) represents an image as a sequence of discrete codes. A short sequence length is important for an AR model to reduce its computational costs to consider long-range interactions of codes. However, we postulate that previous VQ cannot shorten the code sequence and generate high-fidelity images together in terms of the rate-distortion trade-off. In this study, we propose the two-stage framework, which consists of Residual-Quantized VAE (RQ-VAE) and RQ-Transformer, to effectively generate high-resolution images. Given a fixed codebook size, RQ-VAE can precisely approximate a feature map of an image and represent the image as a stacked map of discrete codes. Then, RQ-Transformer learns to predict the quantized feature vector at the next position by predicting the next stack of codes. Thanks to the precise approximation of RQ-VAE, we can represent a 256$\times$256 image as 8$\times$8 resolution of the feature map, and RQ-Transformer can efficiently reduce the computational costs. Consequently, our framework outperforms the existing AR models on various benchmarks of unconditional and conditional image generation. Our approach also has a significantly faster sampling speed than previous AR models to generate high-quality images.

1. Introduction

VQ enables autoregressive models to generate high-resolution images by representing images as sequences of discrete codes. The quantized codes are predicted sequentially by the AR model.

  • VQ represents an image as a sequence of discrete codes for autoregressive generation.
  • The image feature map is quantized before its codes are rearranged into an ordered sequence.
  • An autoregressive model is trained to sequentially predict the codes in that sequence.

POSTECH

The paper addresses the rate-distortion and computational challenges of autoregressive image generation with RQ-VAE and RQ-Transformer. Residual quantization preserves fidelity with short code sequences, while the transformer predicts stacked codes and uses exposure-bias mitigation techniques.

  • Shorter code sequences reduce autoregressive computational costs because predictions use codes from previous positions.
  • Previous VQ methods face a rate-distortion limitation when reducing feature-map resolution while preserving reconstructed-image quality.
  • RQ-VAE recursively quantizes residuals with a fixed-size codebook, representing each feature vector as a stack of discrete codes.The method composes as many vectors as the codebook size raised to the quantization depth.
  • RQ-VAE enables 256×256 images to use 8×8 feature maps while retaining precise feature-map approximation.
  • RQ-Transformer predicts the next stack of RQ-VAE codes to estimate the next quantized feature vector.
  • Soft labeling and stochastic sampling are introduced to reduce exposure bias in RQ-Transformer training.

2. Related Work

Prior work uses autoregressive models over discrete image codes, but pixel-level modeling is inefficient for high-resolution images. VQ-based methods improve efficiency, while composite quantization provides a related route to precise approximation under fixed codebook constraints.

  • AR Modeling for Image Synthesis: Pixel-level autoregressive image modeling is infeasible for high-resolution images because generation is slow and quality is low.
  • AR Modeling for Image Synthesis: VQ-VAE-based methods represent images as discrete codes and use autoregressive models to predict those codes.
  • AR Modeling for Image Synthesis: VQ-GAN improves reconstructed-image perceptual quality with adversarial and perceptual losses.
  • AR Modeling for Image Synthesis: When feature-map resolution is further reduced, VQ-GAN cannot precisely approximate image feature maps with its limited codebook size.
  • VQs in Other Applications: Composite quantization represents vectors as compositions of codes to achieve precise approximation under a given codebook size.
  • VQs in Other Applications: RQ iteratively quantizes a vector and its residuals, while RQ-VAE uses one shared codebook across quantization steps.

3. Methods

The method combines RQ-VAE, which compresses images into stacked residual codes while preserving feature-map fidelity, with RQ-Transformer, which predicts code stacks efficiently across spatial positions.

  • RQ-VAE: Reducing feature-map resolution lowers autoregressive computation, but conventional VQ requires a codebook of size K^4 to halve both spatial dimensions while preserving bit rate.Large VQ codebooks can suffer codebook collapse and unstable training.
  • RQ-VAE: RQ-VAE replaces VQ with residual quantization to represent each feature vector using an ordered stack of D codes from a fixed-size codebook.Residual quantization recursively selects codes for successive residuals, producing a coarse-to-fine approximation.
  • RQ-VAE: RQ with depth D provides up to K^D partitions, matching the capacity of VQ with an exponentially larger codebook size K^D.This allows increased quantization depth to replace exponentially growing VQ codebooks.
  • RQ-VAE: RQ-VAE encodes a feature map as M ∈[K]H×W×D and reconstructs the image from the final quantized feature map through the decoder.The quantized feature maps at intermediate depths are also defined, with the depth-D map denoted by ˆZ.
  • RQ-VAE: RQ-VAE’s precise approximation supports more aggressive downsampling and more realistic reconstructions at a fixed downsampling factor.The paper connects this precision to lower autoregressive computation, faster generation, and better learning of long-range code interactions.
  • RQ-Transformer: RQ-Transformer reshapes the spatial code map into T rows of D codes and predicts the next D-code stack using spatial and depth transformers.The spatial transformer summarizes previous positions, while the depth transformer predicts codes within the next position.
  • RQ-Transformer: Soft labeling and stochastic sampling of RQ-VAE codes are introduced to address exposure bias during RQ-Transformer training.The approach uses the geometric relationship of RQ-VAE code embeddings instead of scheduled sampling.

4. Experiments

Experiments evaluate RQ-Transformer on unconditional and conditional image generation, sampling efficiency, and RQ-VAE reconstruction quality. Across these settings, the framework reports strong generation performance, faster sampling, and improved rate-distortion trade-offs from residual quantization.

  • Unconditional Image Generation: RQ-Transformer outperforms other autoregressive models on unconditional image generation, with larger gains on the larger LSUN-cat and LSUN-bedroom datasets.The authors report marginal improvements on LSUN-church and FFHQ, but significant improvements over autoregressive and diffusion-based models on LSUN-cat and LSUN-bedroom.
  • Conditional Image Generation: RQ-Transformer significantly outperforms previous models on ImageNet class-conditioned generation, achieving 7.55 FID with 3.8B parameters without rejection sampling.A 1.4B-parameter model reaches 8.71 FID when RQ-VAE training increases from 10 to 50 epochs, while rejection sampling further improves reported performance.
  • Conditional Image Generation: RQ-Transformer significantly outperforms VQ-GAN and ImageBART on CC-3M text-conditioned generation across FID and CLIP score.The model uses 23% of ImageBART’s parameters while outperforming it on both reported metrics.
  • Computational Efficiency: 0.02 seconds per image is achieved by RQ-Transformer at batch size 500, making it 7.3× faster than VQ-GAN at batch size 200.At batch sizes 100 and 200, the reported speed-ups over VQ-GAN are 4.1× and 5.6×, respectively.
  • Ablation Study on RQ-VAE: Increasing quantization depth D improves reconstruction quality more effectively than increasing codebook size K, enabling an 8×8×4 code map at fixed K=16,384.The authors choose D=4 for AR modeling because deeper quantization improves reconstruction further but increases computational costs.
  • Ablation Study on RQ-VAE: RQ-VAE performs coarse-to-fine feature-map estimation, adding clearer and finer image details as successive residual-code depths are included.Code embeddings with smaller norms are used at later depths, and overlapping code-usage distributions indicate code sharing across levels.

5. Conclusion

RQ-VAE and RQ-Transformer provide a short discrete representation and efficient autoregressive prediction for high-quality image generation. The approach outperforms previous autoregressive models across several benchmarks, while limitations remain in unconditional generation, bidirectional context modeling, and training cost.

  • RQ-VAE precisely approximates image feature maps with short stacked code sequences, enabling RQ-Transformer to predict codes with low computational costs.The framework uses discrete image representations to support autoregressive high-resolution image generation.
  • Unconditional generation does not outperform StyleGAN2, particularly on small-scale FFHQ, where autoregressive models can overfit.The paper identifies regularization of autoregressive models as a direction for high-resolution generation on small datasets.
  • As quantization depth increases, RQ-VAE reconstructions become clearer and recover finer details from the original image.
  • Autoregressive models capture only unidirectional contexts, limiting image manipulation applications such as inpainting and outpainting.The paper suggests modeling bidirectional contexts could improve image quality and broaden use.
  • Training large-scale autoregressive models remains expensive, energy-intensive, and capable of producing a substantial carbon footprint as model and dataset scales increase.

A.1. Architecture of RQ-VAE

The RQ-VAE architecture follows VQ-GAN but adds residual capacity and down-/up-sampling to produce 8×8 feature maps.

  • RQ-VAE follows VQ-GAN’s architecture for fair comparison and adds two 512-channel residual blocks with a down-/up-sampling block.These additions extract feature maps at 8×8 resolution.

A.2. Architecture of RQ-Transformer

RQ-Transformer uses separate spatial and depth transformers built from self-attention blocks to model stacked RQ-VAE codes.

  • RQ-Transformer consists of spatial and depth transformers, each using a stack of self-attention blocks.The reported models use RQ-VAE codes with an 8×8×4 shape.

A.3. Training Details

The experiments use specified optimization schedules and sampling settings across datasets, with early stopping for FFHQ to address memorization risk.

  • RQ-VAE ImageNet training uses Adam with learning rate 0.00004, 10 epochs, batch size 128, and VQ-GAN adversarial and perceptual-loss settings.The learning rate is linearly warmed up during the first 0.5 epoch without decay, weight decay, or dropout.
  • RQ-Transformers use AdamW, cosine learning-rate scheduling, 0.0005 initial learning rate, weight decay 0.0001, and dataset-dependent training epochs and batch sizes.Self-attention dropout is 0.1 except for the 3.8B-parameter model, which uses 0.3.
  • Unconditional sampling uses dataset-specific top-k or top-p settings for LSUN-cat, LSUN-bedroom, LSUN-church, and FFHQ.
  • FFHQ training uses early stopping when validation loss is minimized because the small dataset allows RQ-Transformer to memorize training samples.Nearest-neighbor analysis is then used to examine whether generated images are memorized or novel.

B.3. Ablation Study on Soft Labeling and Stochastic Sampling

The study evaluates training and sampling choices for RQ-Transformer, while also examining RQ-VAE’s coarse-to-fine reconstruction behavior. Stochastic sampling improves FID, especially when combined with soft labeling, and increasing quantization depth adds fine-grained image information.

  • Soft Labeling and Stochastic Sampling: 13.11 FID is achieved when stochastic sampling and soft labeling are used together, improving over the 14.06 baseline without either technique.Stochastic sampling alone achieves 13.24 FID, while soft labeling alone worsens performance to 14.87 FID.
  • Top-k and Top-p Sampling: Higher top-k values become effective as rejection-sampling acceptance rates decrease because they provide diverse, high-quality samples.The comparison uses varying acceptance rates and top-k values with top-p=1.0.
  • Coarse-to-Fine Approximation: RQ-VAE reconstructions gain more fine-grained information as the number of quantization depths increases.The coarse-to-fine behavior is evaluated using reconstruction, commitment, perceptual, and reconstruction-FID measures on ImageNet validation data.

C.2. The Effects of Adversarial and Perceptual Losses on Training of RQ-VAE

Adversarial and perceptual losses substantially improve RQ-VAE reconstruction clarity and local detail, while residual quantization further improves quality beyond a single code. Shared codebooks also provide better reconstruction than equivalent non-shared codebooks.

  • Residual Quantization Depth: RQ-VAE with D = 4 improves reconstruction quality over D = 1, even when adversarial and perceptual losses are absent.Increasing quantization depth allows the representation to capture more information than a single codebook pass.
  • Adversarial and Perceptual Losses: Adversarial and perceptual training produces clearer reconstructions with finer details than training without these losses.Without these losses, reconstructions are blurry because the codebook does not capture all local image details.
  • Shared versus Non-shared Codebooks: 4.73 rFID with one shared codebook is better than 5.73 rFID with four non-shared codebooks.A shared codebook of size 4,096 reaches 5.94 rFID, close to the non-shared-codebook result.
  • Additional Reconstructions and Samples: RQ-VAE reconstruction examples span LSUN, FFHQ, ImageNet, and CC-3M image-generation settings.The figures include unconditional, class-conditional, text-conditional, rejection-sampled, and nearest-neighbor visualizations.
Loading 2203.01941v2…