VQGAN 和 Codebook 图像合成

VQGAN, short for Vector Quantized Generative Adversarial Network, is an image autoencoder with a twist. An autoencoder is a pair of networks: an encoder that squeezes an image into a compact representation and a decoder that rebuilds the image from it. In VQGAN the compact representation is not a grid of arbitrary numbers but a grid of indices into a codebook, a learned table of vectors that act like a vocabulary of visual parts such as a patch of fur, a stretch of sky or the corner of a window. Once an image is a grid of codebook indices, it looks like a sentence made of tokens, and any model that is good at predicting the next token in text can be trained to generate images.

The method was introduced by Patrick Esser, Robin Rombach and Bjorn Ommer of Heidelberg University in the paper Taming Transformers for High-Resolution Image Synthesis, first posted to arXiv on 17 December 2020 and revised in June 2021. Their motivation was cost. Transformers, the attention-based networks behind modern language models, compare every element of a sequence with every other element, so the cost grows with the square of the sequence length. A 256 by 256 image has 65,536 pixels, far too many to treat as a sequence. Convolutional neural networks, which process images with small sliding filters, are efficient at local texture but weaker at long-range structure. The paper's answer was to let a convolutional VQGAN learn the vocabulary of local parts and let a transformer learn how those parts are arranged.

The mechanism works in two stages. In stage one, the encoder downsamples the image by a factor f; the paper's default is f = 16, so a 256 by 256 image becomes a 16 by 16 grid of feature vectors. Each vector is replaced by its nearest neighbor in the codebook, a step called vector quantization. The default codebook holds 1,024 entries, with larger codebooks of 4,096 to 16,384 entries used for some datasets. The decoder then reconstructs the image from the quantized grid. Because choosing the nearest entry is not differentiable, gradients are copied straight through the quantizer during training, a trick called the straight-through estimator. In stage two, the encoder and codebook are frozen, every training image is converted to its 256 indices, and a transformer is trained to predict each index from the ones before it. To generate, the transformer samples indices one by one and the decoder turns the finished grid into pixels.

What separates VQGAN from its predecessor is the training loss. VQ-VAE, introduced by Aaron van den Oord, Oriol Vinyals and Koray Kavukcuoglu at DeepMind in November 2017, established the codebook idea and showed that discrete latents avoid posterior collapse, a failure where a powerful decoder learns to ignore the latent code. But VQ-VAE was trained mainly to reproduce pixels, which pushes reconstructions toward blur when the code is heavily compressed. VQGAN adds two terms. A perceptual loss compares images in the feature space of a pretrained network rather than pixel by pixel, and a patch-based discriminator, the adversarial half of a generative adversarial network (GAN), penalizes local patches that look fake. The weight on the adversarial term is set automatically from the ratio of gradient sizes at the decoder's last layer. The result is that VQGAN can compress much harder while staying sharp. The paper reports that its ImageNet VQGAN with 16,384 codes on a 16 by 16 grid reaches a reconstruction FID of 4.98, against 32.01 for the DALL-E tokenizer with 8,192 codes on a 32 by 32 grid, while using a sequence one quarter as long. FID, or Frechet Inception Distance, measures how far the statistics of generated images are from real ones; lower is better.

The generation results were strong for their time. With a GPT-2-medium-sized transformer of about 307 million parameters, the paper reports FIDs of 9.6 on the FFHQ face dataset and 10.2 on CelebA-HQ. On class-conditional ImageNet at 256 by 256 pixels it reports an FID near 16 without rejection sampling and 5.20 with it, compared with about 31 for VQ-VAE-2 without rejection, using roughly a tenth of VQ-VAE-2's estimated parameter count. For images larger than the training size, the transformer attends only to a sliding window of nearby tokens, which allowed the authors to produce megapixel images guided by segmentation or depth maps.

VQGAN quickly became a building block. The VQGAN-CLIP method, written up by Katherine Crowson, Stella Biderman and colleagues for ECCV 2022, paired a pretrained VQGAN with OpenAI's CLIP model and optimized the latent grid until CLIP judged the decoded image to match a text prompt, with no task-specific training. Google's ViT-VQGAN (October 2021) replaced the convolutional encoder and decoder with Vision Transformers and improved codebook learning, raising Inception Score on ImageNet from 70.6 for vanilla VQGAN to 175.1 and lowering FID from 17.04 to 4.17. MaskGIT (February 2022) kept VQ tokens but generated them in parallel by iteratively filling in masked positions. Parti (June 2022) scaled an autoregressive Transformer on ViT-VQGAN tokens to 20 billion parameters. The same Heidelberg group's Latent Diffusion paper (December 2021) reused the idea of working in the compressed space of a pretrained autoencoder, but ran a diffusion model there instead of a token predictor.

Compared with alternatives, codebook synthesis has clear trade-offs. Against pixel-space GANs, it separates compression from generation, so the generator can be a large, stable, likelihood-trained transformer rather than an adversarial network that is prone to training collapse. Against pixel-space diffusion, it is far cheaper because the model works on a few hundred tokens rather than tens of thousands of pixels. Against latent diffusion with continuous latents, discrete tokens fit naturally into the same sequence as text tokens, which is why tokenizers of this family appear in models that mix text and images. The cost is that quantization throws information away, so fine detail such as small text or faces in a crowd can be lost before generation even begins.

The main failure modes are well documented. Codebook collapse, where only a small fraction of entries is ever selected, wastes capacity; the Taming Transformers authors note that their ImageNet model effectively used 973 of its 16,384 entries. Autoregressive sampling is slow because a 256-token image needs 256 sequential steps, which motivated MaskGIT's parallel decoding. Raster-order generation, row by row, can produce inconsistent global layouts, and outputs depend heavily on sampling settings such as top-k. Finally, a 16-fold downsampling limits how much detail any downstream model can express.

In practice the requirements are modest by current standards. The paper reports that most of its transformers were sized to train with batch size 2 on a 12 GB GPU, typically using two to four GPUs, and that an updated ImageNet model trained for 45.8 days on a single A100. Training a new VQGAN still needs a large image dataset matched to the target domain, careful tuning of the adversarial weight and a warm-up period before the discriminator switches on. The authors released code and pretrained models, so encoding and decoding with an existing VQGAN needs far less compute than training one.