Chapter 32 · Section 32.5 · 80 min

VQ-VAE and discrete codes

A VQ-VAE, short for vector-quantized variational autoencoder, is an autoencoder whose code is a short list of integers. An autoencoder is a pair of networks. The encoder compresses the input into a code, and the decoder rebuilds the input from that code (§32.1). In a VQ-VAE the encoder turns an image into a grid of vectors. Each vector is replaced by the closest entry of a learned table of vectors, called the codebook, and only that entry's number is kept. A $256\times256$ photograph can become a $32\times32$ grid of 1,024 integers, each one of 8,192 possible values.

That list of integers is why this section matters. A Transformer reads and writes tokens from a fixed vocabulary, and a grid of real-valued pixels has no vocabulary. Once an image is a sequence of code numbers, the next-token training that builds a language model can also generate images. The same idea turns speech into tokens for the audio models of §37.3. Most autoregressive image generators and many multimodal models start from a tokenizer of this kind, so interviewers for generative and multimodal teams expect you to know how one is trained.

Discrete codes also avoid a failure of the ordinary VAE. A VAE trains an autoencoder as a generative model (§32.3). Its loss adds a reconstruction error to a charge for how much information the code carries about the input. §32.4 showed that when the decoder is strong enough to model the data by itself, the cheapest solution can be a code that carries nothing. That failure is called posterior collapse. In a VQ-VAE the charge for the code is a constant fixed by the architecture, so the optimizer gains nothing by emptying the code.

The model comes from van den Oord, Vinyals & Kavukcuoglu (2017). We build it in seven steps. We find the nearest code with one matrix product. We read the model as a VAE to see what its code costs. We carry a gradient through a choice that has no derivative, and we train the codebook as a running k-means. We keep every code in use. We meet a second way through a discrete choice, which samples the code from a smoothed categorical distribution. Finally we fit a model over the integers, which turns the autoencoder into a generator.

Snap each vector to its nearest code

Start with the smallest case. The section comes back to it several times.

✏️ Quantizing one vector by hand

Take three codes in the plane, $\mathbf{c}_1 = (0, 0)$ , $\mathbf{c}_2 = (2, 0)$ , and $\mathbf{c}_3 = (0, 2)$ , and one encoder output $\mathbf{z}_e = (1.2, 0.3)$ .

The squared distance to $\mathbf{c}_1$ is $1.2^2 + 0.3^2 = 1.44 + 0.09 = 1.53$ .
The squared distance to $\mathbf{c}_2$ is $0.8^2 + 0.3^2 = 0.64 + 0.09 = 0.73$ .
The squared distance to $\mathbf{c}_3$ is $1.2^2 + 1.7^2 = 1.44 + 2.89 = 4.33$ .

The nearest code is number 2, so the quantized vector is $\mathbf{z}_q = (2, 0)$ and the stored code is the integer 2. The decoder sees $(2, 0)$ and never learns that the encoder said $(1.2, 0.3)$ . The only information that gets through is which of the three codes won, at most $\log_2 3 \approx 1.585$ bits.

Now picture every point of the plane. Each point has one nearest code. The set of points nearest to code $k$ is called its Voronoi cell, and the cells of all the codes tile the plane. Every encoder output in a cell is replaced by that cell's code. So the decoder learns only which cell the encoder output fell in, one of $K$ possibilities, and the information that passes is at most $\log_2 K$ bits per vector.

📐 Vector-quantized autoencoder

A vector-quantized autoencoder has four parts.

An encoder $\operatorname{Enc}_\phi$ with parameters $\phi$ that maps an input $\mathbf{x}$ to $T$ vectors $\mathbf{z}_{e,1}, \dots, \mathbf{z}_{e,T} \in \R^{d_z}$ . For an image the positions form a grid, $T = h\times w$ (for example $32\times32$ ). For audio they are a sequence of frames.
A codebook $\mathbf{C}_{\text{vq}} \in \R^{K\times d_z}$ whose rows $\mathbf{c}_1, \dots, \mathbf{c}_K$ are the code vectors. $K$ is the codebook size.
A quantizer that replaces each encoder output by its nearest code,

k^\star_t = \arg\min_{k\in[K]} \lVert \mathbf{z}_{e,t} - \mathbf{c}_k \rVert^2, \qquad \mathbf{z}_{q,t} = \mathbf{c}_{k^\star_t},

with ties broken toward the smaller index. Here $[K]$ is the set $\{1, \dots, K\}$ . 4. A decoder $\operatorname{Dec}_\theta$ with parameters $\theta$ that maps the quantized vectors $\mathbf{z}_{q,1}, \dots, \mathbf{z}_{q,T}$ to a reconstruction, or more precisely to the parameters of a likelihood $p_\theta(\mathbf{x}\mid\mathbf{z}_q)$ .

The discrete code of $\mathbf{x}$ is the index sequence $(k^\star_1, \dots, k^\star_T) \in [K]^T$ . Storing it takes $T\log_2 K$ bits.

In plain words, the encoder proposes a vector at each position, the quantizer rounds it to the nearest of $K$ stored prototypes, and the decoder works only from the prototypes. Finding $k^\star$ is a nearest-neighbor search, the problem of finding the stored vector closest to a query that vector databases solve in §7.4. Reading out $\mathbf{c}_{k^\star}$ is a row lookup by index, exactly what an embedding layer does (§1.2). §1.10 already used such a codebook to turn image patches and sound into tokens. This section shows how to train one inside a network.

Finding the nearest code with one matrix product

A training batch holds $n$ encoder outputs, every position of every image flattened into one list, stacked as the rows of $\mathbf{Z}\in\R^{n\times d_z}$ . A Python loop over the rows is far too slow. Instead we expand the squared distance, as in §1.3,

\lVert \mathbf{z} - \mathbf{c}_k\rVert^2 = \lVert\mathbf{z}\rVert^2 - 2\,\mathbf{z}^\top\mathbf{c}_k + \lVert\mathbf{c}_k\rVert^2 .

In plain words, a squared distance is the two squared lengths minus twice the dot product. For all pairs at once, the dot products form one matrix product $\mathbf{Z}\mathbf{C}_{\text{vq}}^\top \in \R^{n\times K}$ , the same way Exercise 0.1b computed all pairwise distances. The first term is the same for every code, so it cannot change which code is nearest and can be dropped. The last term cannot be dropped.

Taking the largest dot product $\mathbf{z}^\top\mathbf{c}_k$ alone is a classic bug. It favors codes with large norms, because a long code has a large dot product with every vector that points its way. It gives the right answer only when every code has the same norm. That is the normalized lookup we will meet later as a cure for unused codes. The next cell checks both claims.

Python · cell 1

The loop and the matrix product agree, and the matrix product is many times faster. The dot-product shortcut picks the nearest code for only about 1 percent of the vectors once code norms differ. At scale, memory becomes the problem. A million vectors against 16,384 codes in float32 is a distance matrix of $10^6\times16{,}384\times4$ bytes, about 65.5 GB. Production code therefore processes the batch in chunks and keeps only the argmin of each row. A chunk of 8,192 rows takes about 537 MB.

Core

Nearest-code assignment without a loop

Write vq_assign(Z, C). The input Z has shape (n, d) and holds encoder outputs, and C has shape (K, d) and holds the codebook. Return idx, an integer array of shape (n,) with the index of each row's nearest code in Euclidean distance, and Zq = C[idx]. Use the distance expansion and one matrix product, with no Python loop over rows. Break ties toward the smaller index, which np.argmin does for you. The test codebook has codes of very different norms, so the dot-product shortcut fails.

import numpy as np

def vq_assign(Z, C):
    """Z: (n, d) encoder outputs, C: (K, d) codebook -> idx (n,) int, Zq (n, d)."""
    raise NotImplementedError

Tests, hints, verified solutions, and AI help, with Pro at $7.99 a month.

VQ-VAE is a VAE whose posterior is a point

A VQ-VAE trains on the same objective as a VAE, for one particular choice of the encoder's distribution and of the prior. Recall that objective from §32.2. A latent variable model is trained by maximizing the evidence lower bound (ELBO), a lower bound on the log-likelihood $\log p_\theta(\mathbf{x})$ ,

\operatorname{ELBO}(\mathbf{x}) = \mathbb{E}_{q_\phi(\mathbf{z}\mid\mathbf{x})}\big[\log p_\theta(\mathbf{x}\mid\mathbf{z})\big] - D_{\mathrm{KL}}\big(q_\phi(\mathbf{z}\mid\mathbf{x})\,\Vert\,p(\mathbf{z})\big).

In plain words, the first term rewards reconstructing $\mathbf{x}$ from codes drawn from the encoder's distribution $q_\phi$ , called the approximate posterior. The second term charges for how far that distribution sits from the prior $p$ , a fixed distribution over codes. The KL charge is the cost of the information the code carries.

In a VQ-VAE the latent at one position is an index $k\in[K]$ . The encoder's posterior puts all its mass on the nearest code $k^\star$ , a one-hot distribution. During training the prior is uniform, $p(k) = 1/K$ . Here is the KL with real numbers before the general case.

Take $K = 4$ and suppose code 3 wins. Then $q = (0, 0, 1, 0)$ and $p = (\tfrac14, \tfrac14, \tfrac14, \tfrac14)$ . Only the winning term of the KL sum is nonzero, since $0\log 0 = 0$ , so the KL is $1\cdot\log\frac{1}{1/4} = \log 4 \approx 1.386$ nats, exactly 2 bits. Had code 1 won, the answer would be the same. The general case is the same computation,

D_{\mathrm{KL}}\big(q_\phi(\cdot\mid\mathbf{x})\,\Vert\,p\big) = \sum_{k=1}^K q_\phi(k\mid\mathbf{x})\log\frac{q_\phi(k\mid\mathbf{x})}{p(k)} = 1\cdot\log\frac{1}{1/K} = \log K .

With $T$ positions and independent uniform priors the KL is $T\log K$ . For a $32\times32$ grid with $K = 512$ that is $1{,}024\times9 = 9{,}216$ bits, the same for every image and every setting of the parameters. The reconstruction term simplifies too, since the expectation over a one-hot $q_\phi$ has only one outcome.

∂ Proof of the VQ-VAE bound

The marginal likelihood sums over all $K^T$ index sequences, $p_\theta(\mathbf{x}) = \sum_{k} p_\theta(\mathbf{x}\mid k)\,p_\omega(k)$ . Every term is nonnegative, so the sum is at least its single term at $k = k^\star(\mathbf{x})$ . Taking logs gives the bound. For the second statement, put the one-hot $q$ into $\mathbb{E}_q[\log p_\theta(\mathbf{x}\mid k)] - D_{\mathrm{KL}}(q\Vert p_\omega)$ . The expectation is $\log p_\theta(\mathbf{x}\mid k^\star)$ , and the KL is $-\log p_\omega(k^\star)$ because only one term survives. ∎

In plain words, the cost of describing $\mathbf{x}$ is the cost of its code under the prior plus the cost of the details the decoder still has to supply. The bound is tight when the decoder gives almost no probability to $\mathbf{x}$ from any other code. Van den Oord et al. argue that a converged decoder behaves this way, since it was only ever trained on $k^\star$ .

Three consequences follow, and each one returns later in the section.

There is nothing for the KL to collapse to. In §32.4 the optimizer could shrink the KL to zero by making the encoder's distribution equal the prior. The price was a code that says nothing. Here the KL is the constant $T\log K$ , and its gradient is zero, so the objective never rewards ignoring the code. It does not force the decoder to use the code either. But a decoder that ignores $\mathbf{z}_q$ pays more reconstruction loss and gains nothing. That is the sense in which VQ-VAE avoids posterior collapse.

The rate is fixed by the architecture, and usage decides how much of it is spent. In §32.4 the rate was the average KL term, the information the code is charged for. Here it is $\log K$ per position, whatever the parameters. The code $k^\star$ is a deterministic function of $\mathbf{x}$ , so the information it actually carries depends only on how evenly the codes are used. The next theorem makes this exact.

Learning the prior is a second stage of maximizing the same bound. With the encoder and decoder frozen, the only term that depends on $\omega$ is $\log p_\omega(k^\star(\mathbf{x}))$ . Maximizing the bound over $\omega$ is maximum-likelihood training of a model over index sequences. We do that in the last part of the section.

The next result uses two quantities from §0.5. The entropy $H(\mathbf{p}) = -\sum_k p_k\log p_k$ of a distribution is its average surprise, at most $\log K$ for $K$ outcomes. The mutual information $I(\mathbf{x}; k^\star)$ is how much knowing the code reduces uncertainty about the input.

∴ A deterministic code carries exactly the entropy of its usage

Let $\mathbf{x}\sim p_{\text{data}}$ , and let $\bar{\mathbf{p}} \in \R^K$ be the code frequencies at one position, $\bar p_k = \mathbb{P}(k^\star = k)$ . Then

I(\mathbf{x}; k^\star) = H(\bar{\mathbf{p}}) \le \log K ,

with equality exactly when every code is used equally often. §32.4 split the average KL term into two parts, the information the code carries about the input and the KL from the distribution of all codes to the prior. Here the average KL $\log K$ splits into this mutual information plus the KL from the usage to the uniform distribution, $D_{\mathrm{KL}}(\bar{\mathbf{p}}\,\Vert\,\text{uniform}) = \log K - H(\bar{\mathbf{p}})$ .

∂ Proof that a deterministic code carries its usage entropy

Write $I(\mathbf{x}; k^\star) = H(k^\star) - H(k^\star\mid\mathbf{x})$ . The second term is zero because $k^\star$ is a function of $\mathbf{x}$ , so knowing $\mathbf{x}$ leaves no uncertainty about it. The distribution of $k^\star$ is $\bar{\mathbf{p}}$ , so $H(k^\star) = H(\bar{\mathbf{p}})$ . Gibbs' inequality (§0.5), the fact that a KL divergence is never negative, gives $H(\bar{\mathbf{p}}) \le \log K$ with equality only at the uniform distribution. For the split, $D_{\mathrm{KL}}(\bar{\mathbf{p}}\Vert\text{uniform}) = \sum_k \bar p_k\log(K\bar p_k) = \log K - H(\bar{\mathbf{p}})$ , and the two pieces add to $\log K$ . ∎

In plain words, a VQ code can carry up to $\log_2 K$ bits per position, but it carries only as many as the entropy of its usage. A codebook of 16,384 entries in which only 200 are used, each equally often, carries $\log_2 200 \approx 7.6$ bits per position, not 14. That is why codebook usage gets so much attention below. Unused codes are the discrete form of posterior collapse, since both waste the capacity the code was given. An interviewer who asks why VQ-VAE does not suffer from posterior collapse usually follows up with what it suffers from instead, and this is the answer.

Gradients through a nearest-neighbor lookup

Training needs the gradient of the loss with respect to the encoder's parameters, and the path from the loss to the encoder runs through $\mathbf{z}_q = \mathbf{c}_{k^\star(\mathbf{z}_e)}$ . Go back to the three-code example. Move $\mathbf{z}_e$ from $(1.2, 0.3)$ to $(1.25, 0.3)$ and the nearest code is still $\mathbf{c}_2$ , so $\mathbf{z}_q$ does not change at all. Inside a Voronoi cell the derivative of $\mathbf{z}_q$ is zero, and on a cell boundary $\mathbf{z}_q$ jumps, so there is no derivative. The true gradient of the reconstruction loss with respect to the encoder is zero almost everywhere, and backpropagation would deliver nothing.

§0.3 met the same problem with rounding. Its answer was the straight-through estimator (Bengio, Léonard & Courville 2013), which runs the true operation in the forward pass and pretends it was the identity in the backward pass. §0.9 packaged it as a custom autograd function. VQ-VAE applies it to the whole vector.

In plain words, the decoder computes with the code, and the encoder is told what the decoder wanted the code to do. The decoder is fed $\mathbf{z}_q^{\text{st}}$ , so its forward pass is exactly the quantized one. Two side effects are easy to miss. First, the codebook receives no gradient from the reconstruction loss at all, because $\mathbf{z}_q$ appears only inside the stop-gradient. Second, the encoder's gradient is computed at the wrong place, at $\mathbf{z}_q$ rather than at $\mathbf{z}_e$ . The next theorem says how wrong that is.

It needs one term. A gradient is $M$ -Lipschitz if moving the input by a distance $r$ changes the gradient by at most $Mr$ , so $M$ measures how fast the gradient can turn.

∴ What the straight-through gradient approximates

Let $\mathcal{L}_{\text{rec}}(\mathbf{z})$ be the reconstruction loss as a function of the decoder's input, differentiable with an $M$ -Lipschitz gradient $\nabla\mathcal{L}_{\text{rec}}$ . Let $\boldsymbol\delta = \mathbf{z}_q - \mathbf{z}_e$ be the quantization error.

The straight-through gradient $\nabla\mathcal{L}_{\text{rec}}(\mathbf{z}_q)$ is the exact gradient of the surrogate $\mathbf{z}_e \mapsto \mathcal{L}_{\text{rec}}(\mathbf{z}_e + \boldsymbol\delta)$ , in which the quantization error is frozen as a constant.
It differs from the gradient of the unquantized autoencoder, $\nabla\mathcal{L}_{\text{rec}}(\mathbf{z}_e)$ , by at most $M\lVert\boldsymbol\delta\rVert$ .
It is not the gradient of the quantized loss $\mathcal{L}_{\text{rec}}(\mathbf{z}_q(\mathbf{z}_e))$ , which is zero inside every cell. In particular it carries no information about what would happen if $\mathbf{z}_e$ crossed into another cell.

∂ Proof of the straight-through approximation

(1) By the chain rule, $\nabla_{\mathbf{z}_e}\mathcal{L}_{\text{rec}}(\mathbf{z}_e + \boldsymbol\delta) = \nabla\mathcal{L}_{\text{rec}}(\mathbf{z}_e + \boldsymbol\delta) = \nabla\mathcal{L}_{\text{rec}}(\mathbf{z}_q)$ , since $\boldsymbol\delta$ is held constant. (2) The Lipschitz condition says $\lVert\nabla\mathcal{L}_{\text{rec}}(\mathbf{a}) - \nabla\mathcal{L}_{\text{rec}}(\mathbf{b})\rVert \le M\lVert\mathbf{a} - \mathbf{b}\rVert$ . Take $\mathbf{a} = \mathbf{z}_q$ and $\mathbf{b} = \mathbf{z}_e$ . (3) Inside a cell $\mathbf{z}_q(\mathbf{z}_e)$ is constant, so its derivative is zero. ∎

In plain words, the straight-through estimator pretends the encoder output already sits at its code. The pretense is good exactly when the encoder output is close to its code, and the error shrinks in proportion to the distance. Huh et al. (2023) call this difference the gradient gap.

A linear decoder does not remove the error. For $\mathcal{L}_{\text{rec}}(\mathbf{z}) = \lVert\mathbf{W}\mathbf{z} + \mathbf{b} - \mathbf{x}\rVert^2$ the gradient is $2\mathbf{W}^\top(\mathbf{W}\mathbf{z} + \mathbf{b} - \mathbf{x})$ , so the straight-through gradient is off by exactly $2\mathbf{W}^\top\mathbf{W}\boldsymbol\delta$ . The error vanishes when the quantization error is zero, or when the loss is linear in $\mathbf{z}$ , so that its gradient is the same at $\mathbf{z}_q$ and at $\mathbf{z}_e$ .

Python · cell 2

Once the spacing is small, the relative error falls in proportion to it, 0.237, 0.095, and 0.019 at spacings 0.25, 0.1, and 0.02, as part 2 of the theorem predicts. With coarse codes the straight-through gradient often points somewhere else entirely, with an average cosine of 0.43 at spacing 2. So encoder outputs should stay close to their codes. That is the job of one of two extra losses, which come next.

The codebook loss and the commitment loss

The straight-through estimator leaves two things untrained. The codebook gets no gradient from reconstruction, so something must pull each code toward the encoder outputs it represents. And nothing ties the encoder outputs to the codes. The VQ-VAE paper notes that the latent space has no natural scale, so encoder outputs can drift and grow if the codes do not keep up. The quantization error, which bounds the straight-through error, then grows with them. Each problem gets its own squared-distance term, and the stop-gradients decide which side of the distance each term moves.

📐 The VQ-VAE loss

For one input with positions $t = 1, \dots, T$ ,

\mathcal{L}_{\text{VQ-VAE}} = \underbrace{-\log p_\theta\big(\mathbf{x}\mid\mathbf{z}_q^{\text{st}}\big)}_{\text{reconstruction}} \;+\; \sum_{t=1}^T\underbrace{\big\lVert \operatorname{sg}[\mathbf{z}_{e,t}] - \mathbf{c}_{k^\star_t}\big\rVert^2}_{\text{codebook}} \;+\; \beta_{\text{commit}}\sum_{t=1}^T\underbrace{\big\lVert \mathbf{z}_{e,t} - \operatorname{sg}[\mathbf{c}_{k^\star_t}]\big\rVert^2}_{\text{commitment}} .

The decoder is trained by the first term only, the codebook by the second only, and the encoder by the first (through the straight-through estimator) and the third. The weight $\beta_{\text{commit}}$ is the paper's $\beta$ , renamed here so it is never confused with the KL weight $\beta$ of §32.4.

In plain words, the decoder learns to reconstruct from codes. Each code moves toward the encoder outputs it wins. Each encoder output is asked, more gently, to stay near the code it chose. Without the stop-gradients the last two terms would merge into one, $(1+\beta_{\text{commit}})\lVert\mathbf{z}_e - \mathbf{c}_{k^\star}\rVert^2$ , pulling codes and encoder outputs toward each other equally hard. Splitting it lets the codebook follow the data at full strength while the encoder, whose outputs carry the information, is only nudged.

The constant $T\log K$ from the bound is dropped, since it has no gradient. In practice the two VQ terms are averaged rather than summed. The paper averages them over the $T$ positions. DeepMind's Sonnet library takes a mean over every element, and the reconstruction term is often a mean squared error. The scaling matters, because $\beta_{\text{commit}}$ only means something relative to the size of the reconstruction loss, as the paper itself notes.

With the loss fixed, we can write down exactly what each parameter receives. This is the derivation an interviewer asks for, and the stop-gradients make it short.

∴ Gradients of the VQ-VAE loss

Take a batch of $n$ encoder outputs $\mathbf{z}_{e,i}$ with codes $k_i = k^\star_i$ . Let $\mathbf{g}_i = \partial\mathcal{L}_{\text{rec}}/\partial\mathbf{z}_{q,i}$ be the gradient the decoder sends back to its input at row $i$ . With the two VQ terms summed over the batch,

\frac{\partial\mathcal{L}}{\partial\mathbf{z}_{e,i}} = \mathbf{g}_i + 2\beta_{\text{commit}}\big(\mathbf{z}_{e,i} - \mathbf{c}_{k_i}\big), \qquad \frac{\partial\mathcal{L}}{\partial\mathbf{c}_k} = 2\sum_{i\,:\,k_i = k}\big(\mathbf{c}_k - \mathbf{z}_{e,i}\big) = 2\,n_k\big(\mathbf{c}_k - \bar{\mathbf{z}}_k\big),

where $n_k$ is the number of rows assigned to code $k$ and $\bar{\mathbf{z}}_k$ is their mean. A code that wins no rows gets exactly zero gradient.

∂ Proof of the VQ-VAE gradients

For $\mathbf{z}_{e,i}$ , the reconstruction term contributes $\mathbf{g}_i$ through the identity Jacobian of $\mathbf{z}_q^{\text{st}}$ . The codebook term contributes nothing, because $\mathbf{z}_{e,i}$ sits inside $\operatorname{sg}$ . The commitment term contributes $\beta_{\text{commit}}\cdot 2(\mathbf{z}_{e,i} - \mathbf{c}_{k_i})$ . For $\mathbf{c}_k$ , only the codebook terms of the rows with $k_i = k$ contain it outside a stop-gradient, and each contributes $2(\mathbf{c}_k - \mathbf{z}_{e,i})$ . Summing them gives $2n_k\mathbf{c}_k - 2\sum_i\mathbf{z}_{e,i} = 2n_k(\mathbf{c}_k - \bar{\mathbf{z}}_k)$ . ∎

The codebook gradient is a sum over every row that picked the code. That is a scatter-add, the backward pass of any row lookup, as §1.2 proved. The same NumPy trap applies. dC[idx] += rows keeps only one of several rows with the same index, while np.add.at(dC, idx, rows) adds them all.

✏️ One VQ step by hand

Continue the three-code example, $\mathbf{z}_e = (1.2, 0.3)$ , $k^\star = 2$ , $\mathbf{c}_2 = (2, 0)$ , with $\beta_{\text{commit}} = 0.25$ .

The codebook loss is $\lVert\mathbf{z}_e - \mathbf{c}_2\rVert^2 = 0.73$ , and the commitment loss is $0.25\times0.73 = 0.1825$ .
The code's gradient is $2(\mathbf{c}_2 - \mathbf{z}_e) = (1.6, -0.6)$ . With learning rate $\eta$ the code moves by $\eta\,(-1.6, 0.6)$ , toward $\mathbf{z}_e$ . With $\eta = 0.5$ it lands exactly on $\mathbf{z}_e$ , the mean of the one vector it won.
The encoder output's gradient is $\mathbf{g} + 2\cdot0.25\,(\mathbf{z}_e - \mathbf{c}_2) = \mathbf{g} + (-0.4, 0.15)$ . Its step, $-\eta\,\mathbf{g} + \eta\,(0.4, -0.15)$ , follows the decoder's wish and adds a small pull toward $\mathbf{c}_2$ .
Codes $\mathbf{c}_1$ and $\mathbf{c}_3$ receive zero gradient.

The next cell checks these numbers.

Python · cell 3

Interview

The VQ-VAE backward pass, with repeated codes

Write vq_vae_grads(z_e, C, g, beta_commit). The input z_e of shape (n, d) is a batch of encoder outputs, and C of shape (K, d) is the codebook. The input g = dL_rec/dz_q, also of shape (n, d), is the gradient the decoder sends back to its input. Assign each row to its nearest code. Then return (dz_e, dC), the gradients with respect to z_e and C of

\mathcal{L} = \mathcal{L}_{\text{rec}}\big(\mathbf{z}_q^{\text{st}}\big) + \sum_{i=1}^n\big\lVert\operatorname{sg}[\mathbf{z}_{e,i}] - \mathbf{c}_{k_i}\big\rVert^2 + \beta_{\text{commit}}\sum_{i=1}^n\big\lVert\mathbf{z}_{e,i} - \operatorname{sg}[\mathbf{c}_{k_i}]\big\rVert^2 ,

using the straight-through estimator. The tests compare your answer with finite differences of the function that each parameter actually sees, on a batch in which three vectors share one code. This is the coding version of the most common VQ interview question.

import numpy as np

def vq_vae_grads(z_e, C, g, beta_commit):
    """z_e: (n, d), C: (K, d), g: (n, d) = dL_rec/dz_q.  Returns (dz_e (n, d), dC (K, d))."""
    raise NotImplementedError

Tests, hints, verified solutions, and AI help, with Pro at $7.99 a month.

Carrying the gradient across with a rotation

The straight-through estimator moves the gradient from $\mathbf{z}_q$ to $\mathbf{z}_e$ without changing it. The decoder sees only $\mathbf{z}_q$ , so for a given loss every encoder output in a cell gets the same gradient, wherever in the cell it sits. Fifty et al. (2024) proposed a different rule, the rotation trick. It carries the gradient across with the rotation that turns $\mathbf{z}_e$ onto the direction of $\mathbf{z}_q$ , so the gradient keeps its angle to the vector. This is a 2024 refinement. A first reading can skim to the widget below and come back for the details.

Here is a tiny case. Let $\mathbf{z}_e = (1, 0)$ and its code $\mathbf{z}_q = (0, 2)$ , a quarter turn away and twice as long. Suppose the decoder's gradient is $\mathbf{g} = (0, 1)$ , pointing along $\mathbf{z}_q$ , so a descent step would shrink $\mathbf{z}_q$ . The straight-through estimator hands $\mathbf{z}_e$ the same $(0, 1)$ , which pushes it sideways. The rotation trick turns $\mathbf{g}$ back by the quarter turn and scales it by the length ratio 2, giving $(2, 0)$ . That points along $\mathbf{z}_e$ , so the step shrinks $\mathbf{z}_e$ , which is what the decoder asked of $\mathbf{z}_q$ .

📐 The rotation trick

Write $\hat{\mathbf{z}}_e = \mathbf{z}_e/\lVert\mathbf{z}_e\rVert$ and $\hat{\mathbf{z}}_q = \mathbf{z}_q/\lVert\mathbf{z}_q\rVert$ for the unit vectors, and $s = \lVert\mathbf{z}_q\rVert/\lVert\mathbf{z}_e\rVert$ for the ratio of norms. Let $\mathbf{R}\in\R^{d_z\times d_z}$ be the rotation in the plane of $\mathbf{z}_e$ and $\mathbf{z}_q$ that takes $\hat{\mathbf{z}}_e$ to $\hat{\mathbf{z}}_q$ and fixes every vector orthogonal to that plane. With $\mathbf{r} = (\hat{\mathbf{z}}_e + \hat{\mathbf{z}}_q)/\lVert\hat{\mathbf{z}}_e + \hat{\mathbf{z}}_q\rVert$ ,

\mathbf{R} = \mathbf{I} - 2\,\mathbf{r}\mathbf{r}^\top + 2\,\hat{\mathbf{z}}_q\hat{\mathbf{z}}_e^\top .

The rotation trick feeds the decoder $\tilde{\mathbf{z}}_q = \operatorname{sg}[s\mathbf{R}]\,\mathbf{z}_e$ . Its value is $s\mathbf{R}\mathbf{z}_e = s\lVert\mathbf{z}_e\rVert\hat{\mathbf{z}}_q = \mathbf{z}_q$ , so the forward pass is unchanged. Its backward pass sends $\partial\mathcal{L}/\partial\mathbf{z}_e = s\,\mathbf{R}^\top\mathbf{g}$ to the encoder.

In the tiny case, $\mathbf{r} = (1, 1)/\sqrt2$ , and the formula gives $\mathbf{R} = \begin{pmatrix}0 & -1\\ 1 & 0\end{pmatrix}$ , the quarter turn. Then $s\mathbf{R}^\top\mathbf{g} = 2\,(1, 0) = (2, 0)$ , as claimed.

∂ Why this R is the rotation, as two Householder reflections

A Householder reflection $\mathbf{H}_{\mathbf{u}} = \mathbf{I} - 2\mathbf{u}\mathbf{u}^\top$ for a unit vector $\mathbf{u}$ mirrors space across the hyperplane orthogonal to $\mathbf{u}$ . It is symmetric and orthogonal, with determinant $-1$ . Write $c = \hat{\mathbf{z}}_e^\top\hat{\mathbf{z}}_q$ , so $\lVert\hat{\mathbf{z}}_e + \hat{\mathbf{z}}_q\rVert^2 = 2(1 + c)$ and $\mathbf{r}^\top\hat{\mathbf{z}}_e = \mathbf{r}^\top\hat{\mathbf{z}}_q = (1+c)/\sqrt{2(1+c)} = \sqrt{(1+c)/2}$ .

First reflect across the hyperplane orthogonal to $\mathbf{r}$ . Since $2\mathbf{r}(\mathbf{r}^\top\hat{\mathbf{z}}_e) = (\hat{\mathbf{z}}_e + \hat{\mathbf{z}}_q)(1+c)/(1+c) = \hat{\mathbf{z}}_e + \hat{\mathbf{z}}_q$ , we get $\mathbf{H}_{\mathbf{r}}\hat{\mathbf{z}}_e = -\hat{\mathbf{z}}_q$ . Then reflect across the hyperplane orthogonal to $\hat{\mathbf{z}}_q$ , which sends $-\hat{\mathbf{z}}_q$ to $\hat{\mathbf{z}}_q$ . The product $\mathbf{H}_{\hat{\mathbf{z}}_q}\mathbf{H}_{\mathbf{r}}$ is orthogonal with determinant $(-1)^2 = 1$ , so it is a rotation, and it maps $\hat{\mathbf{z}}_e$ to $\hat{\mathbf{z}}_q$ . A vector orthogonal to both $\hat{\mathbf{z}}_e$ and $\hat{\mathbf{z}}_q$ is orthogonal to $\mathbf{r}$ as well, so both reflections leave it alone.

Expanding the product,

\mathbf{H}_{\hat{\mathbf{z}}_q}\mathbf{H}_{\mathbf{r}} = \mathbf{I} - 2\mathbf{r}\mathbf{r}^\top - 2\hat{\mathbf{z}}_q\hat{\mathbf{z}}_q^\top + 4(\hat{\mathbf{z}}_q^\top\mathbf{r})\,\hat{\mathbf{z}}_q\mathbf{r}^\top ,

and the last term is $4\sqrt{(1+c)/2}\;\hat{\mathbf{z}}_q(\hat{\mathbf{z}}_e + \hat{\mathbf{z}}_q)^\top/\sqrt{2(1+c)} = 2\hat{\mathbf{z}}_q\hat{\mathbf{z}}_e^\top + 2\hat{\mathbf{z}}_q\hat{\mathbf{z}}_q^\top$ . Two terms cancel, leaving $\mathbf{R} = \mathbf{I} - 2\mathbf{r}\mathbf{r}^\top + 2\hat{\mathbf{z}}_q\hat{\mathbf{z}}_e^\top$ . Applying it, or its transpose, to a vector costs two dot products and needs no $d_z\times d_z$ matrix. The formula fails only when $\hat{\mathbf{z}}_q = -\hat{\mathbf{z}}_e$ . There $\mathbf{r}$ is undefined, and any rotation by $\pi$ in a plane containing $\hat{\mathbf{z}}_e$ will do.

The construction says how to compute the rotation. The reason to want it is one property that the straight-through estimator lacks.

∂ Proof that the angle is preserved

$\mathbf{R}$ is orthogonal, so $\lVert s\mathbf{R}^\top\mathbf{g}\rVert = s\lVert\mathbf{g}\rVert$ . And $(s\mathbf{R}^\top\mathbf{g})^\top\mathbf{z}_e = s\,\mathbf{g}^\top\mathbf{R}\mathbf{z}_e = s\lVert\mathbf{z}_e\rVert\,\mathbf{g}^\top\hat{\mathbf{z}}_q$ . Dividing by the two norms, $s\lVert\mathbf{g}\rVert$ and $\lVert\mathbf{z}_e\rVert$ , gives $\cos(s\mathbf{R}^\top\mathbf{g}, \mathbf{z}_e) = \mathbf{g}^\top\hat{\mathbf{z}}_q/\lVert\mathbf{g}\rVert = \cos(\mathbf{g}, \mathbf{z}_q)$ . ∎

In plain words, if the decoder wants $\mathbf{z}_q$ to grow, the rotation trick asks $\mathbf{z}_e$ to grow along its own direction. The straight-through estimator would instead ask it to move in a fixed direction, which may point partly sideways for a vector near the edge of the cell. When $\mathbf{z}_e$ and $\mathbf{z}_q$ point the same way, $\mathbf{R} = \mathbf{I}$ , and the trick reduces to the straight-through estimator scaled by $s$ .

🔬 What the rotation trick changes in practice

Under the straight-through estimator, points of a cell that receive the same decoder gradient all move by the same step, so their distances to each other never change. Fifty et al. show that the rotation trick pushes the points of a cell apart or pulls them toward the code, depending on the direction of the gradient. Pushing points across a boundary can bring more codes into use. Their headline result trains the VQGAN autoencoder used by latent diffusion (Rombach et al. 2022) on ImageNet, with a $32\times32$ grid and 16,384 codes. The rotation trick lowered reconstruction FID from 5.0 to 1.1, raised codebook usage (measured per batch) from 2 percent to 27 percent, and cut the quantization error fiftyfold, from 0.012 to 0.00024 (their Table 3).

Reconstruction FID compares real images with their reconstructions. It is the Fréchet Inception distance, a distance between the feature statistics of the two sets in a pretrained image network, and lower is better (§34.6 derives it). These are empirical findings for particular models. The paper also notes a failure mode. The angle between $\mathbf{z}_e$ and $\mathbf{z}_q$ can be obtuse when codes or encoder outputs are forced close to the origin. Then the rotation over-turns the gradient and can do worse than the straight-through estimator.

In the widget the decoder's loss is a tilted bowl with its lowest point at the target square. The dashed arrow is the true downhill step at the encoder output, the step an unquantized autoencoder would take.

Drag $\mathbf{z}_e$ toward the edge of its cell. The straight-through step stays the same everywhere in the cell, and it parts ways with the dashed step. Switch to the rotation trick. The step now turns through the angle between the two lines from the origin, so it keeps its angle to the vector. That is a different goal from matching the dashed step, and as you drag, the step lands nearer to it or farther from it. Then go back to straight-through and switch to the whole plane. Every point of a cell gets the same arrow, the checkerboard pattern that Fifty et al. draw, and the rotation trick breaks that pattern.

Research

The rotation trick, forward matrix and batched backward

Implement the rotation trick of Fifty et al. (2024) in two forms.

rotation_matrix(e, q) returns the $d\times d$ matrix $s\mathbf{R}$ , where $s = \lVert\mathbf{q}\rVert/\lVert\mathbf{e}\rVert$ and $\mathbf{R}$ is the rotation that takes $\hat{\mathbf{e}}$ to $\hat{\mathbf{q}}$ in their plane and fixes every vector orthogonal to both, so that $s\mathbf{R}\mathbf{e} = \mathbf{q}$ .
rotation_trick_backward(z_e, z_q, g) takes three arrays of shape (n, d) and returns the encoder gradients $s\mathbf{R}^\top\mathbf{g}$ , row by row, in $O(nd)$ work and without forming any $d\times d$ matrix.

Both must stay finite when $\mathbf{q}$ points exactly opposite to $\mathbf{e}$ . In that case any rotation by $\pi$ in a plane containing $\hat{\mathbf{e}}$ is acceptable. The tests check the mapping, that $\mathbf{R}$ is a rotation (orthogonal with determinant 1) that moves only the plane of $\mathbf{e}$ and $\mathbf{q}$ , the preserved angle, and the batched form against the matrix and for speed.

import numpy as np

def rotation_matrix(e, q):
    """The rotation-trick Jacobian s R for one pair, with s R e = q and R a rotation (d >= 2)."""
    return np.eye(len(e))                                 # placeholder: the straight-through Jacobian

def rotation_trick_backward(z_e, z_q, g):
    """Rows of z_e, z_q, g are (n, d). Returns s R^T g for every row without forming R."""

Tests, hints, verified solutions, and AI help, with Pro at $7.99 a month.

Learning the codebook as online k-means

Look again at the codebook gradient, $2n_k(\mathbf{c}_k - \bar{\mathbf{z}}_k)$ . It points from the mean of the code's vectors to the code. Hold the assignments fixed, and the codebook loss $\sum_{i: k_i = k}\lVert\mathbf{z}_{e,i} - \mathbf{c}_k\rVert^2$ is a sum of squared distances, which is smallest when the code sits at the mean, $\mathbf{c}_k = \bar{\mathbf{z}}_k$ .

Assigning every vector to its nearest code and then moving each code to the mean of its vectors is one round of Lloyd's algorithm for k-means. §7.4 used the same algorithm to build its search indexes. So VQ-VAE learns its codebook by k-means, with two differences. The data, the encoder outputs, move while it runs. And they arrive one minibatch at a time.

A plain gradient step with learning rate $\eta$ gives $\mathbf{c}_k \leftarrow (1 - 2\eta n_k)\mathbf{c}_k + 2\eta n_k\bar{\mathbf{z}}_k$ . That is a step toward the mean, but its size is proportional to the count $n_k$ . With $\eta = 0.01$ , a code that won 4 vectors moves 8 percent of the way to their mean, and a code that won 40 moves 80 percent. Popular codes jump to their data while rarely chosen codes crawl, and a code with $n_k = 0$ does not move at all. Appendix A.1 of the VQ-VAE paper proposes an update whose step is normalized by the count.

📐 EMA codebook update

Keep, for every code, an exponential moving average (EMA) $N_k$ of its count and an EMA $\mathbf{m}_k$ of the sum of the vectors assigned to it. After assigning a batch, with $n_k$ the batch count of code $k$ , update

N_k \leftarrow \gamma N_k + (1-\gamma)\,n_k, \qquad \mathbf{m}_k \leftarrow \gamma\,\mathbf{m}_k + (1-\gamma)\sum_{i\,:\,k_i = k}\mathbf{z}_{e,i}, \qquad \mathbf{c}_k = \frac{\mathbf{m}_k}{N_k},

with decay $\gamma \in [0, 1)$ , 0.99 in the paper. The codebook loss is then dropped from $\mathcal{L}_{\text{VQ-VAE}}$ , and the optimizer never touches the codebook.

DeepMind's Sonnet implementation divides by Laplace-smoothed counts $\tilde N_k = \frac{N_k + \epsilon}{n + K\epsilon}\,n$ , where $n = \sum_j N_j$ and $\epsilon = 10^{-5}$ . The smoothing keeps the total count at $n$ and makes the division finite for a code whose count has decayed to zero.

Here is one step with numbers, first for a rarely chosen code and then for a popular one.

✏️ One EMA step by hand

Take $\gamma = 0.9$ and a code at $\mathbf{c} = (0, 0)$ with count $N = 4$ and sum $\mathbf{m} = (0, 0)$ . A batch assigns it 4 vectors whose mean is $(1, 0)$ , so their sum is $(4, 0)$ .

The count becomes $0.9\times4 + 0.1\times4 = 4.0$ .
The sum becomes $0.9\times(0, 0) + 0.1\times(4, 0) = (0.4, 0)$ .
The code becomes $(0.4, 0)/4.0 = (0.1, 0)$ , 10 percent of the way to the batch mean.

Now repeat with a popular code that has $N = 40$ and wins 40 vectors with the same mean. Its count stays 40, its sum becomes $(4, 0)$ , and it also lands at $(0.1, 0)$ . Both codes move the fraction $1 - \gamma$ of the way, however many vectors they won. The gradient step above moved them 8 percent and 80 percent.

That equal treatment is no accident. Over many batches the EMA computes exactly the k-means mean of the vectors each code has won, with recent batches weighted more.

∴ The EMA codebook update is online k-means

Suppose the EMA starts from $N_k^{(0)} > 0$ and $\mathbf{m}_k^{(0)} = N_k^{(0)}\mathbf{c}_k^{(0)}$ , and let batch $s$ give code $k$ the count $n_k^{(s)}$ . After $t$ batches, without smoothing, $\mathbf{c}_k^{(t)}$ is the unique minimizer of the exponentially weighted codebook loss

J_k^{(t)}(\mathbf{c}) = \gamma^t N_k^{(0)}\big\lVert\mathbf{c} - \mathbf{c}_k^{(0)}\big\rVert^2 + (1-\gamma)\sum_{s=1}^{t}\gamma^{t-s}\sum_{i\in\text{batch } s,\; k_i = k}\big\lVert\mathbf{z}_{e,i} - \mathbf{c}\big\rVert^2 .

With $\gamma = 0$ one step is exactly Lloyd's update on the current batch. Each update moves the code the fraction $w = (1-\gamma)n_k/N_k^{\text{new}}$ of the way from its old position to the batch mean. In steady state $w$ is about $1-\gamma$ for every code, however popular it is.

∂ Proof that the EMA update is online k-means

Unrolling the recursions gives $N_k^{(t)} = \gamma^tN_k^{(0)} + (1-\gamma)\sum_s\gamma^{t-s}n_k^{(s)}$ and $\mathbf{m}_k^{(t)} = \gamma^tN_k^{(0)}\mathbf{c}_k^{(0)} + (1-\gamma)\sum_s\gamma^{t-s}\sum_{i}\mathbf{z}_{e,i}$ . The loss $J_k^{(t)}$ is a weighted sum of squared distances from $\mathbf{c}$ to a set of points. The old code has weight $\gamma^tN_k^{(0)}$ , and each assigned vector of batch $s$ has weight $(1-\gamma)\gamma^{t-s}$ . Setting its gradient to zero shows that its minimizer is the weighted mean of those points. The numerator of that mean is $\mathbf{m}_k^{(t)}$ and the denominator is $N_k^{(t)}$ .

For the last statement, substitute $\mathbf{m}_k^{\text{old}} = N_k^{\text{old}}\mathbf{c}_k^{\text{old}}$ into one update. This gives $\mathbf{c}_k^{\text{new}} = (1-w)\,\mathbf{c}_k^{\text{old}} + w\,\bar{\mathbf{z}}_k$ with $w = (1-\gamma)n_k/(\gamma N_k^{\text{old}} + (1-\gamma)n_k)$ . In steady state $N_k \approx n_k$ , so $w \approx 1-\gamma$ . ∎

In plain words, each code sits at the mean of every vector it has won over roughly the last $1/(1-\gamma)$ batches, 100 batches for $\gamma = 0.99$ , with recent batches weighted more. That is what k-means would compute if the encoder stopped moving, and the forgetting is what lets the codebook follow an encoder that keeps moving. A code that stops winning keeps its position, because its sum and its count decay together. Practitioners believe the EMA update is more stable than the codebook loss, and the simulation below agrees. An interviewer will expect you to derive it from the k-means mean.

🔬 Smoothing, debiasing, and who used the EMA update

With the Laplace smoothing, an idle code does drift eventually. Once its count falls to the order of $\epsilon$ , the smoothed count stops shrinking while the sum keeps decaying, so the code slides toward the origin. For $\gamma = 0.99$ that takes on the order of a thousand idle steps. The VQ-VAE paper describes the EMA update in its appendix but says it was not used in its own experiments, while VQ-VAE-2 used it throughout with $\gamma = 0.99$ . Libraries such as Sonnet also correct the moving averages for their start at zero, the same bias correction Adam applies to its moment estimates (§0.6).

The EMA rule is closer to gradient descent than it looks. Łańcucki et al. (2020) show that for a code that wins one vector, with its count held at 1, it is exactly a gradient step on the codebook loss with learning rate $(1-\gamma)/2$ . The same comparison works in general. A gradient step with learning rate $\eta_k$ moves code $k$ the fraction $2\eta_k n_k$ of the way to its batch mean, and the EMA moves it the fraction $(1-\gamma)n_k/N_k^{\text{new}}$ . So one EMA step is a gradient step with the code's own learning rate $\eta_k = (1-\gamma)/(2N_k^{\text{new}})$ , smaller for popular codes. It is still a step on the codebook loss, so a code that wins nothing gets no update. The EMA does not by itself rescue dead codes.

Core

The EMA codebook update with Laplace smoothing

Write ema_codebook_update(N, m, Z, idx, gamma, eps). The state is the EMA counts N of shape (K,) and the EMA sums m of shape (K, d). The batch is Z of shape (n, d), with its code indices idx of shape (n,). Return (N_new, m_new, C_new), where

$N_k \leftarrow \gamma N_k + (1-\gamma)n_k$ and $\mathbf{m}_k \leftarrow \gamma\mathbf{m}_k + (1-\gamma)\sum_{i: k_i = k}\mathbf{z}_i$ , and
the codebook is $\mathbf{c}_k = \mathbf{m}_k/\tilde N_k$ with the smoothed counts $\tilde N_k = (N_k + \epsilon)\,n/(n + K\epsilon)$ , where $n = \sum_j N_j$ after the update.

Return the unsmoothed N_new, since the smoothing is only for the division, as in Sonnet. One code in the tests has never been used, so its count is zero.

import numpy as np

def ema_codebook_update(N, m, Z, idx, gamma=0.99, eps=1e-5):
    """N: (K,) EMA counts, m: (K, d) EMA sums, Z: (n, d) encoder outputs, idx: (n,) their codes.
    Returns (N_new, m_new, C_new). N_new is stored unsmoothed; only the division is smoothed."""
    N_new = gamma * N
    m_new = gamma * m
    return N_new, m_new, m_new / N_new[:, None]

Tests, hints, verified solutions, and AI help, with Pro at $7.99 a month.

Codebook collapse and how to prevent it

Everything so far trains only the codes that win. A code with no assigned vectors gets exactly zero gradient, and the EMA leaves it where it is. So a code that starts far from the encoder outputs, or that the encoder outputs drift away from, is never chosen and never updated. Meanwhile the live codes move toward the data and their cells grow, covering the space where the dead codes might have won. This is codebook collapse, also called index collapse. By the usage theorem above, it directly lowers the information the code can carry.

📐 Codebook usage and perplexity

Let $\bar{\mathbf{p}}\in\R^K$ be the empirical frequencies of the codes over a large sample of encoder outputs. The usage is the fraction of codes with $\bar p_k > 0$ . The codebook perplexity is

\operatorname{PPL}_{\text{code}} = e^{H(\bar{\mathbf{p}})} = \exp\Big(-\sum_{k\,:\,\bar p_k > 0}\bar p_k\log\bar p_k\Big),

with the entropy in nats. It is built the same way as the perplexity of a language model in §0.5, so $1\le\operatorname{PPL}_{\text{code}}\le K$ . It is the number of equally used codes that would carry the same information, and $\log\operatorname{PPL}_{\text{code}}$ is the mutual information of the usage theorem, in nats.

Here are both numbers for a small codebook.

Measure both over many vectors, not over one batch. A batch of $n$ vectors can show at most $n$ distinct codes. With $K = 16{,}384$ and a batch of 4,096 positions, the per-batch usage can never exceed 25 percent, even for a perfect codebook. LlamaGen, an image generator built on the Llama language-model architecture (Sun et al. 2024), counts usage over a queue of the last 65,536 assignments.

Python · cell 4

Python · cell 5

Data on eight blobs is gentle. In a 256-dimensional latent space with 16,384 codes, the chance that a randomly placed code lands near any encoder output is far smaller. That is why collapse gets worse as $K$ and $d_z$ grow, as the measurements below show.

Pick a starting position for the codes, an update rule, and whether to restart dead codes, then press Train. Dead codes are drawn hollow. Start far away with no restarts and watch most codes stay hollow. Then turn on restarts, or start the codes on the data.

Each fix attacks one of three causes. The codes may start in a bad place, codes that lose may never get a gradient, or the geometry may leave most codes far from the data.

Start the codes on the data. Initialize the codebook by running k-means on a sample of encoder outputs. Łańcucki et al. (2020) went further. During early training they periodically re-estimated the whole codebook on recent encoder outputs with k-means++, a version of k-means whose starting centers are spread apart. They also added batch normalization before the quantizer and gave the codebook its own, larger learning rate.
Restart dead codes. Whenever a code's average usage falls below a threshold, move it to a random encoder output from the current batch, as the Jukebox music model did (Dhariwal et al. 2020). It costs nothing at inference and guarantees every code gets data.
Look up codes in a small, normalized space. Project each encoder output to a low-dimensional lookup space, and $\ell_2$ -normalize both the projection and the codes. The squared distance then depends only on the cosine similarity. In a few dimensions codes and data overlap far more easily. Yu et al. (2021) introduced this in ViT-VQGAN, projecting from 768 dimensions down to 32 or 8.
Update the whole codebook at once. Write the codebook as $\mathbf{C}_{\text{vq}}\mathbf{W}$ , with a frozen random $\mathbf{C}_{\text{vq}}$ and a learned shared matrix $\mathbf{W}\in\R^{d_z\times d_z}$ . Then every gradient moves every code (Zhu et al. 2024, SimVQ). A related method freezes 100,000 codes, the k-means centers of image features from a pretrained CLIP encoder, and trains only a linear projection of them. It reports usage above 99 percent (Zhu et al. 2024, VQGAN-LC). Huh et al. (2023) propose a related shared affine reparameterization of the codes.
Fix the gradient. The rotation trick above changes how the encoder's gradient crosses the quantizer, and can raise usage.
Give up the learned codebook. Finite scalar quantization and lookup-free quantization, the subject of §32.6, round each channel to a fixed grid, so there is no codebook to collapse.

Interview

Codebook health, and dead-code restarts

Write two utilities that a VQ training loop needs.

codebook_stats(idx, K) returns the usage fraction (codes used at least once, divided by K) and the perplexity $e^{H(\bar{\mathbf{p}})}$ , with the entropy in nats and $0\log0 = 0$ .
restart_dead_codes(C, counts, Z, rng, threshold=1.0) returns a copy of C in which every code whose count is below threshold is replaced by a randomly chosen row of Z, the current encoder outputs. Draw the rows with the generator rng.

Live codes must not change, C must not be modified in place, and no two restarted codes may receive the same row. You may assume Z has at least as many rows as there are dead codes.

import numpy as np

def codebook_stats(idx, K):
    """idx: (n,) code indices in [0, K). Returns (usage fraction, perplexity e^H in nats)."""
    raise NotImplementedError

def restart_dead_codes(C, counts, Z, rng, threshold=1.0):
    """Replace every code with count below threshold by a distinct random row of Z."""

Tests, hints, verified solutions, and AI help, with Pro at $7.99 a month.

Discrete codes without the straight-through estimator

The straight-through estimator is not the only way through a discrete choice. The other way keeps a real probability for each of the $K$ codes and samples the code from a smoothed version of that distribution, which has a gradient. DALL-E, OpenAI's first text-to-image model, tokenized its images this way, with a discrete VAE (dVAE) (Ramesh et al. 2021).

Let $\boldsymbol\ell\in\R^K$ be the encoder's logits at one position, scores whose softmax gives the probability of each code. The smoothing rests on the Gumbel-max trick of §0.4. Draw independent standard Gumbel noise $g_k = -\log(-\log U_k)$ , with each $U_k$ uniform on $(0, 1)$ . Here $g_k$ is noise, not a gradient. Then $\arg\max_k(\ell_k + g_k)$ is an exact sample from $\operatorname{softmax}(\boldsymbol\ell)$ . The argmax has no gradient, so the relaxation replaces it by a softmax with a temperature. This is the Gumbel-softmax, found independently by Jang, Gu & Poole (2017) and, as the Concrete distribution, by Maddison, Mnih & Teh (2017).

✏️ One Gumbel-softmax sample by hand

Take three categories with probabilities $(0.6, 0.3, 0.1)$ , so the logits are $\boldsymbol\ell = (\ln 0.6, \ln 0.3, \ln 0.1) = (-0.511, -1.204, -2.303)$ . Suppose the noise comes out as $\mathbf{g} = (0.1, 1.5, -0.3)$ . Then $\boldsymbol\ell + \mathbf{g} = (-0.411, 0.296, -2.603)$ , and the argmax picks category 2. Over many draws of the noise, the argmax picks category 2 with probability exactly 0.3.

At temperature $\tau = 1$ the softmax of $\boldsymbol\ell + \mathbf{g}$ is $(0.319, 0.646, 0.036)$ , a soft vote for category 2.
At $\tau = 0.1$ the softmax of $(\boldsymbol\ell + \mathbf{g})/0.1$ is $(0.001, 0.999, 0.000)$ , almost the one-hot vector of category 2.

Both are differentiable in $\boldsymbol\ell$ for this fixed noise, so a gradient can flow back to the logits.

The general recipe is the same for any logits and any temperature.

📐 Gumbel-softmax relaxation

For logits $\boldsymbol\ell\in\R^K$ , temperature $\tau > 0$ , and independent standard Gumbel noise $g_1, \dots, g_K$ , the Gumbel-softmax sample is

\mathbf{y} = \operatorname{softmax}\!\Big(\frac{\boldsymbol\ell + \mathbf{g}}{\tau}\Big), \qquad y_k = \frac{\exp\big((\ell_k + g_k)/\tau\big)}{\sum_{j=1}^K\exp\big((\ell_j + g_j)/\tau\big)} .

It lies in the probability simplex, the set of vectors with nonnegative entries that sum to one, and it is a differentiable function of $\boldsymbol\ell$ for fixed noise. Its argmax has exactly the distribution $\operatorname{softmax}(\boldsymbol\ell)$ at every $\tau$ . As $\tau\to0$ , $\mathbf{y}$ tends to the one-hot vector of that argmax, and as $\tau\to\infty$ to the uniform vector. The straight-through Gumbel variant uses the one-hot $\mathbf{y}_{\text{hard}}$ in the forward pass and the gradient of $\mathbf{y}$ in the backward pass, $\mathbf{y}_{\text{hard}} + \mathbf{y} - \operatorname{sg}[\mathbf{y}]$ .

In plain words, add the right noise to the logits, then take a soft argmax instead of a hard one. The noise is drawn independently of $\boldsymbol\ell$ , so this is the reparameterization trick of §32.3, applied to a relaxed categorical variable. The temperature trades bias against variance. A high $\tau$ gives smooth samples, unlike the one-hot codes the model will see at test time, so the gradient is biased but steady. A low $\tau$ gives nearly one-hot samples and a nearly unbiased gradient with high variance. The cell shows the argmax frequencies staying exact while the samples sharpen.

Python · cell 6

At $\tau = 5$ the average sample is close to uniform, even though its argmax is distributed exactly as $(0.6, 0.3, 0.1)$ . At $\tau = 0.1$ the samples are almost one-hot, with a largest entry of 0.963 on average, and their mean is the softmax itself.

The dVAE puts this relaxation inside the ELBO. At each of $32\times32$ positions the encoder outputs $K = 8{,}192$ logits. The decoder's first layer is a $1\times1$ convolution, a matrix applied at every position, whose columns are the code vectors. It turns the relaxed one-hot $\mathbf{y}$ into $\sum_k y_k\mathbf{c}_k$ , a weighted average of code vectors. The KL to the uniform prior is now a genuine, input-dependent term, $D_{\mathrm{KL}}(q_\phi(\cdot\mid\mathbf{x})\Vert\text{uniform}) = \log K - H(q_\phi(\cdot\mid\mathbf{x}))$ per position.

Ramesh et al. annealed $\tau$ from 1 to 1/16 over the first 150,000 updates with a cosine schedule, and report that a linear schedule usually diverged. They raised the KL weight from 0 to 6.6 over the first 5,000 updates. To tokenize images for the Transformer, they took the argmax of the encoder's logits with no noise. The 1,024 tokens of a $256\times256$ image shrank the Transformer's context by a factor of $196{,}608/1{,}024 = 192$ compared with raw pixel values.

📌 DALL-E's KL weight

§32.4 showed that a KL weight $\beta$ above 1 usually buys a simpler code at the price of worse reconstructions. DALL-E's dVAE went against that trade-off. Raising the weight to $\beta = 6.6$ , far above the $\beta = 1$ at which the objective is the ELBO, gave better codebook usage and a smaller final reconstruction error (Ramesh et al. 2021, section 2.1). The KL to a uniform prior is $\log K$ minus the entropy of each position's posterior, so a large weight rewards spreading probability over many codes early in training. The authors speculate that with a small weight, noise from the relaxation pushes the optimizer to use fewer codes early on. One detail matters when you compare numbers. The paper divides the whole loss by the $256\times256\times3$ pixel values, so in its normalized loss the KL weight appears as $\beta/192$ .

The dVAE's codes found a second use. BEiT pretrained Vision Transformers by masking image patches and predicting their DALL-E dVAE tokens (Bao et al. 2021), as §28.5 describes.

Core

Gumbel-softmax samples, soft and hard

Write gumbel_softmax(logits, tau, rng, hard=False) for logits of shape (n, K). Draw independent standard Gumbel noise $g = -\log(-\log U)$ with $U$ uniform from rng.random. Return one relaxed sample per row, $\operatorname{softmax}((\boldsymbol\ell + \mathbf{g})/\tau)$ , computed stably. With hard=True, return instead the one-hot vector of each row's argmax, which is the forward value of the straight-through Gumbel estimator. The tests check that the argmax frequencies match $\operatorname{softmax}(\boldsymbol\ell)$ at two temperatures. That catches two common bugs, dividing only the logits by $\tau$ and a wrong sign in the noise.

import numpy as np

def gumbel_softmax(logits, tau, rng, hard=False):
    """logits: (n, K). One relaxed sample per row: softmax((logits + g) / tau), g standard Gumbel.
    With hard=True return the one-hot of its argmax instead."""
    raise NotImplementedError

Tests, hints, verified solutions, and AI help, with Pro at $7.99 a month.

A prior over codes turns the autoencoder into a generator

A trained VQ-VAE can encode and decode, but it cannot yet generate. During training the prior was uniform and independent across positions. Sampling $T$ independent uniform indices gives garbage, because real index grids are highly structured. A sky code is usually next to another sky code.

The VQ-VAE bound already says what to do. With the encoder and decoder frozen, maximize $\log p_\omega(k^\star(\mathbf{x}))$ over the training set. Put the $T$ positions in an order, raster order for an image (left to right, then top to bottom), and factor the prior by the chain rule,

p_\omega(k_1, \dots, k_T) = \prod_{t=1}^T p_\omega(k_t\mid k_1, \dots, k_{t-1}).

In plain words, the probability of a whole grid is the probability of each code given the codes before it. Train any autoregressive model with next-token cross-entropy over the $K$ classes, exactly as a GPT is trained on text (§6.2). To generate, sample one index at a time with the methods of §6.4, then run the decoder once.

The prior is also where the bits go. Its average cross-entropy on held-out code grids, in bits, is the length an ideal compressor would need to store a grid using that prior. It replaces the $T\log_2 K$ of the uniform prior in the bound.

Here is the whole pipeline at toy scale. The images are $8\times8$ bars, with every row equal or every column equal. The tokenizer cuts each image into a $4\times4$ grid of $2\times2$ patches. It quantizes the patches with a codebook of $K = 6$ , learned by ten rounds of Lloyd's algorithm, each one an assignment with assign_vec from earlier followed by an average. That is exactly a VQ-VAE whose encoder and decoder are the identity on patches. The prior predicts each token from all earlier ones with a separate softmax regression per position, a tiny stand-in for a Transformer.

Python · cell 7

Python · cell 8

Python · cell 9

At real scale the original VQ-VAE used a PixelCNN prior for images. A PixelCNN is a convolutional network that predicts each position from the ones above and to the left (van den Oord et al. 2016). For audio it used a WaveNet, the same kind of network for a sequence of sound samples (van den Oord et al. 2016). Since 2021 the prior is almost always a Transformer, in DALL-E, VQGAN, Parti, LlamaGen, and others, as §34.4 describes. The early-fusion multimodal models of §31.6 put image tokens and text tokens in one sequence, and they read and write exactly these image tokens. Neural audio codecs turn speech into the same kind of tokens (§37.3).

How much does a grid of codes actually store? The numbers below are worth being able to compute on a whiteboard.

✏️ Bits per image at real scale

VQ-VAE on ImageNet at $128\times128$ used a $32\times32$ grid with $K = 512$ . That is 9 bits per code and $1{,}024\times9 = 9{,}216$ bits per image. Raw 8-bit RGB takes $128\times128\times3\times8 = 393{,}216$ bits, so the code is about 42.7 times smaller (the paper truncates it to 42.6, section 4.2).
DALL-E's dVAE uses a $32\times32$ grid with $K = 8{,}192$ . That is 13 bits per code and $13{,}312$ bits per image, against $256\times256\times3\times8 = 1{,}572{,}864$ bits, about 118 times fewer.
A codebook of 16,384 codes at 97 percent usage carries at most $\log_2(0.97\times16{,}384) \approx 13.96$ bits per code, and less if usage is uneven. The prior's cross-entropy, not $\log_2 K$ , is what a sampled image actually costs.

Bits are not the whole story. A code learned without any labels can also line up with units that people name.

Where this breaks

Four limits push the field beyond the plain VQ-VAE.

Biased gradients. The straight-through error is bounded by a multiple of the quantization error, so a coarse codebook trains its encoder with the wrong gradient.
Collapse at scale. Codebooks collapse more readily as $K$ and $d_z$ grow, and the fixes above are patches on that dynamic.
The size of the codebook is a trade-off with the prior. More codes give more bits per token and better reconstructions. But the prior faces a larger softmax and sees each code less often. In LlamaGen's ablation, reconstruction stopped improving past 16,384 codes. Reconstruction FID was 3.02, 2.91, 2.19, and 2.26 for 4,096, 8,192, 16,384, and 32,768 codes (Table 2).
Blur. Squared-error reconstruction from a few codes is blurry, because the best squared-error guess is the average of all plausible images (§32.1).

The next section adds perceptual and adversarial losses and replaces the learned codebook with quantizers that cannot collapse.

In PyTorch

A production quantizer keeps the codebook in buffers updated by the EMA rule, uses the one-line straight-through estimator z_e + (z_q - z_e).detach(), tracks usage, and restarts dead codes. The block trains a tiny linear encoder and decoder through it on 2-D blobs and checks that the encoder receives a gradient. It then shows PyTorch's built-in straight-through Gumbel-softmax.

🔥 In PyTorchfor a machine with PyTorch

The table compares the four ways of training a discrete code that this chapter covers.

	VQ with codebook loss	VQ with EMA	Gumbel-softmax (dVAE)	FSQ, LFQ (see §32.6)
encoder's gradient through the quantizer	straight-through (or rotation trick)	straight-through (or rotation trick)	reparameterized relaxation, biased for $\tau > 0$	straight-through rounding
how the codebook learns	gradient of $\lVert\operatorname{sg}[\mathbf{z}_e] - \mathbf{c}\rVert^2$ , step size proportional to usage	running means, online k-means	gradient of the ELBO through the soft lookup	fixed grid, nothing to learn
extra losses	codebook and commitment	commitment	KL to the prior, weight $\beta$	none (LFQ adds an entropy penalty)
codebook usage	prone to collapse	better, still needs restarts or good initialization	good with a large KL weight	near full by construction
used by	VQ-VAE, VQGAN	VQ-VAE-2, Jukebox, SoundStream	DALL-E, BEiT's tokenizer	MAGVIT-v2 (LFQ), MaskGIT and UViM with FSQ

⚠️ Where VQ implementations go wrong

The most common bugs are quiet ones.

Taking the argmax of $\mathbf{z}^\top\mathbf{c}$ instead of the argmin of the distance silently favors long codes.
Writing dC[idx] += ... drops the gradient of every repeated index.
Forgetting a stop-gradient turns the two VQ terms into one symmetric pull.
Forgetting the straight-through line leaves the encoder with no reconstruction gradient at all. It then learns only from the commitment loss and collapses toward the codes.
Using EMA updates and also passing the codebook to the optimizer updates it twice.
Measuring perplexity on a single small batch understates usage, since a batch of $n$ vectors can show at most $n$ codes.
Mixing bases, an entropy in bits with $e^H$ , gives a wrong perplexity.
Restarting several dead codes at the same vector makes them tie forever.
A $\beta_{\text{commit}}$ tuned with summed losses means something different once the losses are averaged over dimensions.

🧭 What to remember

A VQ-VAE replaces each encoder vector by its nearest codebook entry and keeps only the index, so an image becomes $T$ tokens of $\log_2 K$ bits each. Nearest-code search is one matrix product through the distance expansion, and the code-norm term cannot be dropped.
Read as a VAE, the posterior is one-hot and the training prior uniform, so the KL is the constant $T\log K$ . Nothing rewards ignoring the code, which is why VQ-VAE avoids posterior collapse. The information actually carried is the usage entropy, at most $\log K$ per position.
The argmin has zero gradient almost everywhere. The straight-through estimator copies the decoder's gradient from $\mathbf{z}_q$ to $\mathbf{z}_e$ , with an error bounded by the quantization error, and the rotation trick carries it across with a rotation instead.
The loss adds a codebook term (codes move to their encoder outputs) and a commitment term with $\beta_{\text{commit}} = 0.25$ (encoder outputs stay near their codes). The codebook gradient is a scatter-add, $2n_k(\mathbf{c}_k - \bar{\mathbf{z}}_k)$ , and is zero for unused codes.
The EMA update is online k-means, each code the weighted mean of what it won over about $1/(1-\gamma)$ batches. Codebook collapse comes from codes that never win, and k-means initialization, dead-code restarts, a low-dimensional normalized lookup, and shared reparameterizations fix it.
Gumbel-softmax is the other route through a discrete choice, used by DALL-E's dVAE. A prior trained by next-token cross-entropy over the indices turns the autoencoder into a generator. Next, §32.6 adds perceptual and adversarial losses and quantizers without a learned codebook, to build the image tokenizers that modern generators use.

van den Oord, Vinyals & Kavukcuoglu (2017). Neural Discrete Representation Learning. The VQ-VAE paper. Section 3 gives the one-hot posterior, the constant KL, the straight-through estimator, and the three-term loss, and appendix A.1 the EMA update.
Razavi, van den Oord & Vinyals (2019). Generating Diverse High-Fidelity Images with VQ-VAE-2. Two levels of codes, EMA codebooks, and autoregressive priors with self-attention at $256\times256$ and beyond.
Bengio, Léonard & Courville (2013). Estimating or Propagating Gradients Through Stochastic Neurons for Conditional Computation. Where the straight-through estimator was named and compared with alternatives.
Jang, Gu & Poole (2017). Categorical Reparameterization with Gumbel-Softmax, and Maddison, Mnih & Teh (2017), The Concrete Distribution. The relaxation, its temperature, and the straight-through variant, discovered independently.
Ramesh et al. (2021). Zero-Shot Text-to-Image Generation. Section 2.1 and appendix A describe the dVAE, its temperature and KL schedules, and the surprising effect of a large KL weight.
Łańcucki et al. (2020). Robust Training of Vector Quantized Bottleneck Models, and Huh, Cheung, Agrawal & Isola (2023), Straightening Out the Straight-Through Estimator. Why codebooks collapse and how to stop it, from data-dependent re-initialization to the gradient gap and affine reparameterization.
Fifty et al. (2024). Restructuring Vector Quantization with the Rotation Trick. The rotation trick, its Householder construction, and its effect on usage and reconstruction.
Sun et al. (2024). Autoregressive Model Beats Diffusion: Llama for Scalable Image Generation, and Zhu et al. (2024), Scaling the Codebook Size of VQGAN to 100,000 with a Utilization Rate of 99%. Codebook dimension, size, and usage measured at ImageNet scale.

Made for larger screens

VQ-VAE and discrete codes

Snap each vector to its nearest code

Finding the nearest code with one matrix product

Nearest-code assignment without a loop

VQ-VAE is a VAE whose posterior is a point

Gradients through a nearest-neighbor lookup

The codebook loss and the commitment loss

The VQ-VAE backward pass, with repeated codes

Carrying the gradient across with a rotation

The rotation trick, forward matrix and batched backward

Learning the codebook as online k-means

The EMA codebook update with Laplace smoothing

Codebook collapse and how to prevent it

Codebook health, and dead-code restarts

Discrete codes without the straight-through estimator

Gumbel-softmax samples, soft and hard

A prior over codes turns the autoencoder into a generator

Where this breaks

In PyTorch

Read more

Watch next