Autoregressive & Token-Based Image Generation

"They asked me to draw a picture, so naturally I spelled it. Two hundred and fifty-six little words, each one a patch of the sky, written left to right like any civilized sentence. The diffusion models think I am old-fashioned. I think they are just bad at reading."

An Autoregressive Image Model That Treats Pictures as Prose

Every system so far in this chapter denoises a continuous latent. There is an entirely different way to generate an image, one that descends from language modeling rather than from the noise processes of Chapter 33. The idea is to make an image discrete, a finite sequence of symbols from a fixed vocabulary, and then reuse the full transformer generation machinery of sequence models. This section is about that path: how it works, how it differs from diffusion, and why it matters more in 2026 than it did in 2023.

1. Making an Image Discrete: VQ Tokenization Intermediate

The enabling trick is the vector-quantized autoencoder (VQ-VAE / VQGAN), a discrete cousin of the VAE you built in Chapter 31. An encoder maps the image to a grid of feature vectors; each vector is then snapped to its nearest entry in a learned codebook of, say, 8192 vectors. The image is now represented by a grid of integer indices, one per spatial location, exactly like a 16-by-16 sentence of 256 "visual words". A decoder reconstructs pixels from the quantized grid. The codebook lookup is the only difference from a continuous autoencoder, and it is what turns a picture into a sequence of tokens that a transformer can model.

Figure 34.4.1 makes the conversion concrete: a continuous image becomes a small grid of integers drawn from a fixed vocabulary. From here, generating an image is generating that grid of integers, and the entire generative problem reduces to a familiar sequence-modeling problem. The "snap to nearest codebook entry" step is a hard lookup with no usable gradient, so ordinary backpropagation cannot train through it. The straight-through gradient estimator is the fix: on the backward pass it simply copies the gradient from the quantized vector back to the encoder as if the lookup were the identity. That one trick lets the codebook lookup train end to end and is the only piece of machinery beyond the ordinary VAE; the rest is the autoencoder you already know.

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

class VectorQuantizer(nn.Module):
    """Snap each encoder feature vector to its nearest codebook entry."""
    def __init__(self, n_codes=8192, dim=256):
        super().__init__()
        self.codebook = nn.Embedding(n_codes, dim)

    def forward(self, z):                 # z: (B, dim, H, W) encoder features
        b, c, h, w = z.shape
        flat = z.permute(0, 2, 3, 1).reshape(-1, c)            # (B*H*W, dim)
        # nearest codebook vector by squared Euclidean distance
        d = (flat.pow(2).sum(1, keepdim=True)
             - 2 * flat @ self.codebook.weight.t()
             + self.codebook.weight.pow(2).sum(1))
        ids = d.argmin(1)                                      # token id per cell
        zq = self.codebook(ids).view(b, h, w, c).permute(0, 3, 1, 2)
        zq = z + (zq - z).detach()        # straight-through estimator for gradients
        return zq, ids.view(b, h, w)

vq = VectorQuantizer()
z = torch.randn(1, 256, 16, 16)           # encoder output for one image
zq, token_grid = vq(z)
print("quantized features:", tuple(zq.shape))
print("token grid:", tuple(token_grid.shape),
      "-> sequence length:", token_grid.numel())

2. Generating the Tokens, Left to Right: Parti Intermediate

With the image reduced to a sequence of 256 integers, text-to-image generation becomes machine translation: encode the prompt into text tokens, then generate the image tokens one at a time, each conditioned on the prompt and on the image tokens generated so far. Parti is the canonical autoregressive system. It is a standard encoder-decoder transformer, the same architecture used for translation, with the prompt as the source sequence and the image-token grid (flattened in raster order) as the target sequence. Generation is the autoregressive next-token loop of any language model: predict a distribution over the 8192 codebook entries, sample, append, repeat 256 times, then hand the completed grid to the VQ decoder for pixels.

The appeal is unification: image generation reuses, verbatim, the transformer stack, the scaling laws, and the training infrastructure of language models. Parti scaled to 20 billion parameters and showed that image quality improves predictably with scale, just as in language. The cost is the same cost language models pay: generating $n$ tokens takes $n$ sequential forward passes. A 256-token image needs 256 steps, and higher resolutions explode quadratically, which is the weakness the next subsection attacks.

3. Generating in Parallel: MUSE and Masked Prediction Advanced

MUSE keeps the discrete tokens but throws out the left-to-right order. It is trained like the masked-language-modeling objective of BERT: hide a random subset of the image tokens and predict them all at once from the visible ones and the text. At generation time it starts from a fully masked grid and unmasks in a handful of rounds: predict every masked token in parallel, keep the most confident predictions, re-mask the rest, and repeat. A 256-token image that Parti generates in 256 sequential steps, MUSE produces in roughly 8 to 24 parallel rounds, an order-of-magnitude speedup. This parallel-unmasking schedule is the discrete-token analogue of the few-step samplers that accelerated diffusion in Chapter 33.

import torch

# MUSE-style generation: start from an all-masked token grid and fill it in
# over a handful of parallel rounds, committing only the most confident
# predictions each round, instead of one token per step left to right.
def muse_decode(predict_logits, n_tokens=256, n_steps=12, vocab=8192,
                text_cond=None):
    """Iterative parallel unmasking. predict_logits: model -> (n_tokens, vocab)."""
    MASK = vocab                                   # special mask id
    tokens = torch.full((n_tokens,), MASK)         # start fully masked
    for step in range(n_steps):
        logits = predict_logits(tokens, text_cond)         # predict ALL positions
        probs = logits.softmax(-1)
        sampled = torch.multinomial(probs, 1).squeeze(-1)  # one sample per cell
        conf = probs.gather(1, sampled[:, None]).squeeze(-1)
        # keep an increasing fraction of the most confident predictions
        keep = int(n_tokens * (step + 1) / n_steps)
        masked = (tokens == MASK)
        conf = torch.where(masked, conf, torch.full_like(conf, -1.0))
        keep_idx = conf.topk(min(keep, masked.sum().item())).indices
        tokens[keep_idx] = sampled[keep_idx]               # commit the confident ones
    return tokens                                          # complete token grid

Figure 34.4.2 contrasts the two orders on a single token grid: Parti commits one cell per forward pass, while MUSE commits the most confident handful of cells each round and re-masks the rest. That confidence-gated commit raises an obvious question. Why iterate at all, rather than predict all 256 tokens in a single parallel pass? Because each position is predicted independently from the same context, so a one-shot commit treats the tokens as if they were statistically independent when in fact they are tightly coupled (the patch left of a fox's ear constrains the patch right of it). Committing everything at once yields locally plausible but globally incoherent tiles. The fix is to commit only the few most confident predictions each round and re-mask the rest, so the next round predicts the still-uncertain positions conditioned on the now-fixed ones; coherence is bought back over a dozen rounds instead of all at once. This is the same dependency problem autoregression sidesteps by construction (each token sees every earlier one), and the confidence schedule is MUSE's way of recovering most of that benefit while keeping the parallelism.

The MUSE decode in this listing is the practical reason token-based generation became competitive on speed: 12 parallel rounds instead of 256 serial steps. It pairs the language-model integration of autoregression with a sampling cost closer to fast diffusion, which is exactly the combination the field wanted.

from diffusers import VQModel
import torch

# A pretrained VQ tokenizer: encode an image to tokens, decode tokens to pixels.
vq = VQModel.from_pretrained("CompVis/ldm-celebahq-256",
                             subfolder="vqvae").eval()
with torch.no_grad():
    latents = vq.encode(torch.randn(1, 3, 256, 256)).latents
    recon = vq.decode(latents).sample
print("vq latents:", tuple(latents.shape))

4. Why Tokens Are Resurging: Unified Multimodal Models Intermediate

Diffusion appeared to win the text-to-image race by 2023, so why does the token path matter? Because it is the only path that lets a single model both read and write images in one sequence. If an image is a sequence of tokens and text is a sequence of tokens, a single transformer can ingest interleaved text and image tokens and emit either, the foundation of modern multimodal models that answer questions about an image and then generate a new one in the same conversation. The features-and-descriptors thread of Chapter 10, which learned in Chapter 25 to replace hand-crafted descriptors with learned ones, reaches its conclusion here: the universal representation is now a discrete token an LLM can manipulate. Diffusion still produces the highest single-image fidelity, but the token path owns the unified-model frontier, and the two are increasingly hybridized.