VQ-VAE: learn a discrete dictionary

An interactive field guide

What if an encoder had to choose from a finite dictionary? VQ-VAE turns continuous features into discrete indices, then learns how those indices occur together.

Before this article: Variational autoencoders

What was already possible

A Gaussian VAE uses continuous samples and a differentiable reparameterization. Discrete nearest-neighbor selection breaks that gradient path. Vector quantization was already a compression tool; the new challenge is to train the encoder, dictionary, and decoder together.

What changes, and why

The encoder emits z_e at each location. Select the closest dictionary vector e_k and send that vector to the decoder. Store k as the discrete code. Use a straight-through gradient estimator for the encoder, a dictionary loss for e_k, and a commitment loss to keep z_e near its selected code.

Read from left to right. Select a module to connect its role, tensor shape, and calculation. Snippets are teaching pseudocode. Sources provide the paper method or audited implementation.

1. Encodez_e: [B,h,w,D]2. Select and route gradientsindices: [B,h,w]; e: [K,D]3. Decode and fit a priorimage / categorical probabilities

Encode

Each spatial vector is quantized separately with a shared dictionary.

z_e = encoder(x)

Select and route gradients

Forward uses a dictionary vector; backward uses the straight-through estimator.

ids = squared_distances(z_e, e).argmin(-1)
z_q = e[ids]
z_st = z_e + stopgrad(z_q - z_e)

Decode and fit a prior

Reconstruction trains the tokenizer. A separate prior predicts indices.

x_hat = decoder(z_st)
L = recon(x_hat,x) + norm(stopgrad(z_e)-z_q)**2
L += beta * norm(z_e-stopgrad(z_q))**2

Inspect the pinned source

$$k=\arg\min_j\|z_e-e_j\|_2^2,\qquad z_q=e_k$$

Work through the mechanism

In the demo, z_e=0.8 and the four entries are −3, −1, 1, 3. The selected vector is 1 and squared error is 0.04. Moving z_e across a boundary changes the index abruptly. This is why ordinary differentiation through argmin cannot train the encoder.

Change a parameter, inspect the computed values, then compare this small example with the model described above.

Analytical teaching model · controls change the example, not a trained research model.

Encoder value z

0.8

Codebook entries K

4

Reset

Orange is the encoder output. Green is its selected code.

Computed values for the current example

Quantity	Value
Selected index	2
z_q	1.000
(z − z_q)²	0.040

Training is not sampling

The forward expression z_e + stopgrad(z_q−z_e) equals z_q. Its derivative with respect to z_e is one. The codebook term stops gradients to the encoder; the commitment term stops gradients to the codebook. This is a biased gradient estimator, not the exact derivative of nearest-neighbor lookup.

During tokenizer training, the deterministic posterior and uniform categorical prior give a constant KL of log K per location. Afterward, a PixelCNN prior is fitted to the codes. Sampling arbitrary independent uniform IDs is not equivalent to sampling this learned prior.

What this does not establish

A dictionary can contain unused entries. More entries improve potential capacity but can make prior modeling harder. Original-paper gradient updates and the EMA alternative should not be silently mixed. VQ-VAE-2 next changes the hierarchy, not the basic nearest-code idea.

Open the derivation

The hard posterior assigns probability one to one code. Under a uniform K-entry prior, its KL is log K, independent of the encoder parameters except at code switches. This does not learn the image-wide code distribution.

$$D_{KL}(q(k|x)\Vert 1/K)=\log K$$

Use sg for stop-gradient. The middle term moves the codebook toward the encoder. The last term moves the encoder toward the codebook. Straight-through routing handles the reconstruction gradient separately.

$$L=L_{\mathrm{rec}}+\|\operatorname{sg}[z_e]-e_k\|^2+\beta\|z_e-\operatorname{sg}[e_k]\|^2$$

Check your understanding

Which loss directly moves the selected dictionary vector?

The dictionary loss. The reconstruction straight-through path updates the encoder and decoder; commitment anchors the encoder.

Sources and next steps

• Method / pinned implementation source
• VQ-VAE-2
• VQGAN
• UniTok

Return to the learning map

Diagrams and snippets are teaching reconstructions. Paper reports, implementation observations, and analytical examples are distinguished. No research model runs in this page.