Source-linked AI summary
Vector-quantized Image Modeling with Improved VQGAN
Jiahui Yu, Xin Li, Jing Yu Koh, Han Zhang, Ruoming Pang, James Qin, Alexander Ku, Yuanzhong Xu, Jason Baldridge, Yonghui Wu
TL;DR
The paper asks whether language-model-style pretraining can support image generation and understanding. It introduces VIM, which autoregressively models image tokens produced by an improved ViT-VQGAN, and reports stronger synthesis results than vanilla VQGAN.
Problem
Language-model pretraining improves generation and understanding in NLP, motivating investigation of whether analogous pretraining can support image generation and understanding.
Method
VIM uses an improved Vision-Transformer-based VQGAN to encode images into discrete tokens, then trains a Transformer to predict rasterized tokens autoregressively.
Results
IS 175.1 and FID 4.17 on ImageNet class-conditioned synthesis improve over vanilla VQGAN's IS 70.6 and FID 17.04.
Takeaways & Limitations
Better image quantization improves both computational efficiency and reconstruction quality, supporting image synthesis and representation evaluation with VIM.
Takeaways & Limitations
For unsupervised learning, perceptual loss is excluded because supervised VGG pretraining may leak supervision into linear-probe measurements.
Abstract
from arXiv · showhide
Pretraining language models with next-token prediction on massive text corpora has delivered phenomenal zero-shot, few-shot, transfer learning and multi-tasking capabilities on both generative and discriminative language tasks. Motivated by this success, we explore a Vector-quantized Image Modeling (VIM) approach that involves pretraining a Transformer to predict rasterized image tokens autoregressively. The discrete image tokens are encoded from a learned Vision-Transformer-based VQGAN (ViT-VQGAN). We first propose multiple improvements over vanilla VQGAN from architecture to codebook learning, yielding better efficiency and reconstruction fidelity. The improved ViT-VQGAN further improves vector-quantized image modeling tasks, including unconditional, class-conditioned image generation and unsupervised representation learning. When trained on ImageNet at \(256\times256\) resolution, we achieve Inception Score (IS) of 175.1 and Fr'echet Inception Distance (FID) of 4.17, a dramatic improvement over the vanilla VQGAN, which obtains 70.6 and 17.04 for IS and FID, respectively. Based on ViT-VQGAN and unsupervised pretraining, we further evaluate the pretrained Transformer by averaging intermediate features, similar to Image GPT (iGPT). This ImageNet-pretrained VIM-L significantly beats iGPT-L on linear-probe accuracy from 60.3% to 73.2% for a similar model size. VIM-L also outperforms iGPT-XL which is trained with extra web image data and larger model size.
1 INTRODUCTION
The paper introduces Vector-quantized Image Modeling (VIM), which applies autoregressive Transformer prediction of discrete image tokens to image generation and understanding. Its two-stage design combines an improved ViT-VQGAN image quantizer with Transformer-based token modeling, emphasizing efficiency and reconstruction quality.
- VIM applies autoregressive modeling of discrete image tokens to both image generation and image understanding.The approach is motivated by the success of generative pretraining in language modeling.
- Stage 1 encodes 256×256 images into 32×32 discretized latent codes using a Vision-Transformer-based VQGAN with codebook size 8192.The proposed ViT-VQGAN improves architecture and codebook learning while combining logit-laplace, ℓ2, adversarial, and perceptual losses.
- Stage 2 trains a Transformer to autoregressively predict 1024 rasterized image tokens encoded by the learned ViT-VQGAN.A prepended class-id token supports class-conditioned synthesis, while averaged intermediate features support linear-probe evaluation.
- Better quantization improves VIM by speeding Stage 2 training and reducing information loss from the original image in pixel space.The paper identifies computational efficiency and reconstruction quality as key properties of the image quantizer.
- The evaluation covers quantizer reconstruction quality, unconditional and class-conditioned image synthesis, and representation quality through linear-probe accuracy.The study uses ℓ1, ℓ2, log-laplace, IS, and FID metrics, with linear separability as the representation criterion.
2 RELATED WORK
Related work includes image synthesis methods that model discrete image representations and generative-pretraining approaches for image recognition. These lines of work motivate evaluating image-generation models for representation learning as well as synthesis.
- Image Synthesis: VQVAE decomposes image generation into discrete representation learning followed by autoregressive modeling of latent variables.The related methods include VQVAE, VQGAN, and other generative architectures.
- Image Recognition with Generative Pretraining: Generative image models have been studied for representation learning, but their performance has often not surpassed self-supervised methods based on auxiliary classification tasks.BigBiGAN is cited as an example of a generation-based model matching other self-supervised methods.
3 VECTOR-QUANTIZED IMAGES WITH VIT-VQGAN
ViT-VQGAN replaces CNN-based VQGAN components with Vision Transformers and introduces codebook-learning improvements to improve efficiency, reconstruction quality, and codebook usage. Its training combines vector-quantization, reconstruction, perceptual, adversarial, and logit-laplace objectives, with perceptual loss excluded for unsupervised evaluation to avoid supervised-information leakage.
- 3.1 VQGAN WITH VISION TRANSFORMERS: ViT-VQGAN replaces VQVAE and VQGAN CNN encoders and decoders with Vision Transformers, improving accelerator efficiency and reconstruction quality.The encoder maps 256×256 images into 32×32=1024 image tokens, and the decoder reverses this mapping.
- 3.1 VQGAN WITH VISION TRANSFORMERS: ViT-VQGAN achieves a better speed-quality trade-off than CNN-VQGAN, further accelerating Stage 2 training.Throughputs are benchmarked using the same 128 CloudTPUv4 devices.
- 3.2 CODEBOOK LEARNING: Low codebook usage creates dead codes, worse reconstructions, and reduced diversity, motivating factorized and ℓ2-normalized code improvements.Factorized codes reduce the lookup space before projecting matched codes into the embedding space, while normalization maps latent variables onto a sphere.
- 3.2 CODEBOOK LEARNING: ℓ2-normalized codes convert Euclidean distance into cosine similarity, improving training stability and reconstruction quality.The normalization is applied to both encoded latent variables and codebook latent variables.
- 3.3 VIT-VQGAN TRAINING LOSSES: The default objective combines vector-quantization, adversarial, perceptual, logit-laplace, and ℓ2 losses with weights 1, 0.1, 0.1, 0.1, and 1.0, respectively.The authors report that logit-laplace loss contributes to codebook usage, while ℓ2 and perceptual losses contribute significantly to FID.
- 3.3 VIT-VQGAN TRAINING LOSSES: Perceptual loss is excluded from unsupervised-learning quantizers because its supervised VGG pretraining might leak supervision into Stage 2 linear-probe evaluation.The loss is retained for unconditional and class-conditioned synthesis because it produces higher-fidelity reconstructions.
4 VECTOR-QUANTIZED IMAGE MODELING
VIM encodes images as rasterized discrete tokens with ViT-VQGAN, then trains a decoder-only Transformer to model those tokens autoregressively. The same pretrained Transformer supports unconditional and class-conditioned generation, while averaged intermediate features enable linear probing.
- Vector-quantized image modeling: A decoder-only Transformer predicts each image token autoregressively using causal attention over the token sequence.The model minimizes the negative log-likelihood of the image-token data.
- Vector-quantized image modeling: Images are encoded into 32×32 discrete latent codes with an 8192-entry codebook, then flattened into 1024 raster-ordered tokens.The ViT-VQGAN provides the Stage 1 image quantizer for Stage 2 modeling.
- Image synthesis: Unconditional generation samples tokens from the output softmax and decodes them through the ViT-VQGAN decoder.The default setup uses 256×256 inputs, 32×32 latent codes, an 8192-sized codebook, and sequence length 1024.
- Image synthesis: Class-conditioned synthesis prepends a class-id token before the image tokens and autoregressively decodes the remaining tokens.Separate embeddings are learned for class-id and image tokens.
- Unsupervised learning: For image understanding, a frozen intermediate Transformer output is averaged across token positions and mapped through a learnable softmax classifier.The most discriminative feature typically comes from a middle Transformer block.
5 EXPERIMENTS
Experiments evaluate ViT-VQGAN across quantization, image synthesis, and unsupervised representation learning. The proposed system improves reconstruction and generation metrics over baselines and substantially raises ImageNet linear-probe accuracy over iGPT.
- Image quantization: ViT-VQGAN is trained separately on CelebA-HQ, FFHQ, and ImageNet with Small, Base, and Large encoder-decoder configurations.Training uses batch size 256, 500,000 steps, and CloudTPUv4 hardware.
- Image quantization: ViT-VQGAN achieves better reconstruction FID than vanilla VQGAN with an 8192-sized codebook, even without hierarchical codebooks or Gumbel-Softmax.The comparison uses validation reconstructions on ImageNet, CelebA-HQ, and FFHQ.
- Image quantization: ViT-based VQGANs outperform CNN-based VQGANs in quality and throughput, while StyleGAN discrimination improves stability and reconstruction quality over PatchGAN.Among ViT model sizes, the smallest model has lower quality but higher throughput than larger configurations.
- Image quantization: Factorized codes with latent dimensions reduced from 256 to 16 or 8 improve reconstruction quality, whereas removing ℓ2-normalization substantially worsens results.These findings come from the codebook-learning ablations.
- Image synthesis: 7.0 FID on CelebA-HQ and 5.3 FID on FFHQ are achieved for unconditional synthesis without top-k or nucleus sampling heuristics.These scores are reported as significantly better than VQGAN.
- Image synthesis: 175.1 IS and 4.17 FID on ImageNet improve over vanilla VQGAN's 70.6 IS and 17.04 FID for 256×256 class-conditioned synthesis.Classifier-based rejection sampling further reaches 3.04 FID and 321.7 IS.
- Unsupervised learning: 73.2% ImageNet linear-probe accuracy for VIM-L improves over iGPT-L's 60.3% at similar model size and also surpasses the larger, web-data-trained iGPT-XL.Alternative CNN-based VQGAN and DALL-E dVAE quantizers perform worse than ViT-VQGAN in the comparison.
6 ETHICS
The paper highlights dataset and model-generation biases, privacy and copyright concerns, and the potential misuse of image generation systems.
- FFHQ: FFHQ has broad demographic coverage, but documented imbalances in skin tone and correlations among attributes remain.The cited analysis also examines biases in face-generation datasets and algorithms.
- CelebA-HQ: CelebA-HQ biases generations toward attractive people, especially those aged twenty to forty, despite coverage across perceived genders and ethnicities.The authors report that their generations show a similar age-related bias.
- ImageNet: ImageNet provides a significant but biased sample of scenes and objects because of biases in image selection and label annotation.The paper’s ImageNet generation results therefore cover only part of the world’s visual diversity.
- Broader concerns: Image generation models raise concerns about demographic biases, privacy, copyright, and possible creation of fake media.The paper points to prior work documenting problematic biases and potential misuse.
A LINEAR-PROBE ON IMAGENET
Figure 3 reports linear-probe accuracy from different layers of a pretrained VIM-Base Transformer model.
- Linear-probe evaluation: Figure 3 compares linear-probe accuracy across layers of a pretrained VIM-Base Transformer.The figure caption identifies the evaluated model and probing metric, but provides no values in the supplied passage.
B MODEL SIZES OF CLASS-CONDITIONED IMAGENET SYNTHESIS
The paper compares class-conditional ImageNet synthesis across Stage 2 Transformer sizes using FID without rejection sampling.
- Model-size comparison: The authors present class-conditioned ImageNet synthesis results for different Stage 2 Transformer sizes and compare them with VQGAN.The supplied passage does not include the corresponding numerical results.
- Model-size comparison: Table 8 reports FID for class-conditional ImageNet synthesis across different Stage 2 Transformer sizes.The supplied caption specifies that results exclude rejection sampling.
C IMPLEMENTATION DETAILS OF FACTORIZED CODEBOOK
The factorized codebook implementation reduces encoded embeddings to a lower-dimensional space for code lookup, with further details illustrated in Figure 4.
- Factorized codebook: A linear projection reduces each encoded embedding to a low-dimensional variable space before code lookup.This projection is part of the factorized codebook implementation.
- Factorized codebook: Figure 4 illustrates the factorized codes and codebook details used by the implementation.The supplied caption does not specify the figure’s axes or numerical comparisons.
D MORE SAMPLES ON CLASS-CONDITIONED IMAGENET SYNTHESIS
Figures 5 and 6 present uncurated samples from class-conditioned generation trained on ImageNet.
- Figure 5 presents an uncurated sample set from class-conditioned generation trained on ImageNet.
- Together, the figures provide additional visual samples from the ImageNet-trained class-conditioned generation setting.
- Figure 6 presents an uncurated sample set from class-conditioned generation trained on ImageNet.