12.4.2 VQ-VAE & VQ-GAN

Overview

The VAE family (Lesson 12.2.1) compresses images into continuous latent vectors — points in a 16- or 64-dimensional real-valued space. That works well for the operations you’d want to do directly on the latent (interpolation, sampling), but it also locks you out of an entire family of downstream models: anything that wants to predict the next element in a sequence, like a transformer language model. Language models work on discrete tokens because their training objective is “predict the next token from a vocabulary.” Continuous latents have no vocabulary.

VQ-VAE, introduced by Van den Oord, Vinyals and Kavukcuoglu in 2017 [1], plugged this gap. It’s a VAE where the bottleneck is discrete: the encoder outputs continuous vectors, but each one is snapped to its nearest neighbour in a learned codebook of, say, 1,024 reference vectors before going into the decoder. The codebook indices are integers. An image becomes a grid of integers; a grid of integers is a sequence; a sequence can be modelled by a transformer.

VQ-GAN, introduced by Esser, Rombach and Ommer in 2020 [2], improved the reconstruction quality by adding a GAN-style adversarial loss on top — same encoder/decoder/codebook structure, plus a PatchGAN discriminator (like the one in Lesson 12.1.4) that pushes the decoder to produce sharper outputs. The combination — sharp VQ-GAN tokeniser + a transformer trained on the resulting token sequences — is the architecture behind Taming Transformers (Lesson 12.6.1) and the structural template for DALL-E 1, Parti, Muse, and many others.

Learning objectives

1. Explain why a continuous VAE bottleneck blocks transformer-based generation and how vector quantisation fixes it.
2. Implement the straight-through estimator for backpropagation through a non-differentiable nearest-neighbour lookup.
3. Describe the three losses in VQ-VAE training (reconstruction, codebook, commitment) and what each pulls toward what.
4. Predict how adding a GAN discriminator (VQ-GAN) changes the quality vs codebook-usage trade-off.

Quick start — train a tiny VQ-VAE on MNIST

60-second training run, 64-entry codebook, MNIST. After 3 epochs you’ll see most digits cleanly reconstructed through a discrete bottleneck.

pip install torch torchvision
python vqvae_demo.py

Expected output: three lines of training loss, then a final report of how many codebook entries the model is actually using (typically 30–50 of the 64 available — the rest are “dead”).

Core concepts

Concept 1 — Why discrete? Because transformers expect tokens

A standard VAE encodes an image to a continuous latent vector z ∈ R^d. To generate new images you sample z ~ N(0, I) and decode. This is fine for unconditional generation but it gives you nothing to model autoregressively — you can’t “predict the next dimension of z” the way a language model predicts the next word, because the dimensions of a continuous latent aren’t ordered or discrete.

VQ-VAE introduces a discrete bottleneck. The encoder outputs a (C, H, W) feature map of continuous vectors; each spatial location is snapped to its nearest neighbour in a learned codebook of K reference vectors. The output is no longer a continuous tensor but a grid of integers — codebook indices — of shape (H, W). An MNIST image at 7 × 7 resolution becomes a 49-long sequence of integers between 0 and K - 1.

Once images are sequences of integers, you can train a transformer to model them. The transformer learns P(token_n | token_1, ..., token_{n-1}); sampling proceeds left-to-right (or in any defined raster order); the resulting token sequence gets decoded back to an image by the VQ-VAE’s decoder.

Concept 2 — Vector quantisation and the straight-through estimator

The quantisation step is mechanically simple: for each encoder output vector z_e, find the nearest codebook entry e_k and replace z_e with e_k. The catch: nearest-neighbour lookup is not differentiable. You cannot backpropagate through argmin.

The fix is the straight-through estimator [3]. In the forward pass, send the quantised vector through the decoder normally. In the backward pass, copy the gradient that flowed into the quantised vector straight back to the encoder, as if quantisation were the identity function.

class VectorQuantizer(nn.Module):
    def forward(self, z_e):
        # Forward: nearest neighbour in codebook
        distances = (z_e.pow(2).sum(1, keepdim=True)
                     - 2 * z_e @ self.codebook.weight.t()
                     + self.codebook.weight.pow(2).sum(1))
        indices = distances.argmin(dim=1)
        z_q = self.codebook(indices)

        # Backward: straight-through. Forward returns z_q; backward sees identity.
        z_q = z_e + (z_q - z_e).detach()
        return z_q, indices

The key line is z_q = z_e + (z_q - z_e).detach(). Numerically the output equals z_q (the quantised vector). Symbolically the gradient flows through z_e — (z_q - z_e).detach() is a constant from autograd’s perspective. The encoder gets gradient as if the decoder were directly operating on its output.

This isn’t strictly correct (the true gradient is zero almost everywhere), but it’s a useful biased estimator and it works in practice. The codebook itself is updated separately via a dedicated codebook loss (Concept 3).

Concept 3 — Three losses balance encoder, codebook, and decoder

VQ-VAE training combines three losses:

L = L_recon + L_codebook + β · L_commitment

Reconstruction loss (L_recon) is straight MSE between the input and the decoder’s output. This pulls the encoder, codebook, and decoder together to minimise reconstruction error. Backpropagation flows through the decoder, through the straight-through quantisation, and into the encoder.

Codebook loss (L_codebook = ‖sg[z_e] - e_k‖² where sg[·] is stop-gradient) pulls the codebook entries toward the encoder outputs. The encoder’s output is detached (sg) so this loss only updates the codebook. Effectively: “for each chosen codebook entry, move it closer to the encoder vectors that mapped to it.”

Commitment loss (L_commitment = ‖z_e - sg[e_k]‖²) pulls the encoder outputs toward whichever codebook entries they were assigned to. The codebook is detached so this loss only updates the encoder. Effectively: “encoder, commit to the codebook entries that already exist instead of drifting to new regions of latent space.”

The β coefficient (default 0.25 per the paper) controls how strongly the encoder is pressured to stick close to codebook entries. Too low and the encoder drifts; too high and the codebook can’t refine its positions.

def vq_loss(z_e, z_q, commitment_cost=0.25):
    codebook_loss = F.mse_loss(z_q, z_e.detach())
    commitment_loss = F.mse_loss(z_q.detach(), z_e)
    return codebook_loss + commitment_cost * commitment_loss

One common failure mode: codebook collapse. Most codebook entries get used early in training; a few unlucky ones don’t get any encoder vectors assigned to them; with no gradient signal, they stay where they are forever while the encoder learns to avoid them. After a few epochs you have a 1,024-entry codebook with only 200 entries in active use. Modern implementations mitigate this with codebook reset tricks (replace dead entries with active ones) or with exponential-moving-average updates instead of gradient descent on the codebook.

Concept 4 — VQ-GAN: adding a discriminator for sharpness

VQ-VAE produces reconstructions but they’re slightly blurry — the standard MSE-loss-produces-soft-outputs problem from Lesson 12.2.1. Esser et al.’s VQ-GAN [2] fixed this by stacking a PatchGAN discriminator (same patch-based discriminator as Pix2Pix in 12.1.4) on the reconstruction and adding an adversarial loss:

L_VQGAN = L_recon + L_codebook + β · L_commitment
        + λ · L_perceptual         (LPIPS feature loss)
        + γ · L_GAN                (PatchGAN adversarial)

The perceptual loss replaces MSE on raw pixels with MSE on intermediate features of a pre-trained VGG (the same idea as NST in 12.4.1), which is more aligned with how humans judge image similarity. The GAN loss is computed on the patch-level discriminator’s classifications of real vs reconstructed patches. Together they produce reconstructions that are not just numerically close to the input but perceptually close, with crisp textures.

The architectural shape is identical to VQ-VAE — same encoder, same codebook, same straight-through quantisation, same decoder. The only differences are the loss function and the addition of the discriminator network. The cost is training complexity: GAN training is unstable in ways VAE training isn’t; warm-up schedules and careful loss weighting matter.

Exercises

python vqvae_demo.py

import torch
import torch.nn as nn
import torch.nn.functional as F

class VectorQuantizer(nn.Module):
    def __init__(self, codebook_size: int, embedding_dim: int, commitment_cost: float):
        super().__init__()
        # TODO: create codebook as nn.Embedding(codebook_size, embedding_dim)
        # TODO: initialise its weights uniformly in [-1/K, 1/K]
        # TODO: store commitment_cost as self attribute
        pass

    def forward(self, z_e):
        b, c, h, w = z_e.shape
        z_e_flat = z_e.permute(0, 2, 3, 1).reshape(-1, c)

        # TODO: compute pairwise squared distances between z_e_flat and codebook
        # (use the (a-b)^2 = a^2 - 2ab + b^2 expansion)

        # TODO: indices = argmin over codebook dim
        # TODO: z_q_flat = self.codebook(indices)

        # TODO: codebook_loss = MSE(z_q_flat, z_e_flat.detach())
        # TODO: commitment_loss = MSE(z_q_flat.detach(), z_e_flat)
        # TODO: loss = codebook_loss + self.commitment_cost * commitment_loss

        # TODO: apply straight-through estimator
        # z_q_flat = z_e_flat + (z_q_flat - z_e_flat).detach()

        # TODO: reshape z_q_flat back to (b, c, h, w) shape

        # TODO: return z_q, loss, indices.reshape(b, h, w)
        pass

class VectorQuantizer(nn.Module):
    def __init__(self, codebook_size, embedding_dim, commitment_cost):
        super().__init__()
        self.codebook = nn.Embedding(codebook_size, embedding_dim)
        nn.init.uniform_(self.codebook.weight, -1.0/codebook_size, 1.0/codebook_size)
        self.commitment_cost = commitment_cost

    def forward(self, z_e):
        b, c, h, w = z_e.shape
        z_e_flat = z_e.permute(0, 2, 3, 1).reshape(-1, c)

        d = (z_e_flat.pow(2).sum(1, keepdim=True)
             - 2 * z_e_flat @ self.codebook.weight.t()
             + self.codebook.weight.pow(2).sum(1))
        indices = d.argmin(dim=1)
        z_q_flat = self.codebook(indices)

        codebook_loss = F.mse_loss(z_q_flat, z_e_flat.detach())
        commitment_loss = F.mse_loss(z_q_flat.detach(), z_e_flat)
        loss = codebook_loss + self.commitment_cost * commitment_loss

        z_q_flat = z_e_flat + (z_q_flat - z_e_flat).detach()
        z_q = z_q_flat.reshape(b, h, w, c).permute(0, 3, 1, 2)
        return z_q, loss, indices.reshape(b, h, w)

Summary

Common pitfalls

• Forgetting the straight-through trick (z_e + (z_q - z_e).detach()) — the encoder receives no gradient and never trains.
• Mixing up stop-gradient placement in the codebook vs commitment losses — swapping them breaks the encoder/codebook decoupling and training stalls.
• Codebook collapse (most entries dead). Reset dead entries periodically or use EMA codebook updates instead of SGD.
• Treating VQ-VAE as a generative model on its own. It’s a tokeniser; you need a second-stage model (transformer, diffusion-over-tokens, etc.) to actually generate novel sequences.
• Using a too-small commitment cost (< 0.1). The encoder drifts away from the codebook and reconstruction quality degrades over training rather than improving.

References