Interactive Explainer

Vector Quantization (VQ-VAE)

A plain autoencoder has a continuous latent — any point in space is legal. VQ-VAE instead keeps a small learned codebook of $K$ vectors and forces every encoding to snap to its nearest entry. The latent becomes discrete: an index into a lookup table. Drag the encoding point below and watch it jump to the closest code; drag the codes and watch their territories — the Voronoi cells — redraw.

~8 min Deep Learning · Generative Models · Discrete Latents

Prelude

One continuous vector becomes one integer

The encoder produces a continuous vector $z \in \mathbb{R}^d$. Vector quantization replaces it with the nearest of $K$ learned code vectors $e_1, \dots, e_K$ (the codebook):

k^{\star} = \arg\min_{k}\; \lVert z - e_k \rVert, \qquad z_q = e_{k^{\star}}.

The decoder then works from $z_q$ alone, so all it ever receives is one of $K$ possible vectors — equivalently, the single integer $k^{\star}$. Nearest-neighbour assignment carves the plane into Voronoi cells: every point in a cell shares the same closest code. Quantization is just: which cell did $z$ land in?

The Lab

Snap to the nearest code

The coloured regions are the Voronoi cells of the codebook; the large ringed dots are the code vectors $e_k$. The dark diamond is the encoding $z$ — drag it anywhere and the arrow shows it snapping to $z_q$, the code owning whatever cell it is in. Drag the codes to move the cell boundaries. Hit Quantize a cloud to drop a batch of encodings and colour each by the code it is assigned to.

Codebook size $K$ = 5

Drag the dark diamond ($z$) or any ringed code ($e_k$). The arrow points from $z$ to its quantized value $z_q$ — the nearest code. Toggle the cloud to see a whole batch assigned at once.

Assigned code $k^{\star}$ —

Quantization error $\lVert z - z_q\rVert$ —

Bits to store index $\log_2 K$ —

Drag $z$ slowly across a boundary. The assigned code and the arrow flip the instant you cross into a new cell — the map from continuous $z$ to discrete $k^{\star}$ is piecewise constant, with jumps exactly at the Voronoi seams. That discreteness is the whole point, and also the reason gradients need a special trick (below).

The Payoff

How do you train through a hard argmin?

The snap is a non-differentiable step: nudging $z$ within a cell does not change $z_q$ at all, so $\partial z_q / \partial z = 0$ almost everywhere. That would block gradients from ever reaching the encoder. VQ-VAE uses the straight-through estimator: on the backward pass, pretend the quantizer was the identity and copy the decoder's gradient straight from $z_q$ onto $z$.

z_q = z + \operatorname{stopgrad}\!\big(e_{k^{\star}} - z\big) \;\;\Longrightarrow\;\; \frac{\partial z_q}{\partial z} \approx 1 .

Forward: the decoder sees the true discrete code $e_{k^{\star}}$.
Backward: the gradient flows as if $z_q = z$, so the encoder still learns.
Codebook loss: extra terms $\lVert \operatorname{sg}[z] - e_k \rVert^2$ and a commitment term $\beta\lVert z - \operatorname{sg}[e_k]\rVert^2$ pull the code to its assigned encodings and keep the encoder from drifting away.

Why discretize at all? A discrete latent grid is a sequence of integers — exactly what an autoregressive prior (a PixelCNN or a Transformer) can model. VQ-VAE and its successor VQ-GAN turn images into token grids so that the same machinery used for language can generate pictures. The codebook is the vocabulary.