Core ML

Vector-Quantized Models

Snap every encoder output to the nearest entry in a fixed, learned codebook, turning a continuous latent space into a finite set of discrete symbols.

An encoder's continuous output gets replaced by whichever entry in a fixed codebook sits closest to it, so the model works with a finite set of discrete codes instead of raw continuous vectors

Why Does This Exist?

Variational autoencoders give every image a point in a continuous latent space — any real-valued vector is a valid code, in principle. That's exactly the property that makes a different problem hard: generating a new image autoregressively, one piece at a time, the way a language model generates one token at a time. Autoregressive generation over a continuous space has no natural notion of "the next value" — there's no fixed vocabulary to pick from, no way to run the same next-token machinery that works so well for text.

Text doesn't have this problem because words already come from a finite vocabulary. Vector-quantized models give images, audio, and other continuous data that same property: force every latent value onto one of a fixed, finite set of learned codes, and suddenly a discrete, autoregressive generation loop — the same kind that works for text — becomes possible.

Think of It Like This

A paint-by-numbers set with a fixed palette

A photograph has effectively infinite possible colors. A paint-by-numbers kit doesn't — it ships with a fixed set of numbered paint pots, say sixty-four of them, and every region of the image gets assigned to whichever pot's color is the closest match. The image loses some fidelity in that snap-to-nearest-pot step, but what you gain is enormous: the whole painting can now be described as a grid of pot numbers, a finite, countable vocabulary, instead of a grid of continuous color values.

Vector quantization does the same thing to a learned representation: it ships with a codebook of vectors instead of paint pots, and every point the encoder produces gets snapped to whichever codebook entry is nearest.

How It Actually Works

The codebook and the nearest-neighbor snap

A codebook is a fixed-size table of learned vectors, $e_1, \dots, e_K$ — typically a few hundred to a few thousand entries. Given the encoder's continuous output $z$ for some input patch, vector quantization replaces it with whichever codebook entry is closest:

z_q = e_k, \quad k = \arg\min_j \lVert z - e_j \rVert^2

$z_q$ is the quantized code that actually gets passed to the decoder. The codebook entries themselves are learned parameters, updated during training just like any other weight, so the "palette" adapts to whatever kinds of patches the data actually contains rather than being fixed in advance.

The gradient problem, and the straight-through fix

argmin is not differentiable — a tiny nudge to $z$ either doesn't change which codebook entry is nearest, or it jumps discontinuously to a different one. Backpropagating through the quantization step directly is impossible in the ordinary sense.

The fix, again, is a straight-through estimator: in the forward pass, use $z_q$ (the snapped, discrete code) for everything downstream. In the backward pass, treat the quantization step as if it had been the identity function, and let the gradient that arrives at $z_q$ pass straight through to $z$ unchanged. The encoder still receives a usable training signal, even though the function it's actually being credited with computing wasn't the one that ran forward. A separate loss term pulls each codebook entry toward the encoder outputs assigned to it, which is what actually updates the codebook itself.

What becomes possible once the latent space is discrete

Once every patch of an image maps to one of $K$ codebook indices, the entire image becomes a grid of integers — a sequence, exactly like a sequence of word tokens. A second, separate model (commonly a transformer) can then be trained autoregressively over that grid: predict the next code index given the ones already generated, precisely the same objective language models use for the next word. VQ-VAE introduced the quantization mechanism itself; VQGAN paired it with an adversarial loss (borrowing from GANs) to sharpen the reconstructions the decoder produces from those discrete codes.

The tradeoff against a continuous latent space

Snapping to the nearest codebook entry throws away whatever information distinguished $z$ from $z_q$ — some reconstruction fidelity is lost at the quantization step no matter how well-trained the codebook is. What's gained in exchange is a representation that autoregressive, next-token-style generation can actually operate on, which a continuous latent space fundamentally cannot support in the same way.

Show Me the Code

Fitting a small codebook to data and measuring how reconstruction error shrinks as the codebook grows.

1import numpy as np2
3
4def quantize(data: np.ndarray, codebook: np.ndarray) -> np.ndarray:5    dists = ((data[:, None, :] - codebook[None, :, :]) ** 2).sum(axis=2)6    return codebook[dists.argmin(axis=1)]7
8
9rng = np.random.default_rng(1)10data = rng.normal(0, 1, size=(2000, 4))11for k in (4, 16, 64):12    codebook = rng.normal(0, 1, size=(k, 4))13    for _ in range(5):  # a few rounds of Lloyd's algorithm to fit the codebook14        idx = ((data[:, None, :] - codebook[None, :, :]) ** 2).sum(axis=2).argmin(axis=1)15        for c in range(k):16            if (idx == c).any():17                codebook[c] = data[idx == c].mean(axis=0)18    mse = float(((data - quantize(data, codebook)) ** 2).mean())19    print(f"codebook size {k}: reconstruction MSE = {mse:.4f}")20# -> codebook size 4: reconstruction MSE = 0.639421# -> codebook size 16: reconstruction MSE = 0.344722# -> codebook size 64: reconstruction MSE = 0.1692

Quadrupling the codebook from 4 to 16 entries roughly halves the reconstruction error, and quadrupling again to 64 nearly halves it again — a bigger palette snaps less aggressively, at the cost of a larger discrete vocabulary for the downstream autoregressive model to predict over.

Watch Out For

Codebook collapse — most entries never get used

If a small handful of codebook entries end up absorbing nearly every encoder output, the effective vocabulary shrinks far below the codebook's actual size, and the model loses representational capacity it was supposed to have. This is common enough to be an expected failure mode, not an edge case; monitoring codebook usage — how many entries are actually selected across a batch — during training catches it early, and techniques like resetting unused entries mid-training are standard mitigations.

Expecting the straight-through estimator's inconsistency to break training

Like in quantization-aware training, passing the gradient straight through a non-differentiable step feels like it should cause instability, and in practice it's a well-established, stable technique here too. Treat instability during training as a codebook-size, learning-rate, or initialization issue before suspecting the straight-through estimator itself.

The Quick Version

Vector quantization replaces an encoder's continuous output with the nearest entry in a fixed, learned codebook, producing a discrete latent representation.
The codebook lookup isn't differentiable, so a straight-through estimator lets gradients pass to the encoder as if quantization were the identity function.
A discrete latent grid can be modeled autoregressively by a second model, the same way a language model predicts the next word — something a continuous latent space can't directly support.
VQ-VAE introduced the mechanism; VQGAN added an adversarial loss for sharper reconstructions.
Codebook collapse, where most entries go unused, is a common failure worth monitoring directly.