Vision & Multimodal advanced 8 min read 10 flashcards

Discrete Visual Tokenisers

How VQ-VAE, VQGAN and FSQ turn an image into a short sequence of integers, why codebook collapse wrecks half of them, and what the compression ratio costs in reconstruction fidelity.

An autoregressive transformer needs a finite vocabulary. Pixels are not one: a 256×256 RGB image is 196,608 integers in \([0, 255]\), and modelling that sequence directly is hopeless. A visual tokeniser is the compressor that makes the problem tractable, mapping the image to something like 256 or 1,024 integers drawn from a codebook of a few thousand entries. Everything downstream, including whether a unified model can generate images at all, is downstream of how well that compressor works.

Vector quantisation, and the gradient problem it creates

VQ-VAE encodes the image to a grid of continuous vectors \(z_e(x)\), then replaces each with its nearest neighbour in a learned codebook \(\{e_k\}_{k=1}^{K}\) (van den Oord et al., 2017, arXiv:1711.00937):

\[q(x)_{ij} = \arg\min_k \lVert z_e(x)_{ij} - e_k \rVert_2\]

The \(\arg\min\) has zero gradient almost everywhere, so training uses a straight-through estimator: copy the gradient from the decoder input straight back to the encoder output as if quantisation were the identity. Two auxiliary terms keep the two sides from drifting apart, a codebook loss pulling \(e_k\) toward the encoder outputs assigned to it, and a commitment loss \(\beta \lVert z_e(x) - \mathrm{sg}[e] \rVert^2\) penalising the encoder for producing vectors far from any code.

VQGAN kept this backbone and changed the objective: add a patch discriminator and a perceptual loss so the decoder is optimised for what looks right rather than for L2, then train an autoregressive transformer over the resulting code grid (Esser, Rombach and Ommer, 2021, arXiv:2012.09841). This is the recipe underneath most token-based image generation, and downsampling factor \(f\) is the knob: \(f{=}16\) turns 256×256 into a 16×16 grid, so 256 tokens per image.

Codebook collapse

The failure that dominates practice is that most of the codebook goes unused. Early in training, a handful of codes sit near the bulk of encoder outputs, win every nearest-neighbour lookup, and receive all the gradient; the rest are never selected, never updated, and stay dead forever. A nominal vocabulary of 8,192 can end up with a few hundred live entries, and effective capacity collapses with it. The standard countermeasures, EMA codebook updates, code reseeding, splitting overused codes, entropy penalties and low-dimensional lookup, are all patches on the same wound.

Finite scalar quantisation removes the wound instead. FSQ projects the latent to a handful of dimensions, typically fewer than ten, and quantises each one independently to a small fixed set of scalar values; the implicit codebook is the Cartesian product of those sets. There is no learned codebook to collapse, so commitment losses, reseeding and entropy penalties all disappear, and it reaches comparable codebook sizes and competitive results on MaskGIT-style generation and on UViM tasks (Mentzer et al., 2023, arXiv:2309.15505).

The rate-fidelity trade nobody escapes

The tokeniser sets a hard information bottleneck. With \(K\) codes on an \(n \times n\) grid the representation carries at most \(n^2 \log_2 K\) bits: a 16×16 grid with 8,192 codes is 3,328 bits, about 416 bytes, for the whole image. Anything not in those bits is invented by the decoder.

That is why token-based image models render plausible-looking but wrong text inside images, why faces at small scale come out subtly off, and why measuring a tokeniser by reconstruction FID matters more than measuring the generator. Push the grid finer and fidelity improves while the transformer's sequence length grows quadratically; push it coarser and generation gets cheap and lossy. Continuous-latent approaches sidestep the quantiser entirely at the price of needing a diffusion or flow objective rather than a categorical one.

Check yourself

10 flashcards for this concept

Click a card to reveal the answer.

Drill the whole track