12.6.1 Taming Transformers

Overview

Lesson 12.4.2 ended with a question: VQ-GAN turns images into sequences of integer codebook indices, but the VQ-GAN itself isn’t a generative model — it’s an autoencoder. To actually generate novel images, you need a second-stage model that learns the distribution over those integer sequences and can sample from it.

Esser, Rombach and Ommer’s “Taming Transformers for High-Resolution Image Synthesis” (2021) [1] — the same paper that introduced VQ-GAN — answered the question with a GPT-style autoregressive transformer trained on the token sequences a VQ-GAN encoder produces. The combination is structurally beautiful: VQ-GAN handles the what each pixel looks like problem (mapping continuous RGB to discrete tokens and back); the transformer handles the which tokens go together problem (modelling the joint distribution over token sequences). Generation runs left-to-right, one token at a time, exactly like a language model generates text.

This architectural template — discrete tokeniser + autoregressive transformer — became the foundation for OpenAI’s DALL-E 1 [2], Google’s Parti [3], and the masked-prediction variant Muse [4]. Today most state-of-the-art text-to-image models either use this template directly or are explicitly motivated by it.

Learning objectives

1. Describe the two-stage VQ-GAN + Transformer architecture and which stage handles which problem.
2. Explain causal self-attention and why the lower-triangular attention mask is what makes autoregressive generation possible.
3. Train a tiny autoregressive transformer over token sequences and sample from it left-to-right.
4. Predict how raster ordering vs other token orderings affects what the transformer can learn.

Quick start — train ImageGPT on synthetic token sequences

A minimal GPT-style transformer that learns to predict the next integer in a 49-long sequence (a 7×7 token grid). For this demo the training data is synthetic random integers; in a real setup the training sequences come from running a trained VQ-VAE/VQ-GAN encoder over an image dataset.

pip install torch
python transformer_image_demo.py

5 epochs, ~30 seconds on CPU. Output: per-epoch loss, then 4 sampled sequences. With the random training data the model can’t learn meaningful structure; the real value is seeing the architecture’s pieces line up.

Core concepts

Concept 1 — Two stages: tokeniser and generator

The Taming Transformers architecture decouples two responsibilities:

Stage 1 is what turns the continuous, high-dimensional image domain into a discrete, low-dimensional token domain that a transformer can model. The VQ-GAN encoder takes a 256 × 256 × 3 image and produces, say, a 16 × 16 = 256-long sequence of codebook indices. The decoder takes a sequence of indices back to an RGB image.

Stage 2 is just a language model over the token vocabulary. The vocabulary size is the VQ-GAN’s codebook size (typically 1,024 to 16,384). The sequence length is the spatial grid size (256 for 16 × 16, 1,024 for 32 × 32). Train it on the token sequences from your image dataset; sample new sequences autoregressively; decode each sample with the VQ-GAN to get a generated image.

The two stages are trained separately. Stage 1 is trained first, with the full VQ-GAN losses (reconstruction + perceptual + GAN — see 12.4.2). Stage 2 is trained second, with the standard next-token-prediction loss, while the VQ-GAN’s weights are frozen. The stages can be developed independently and each can be improved without retraining the other.

Concept 2 — Causal self-attention and the autoregressive constraint

The transformer’s training objective is next-token prediction: given tokens t_1, ..., t_{n-1}, predict t_n. For this to work, the model must not be able to “cheat” by looking at future tokens during training. The mechanism that enforces this is the causal self-attention mask.

In a standard attention layer, every position in the sequence attends to every other position. For autoregressive training, you mask out attention to future positions by setting those attention weights to negative infinity (so the softmax sends them to zero):

# Lower-triangular mask blocks future positions
mask = torch.triu(torch.ones(seq_len, seq_len), diagonal=1).bool()
# In the attention call:
attention_output, _ = nn.MultiheadAttention(...)(
    query=x, key=x, value=x, attn_mask=mask, need_weights=False,
)

The mask is a (seq_len, seq_len) boolean matrix where mask[i, j] = True (block) if j > i (j is in the future from position i’s perspective). The attention layer interprets True positions as “do not attend to this.”

With this mask, position i can attend to positions 1 through i (everything that came before, including itself), but not to positions i + 1 through seq_len. The model can be trained on the entire sequence in parallel (one forward pass produces predictions for every position simultaneously) but each prediction only depends on past tokens.

Sampling at inference time is sequential. Start with a special “begin” token (or just zero), predict the distribution over the next token, sample from it, append the sampled token, repeat until you have a full sequence:

@torch.no_grad()
def generate(model, num_samples, seq_len, temperature=1.0):
    tokens = torch.zeros(num_samples, 1, dtype=torch.long)
    for i in range(seq_len - 1):
        logits = model(tokens)[:, -1, :] / temperature
        probs = torch.softmax(logits, dim=-1)
        next_token = torch.multinomial(probs, num_samples=1)
        tokens = torch.cat([tokens, next_token], dim=1)
    return tokens

The temperature parameter controls sample diversity: low values (< 1.0) make sampling more deterministic and conservative; high values (> 1.0) make it more random and creative. Same parameter that text language models expose.

Concept 3 — Raster ordering: turning 2D image structure into a 1D sequence

The VQ-GAN produces a 2D grid of tokens (e.g. 16 × 16). The transformer expects a 1D sequence. The standard choice is raster order: read the grid row by row, left to right (the same way you’d read a page in English). Token 0 is the top-left; token 15 is top-right; token 16 is the start of the second row; token 255 is bottom-right.

This is arbitrary in principle — you could use any ordering — but raster has been overwhelmingly the standard. The model has to learn the 2D spatial structure through the 1D ordering: it learns that token i + 16 is “the position directly below token i,” that local context spans positions (i - 1, i, i + 1, i + 15, i + 16, i + 17), and so on.

Other orderings have been tried. Spiral orderings let the model build outward from the centre. Z-order (Morton encoding) preserves 2D locality better — adjacent positions in the sequence tend to be spatially close in the image. MaskGIT [4] uses a non-autoregressive parallel decoding scheme that doesn’t have a fixed ordering at all. In practice, raster + a sufficiently large model works well enough that the field has mostly stopped exploring alternatives.

Exercises

python transformer_image_demo.py

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

class CausalSelfAttention(nn.Module):
    def __init__(self, embed_dim=128, num_heads=4, seq_len=49):
        super().__init__()
        self.num_heads = num_heads
        self.head_dim = embed_dim // num_heads

        # TODO: q, k, v projections (each Linear(embed_dim, embed_dim))
        # TODO: output projection Linear(embed_dim, embed_dim)

        # TODO: causal mask buffer (lower-triangular, True where blocked)

    def forward(self, x):
        b, s, d = x.shape
        # TODO: project x to q, k, v
        # TODO: reshape each to (b, num_heads, s, head_dim)
        # TODO: compute attention scores: (q @ k.transpose(-2, -1)) / sqrt(head_dim)
        # TODO: apply causal mask: scores.masked_fill_(mask[:s, :s], float('-inf'))
        # TODO: softmax over last dim
        # TODO: out = attn @ v, reshape to (b, s, d)
        # TODO: apply output projection
        pass

class CausalSelfAttention(nn.Module):
    def __init__(self, embed_dim=128, num_heads=4, seq_len=49):
        super().__init__()
        self.num_heads = num_heads
        self.head_dim = embed_dim // num_heads

        self.qkv = nn.Linear(embed_dim, 3 * embed_dim)
        self.proj = nn.Linear(embed_dim, embed_dim)

        mask = torch.triu(torch.ones(seq_len, seq_len), diagonal=1).bool()
        self.register_buffer("causal_mask", mask)

    def forward(self, x):
        b, s, d = x.shape
        qkv = self.qkv(x).reshape(b, s, 3, self.num_heads, self.head_dim)
        q, k, v = qkv.unbind(dim=2)
        # (b, h, s, dh)
        q, k, v = [t.transpose(1, 2) for t in (q, k, v)]

        scores = (q @ k.transpose(-2, -1)) / math.sqrt(self.head_dim)
        scores = scores.masked_fill(self.causal_mask[:s, :s], float("-inf"))
        attn = F.softmax(scores, dim=-1)

        out = attn @ v                       # (b, h, s, dh)
        out = out.transpose(1, 2).reshape(b, s, d)
        return self.proj(out)

Summary

Common pitfalls

• Forgetting to mask future positions in causal self-attention. Without the mask, the model “cheats” during training by looking at the answer and produces garbage at inference time.
• Confusing the tokeniser’s vocab size with the transformer’s vocab size — they must match exactly. A 1024-entry VQ-GAN paired with a 512-vocab transformer can’t sample legal tokens.
• Sampling with temperature = 1.0 and complaining about noisy outputs. The sweet spot for image-token transformers is usually 0.7-0.9.
• Naively scaling to longer sequences. Attention is O(N²) in sequence length; doubling the sequence quadruples the compute. Past 1024-2048 tokens you need sparse-attention or parallel-decoding variants.
• Training the transformer and the VQ-GAN jointly. They can be trained jointly in principle but stability is poor; the canonical recipe is sequential: train the VQ-GAN first to convergence, then freeze it and train the transformer.

References