VQ-VAE

The VQ-VAE (vector-quantised variational autoencoder) is a generative model with a discrete latent code, introduced by van den Oord, Vinyals and Kavukcuoglu at NIPS 2017. An autoencoder’s encoder maps an image to a grid of vectors; each vector is replaced by its nearest entry in a learned codebook of $K$ vectors, and a decoder reconstructs the image from the chosen entries. The image is thereby stored as a small grid of integers, one codebook index per position. New images come from a second model, an autoregressive prior fitted to those index grids after the autoencoder is trained. VQ-VAE-2 (Razavi, van den Oord & Vinyals, NeurIPS 2019) stacks two or three such grids at different scales and scales up the prior. For images, both use convolutional encoder–decoders trained with a squared-error reconstruction loss; they are the starting point of the discrete image tokenizers that VQGAN and later two-stage generators build on.

Two pieces of the model are borrowed from elsewhere. The codebook is learned by a form of k-means clustering, which signal processing calls vector quantization, and the model is read as a variational autoencoder whose posterior is a deterministic one-hot choice. The VAE article summarizes the model in a paragraph; this article covers the quantizer, its gradient, the codebook updates, the prior, VQ-VAE-2 and the failure in which most of the codebook goes unused.

Problem Addressed

In 2017 continuous latents trained with the Gaussian reparameterization trick were the norm for VAEs. Discrete latents were trained with sampling-based estimators such as NVIL and VIMCO, or with continuous relaxations (Concrete, Gumbel-softmax); van den Oord et al. (§2) note that none of these had closed the gap to continuous VAEs and that most were tested on small datasets such as MNIST with latents of fewer than about eight dimensions. A second problem was posterior collapse: a VAE with a powerful autoregressive decoder tends to ignore its latent (see the variational autoencoder article). The authors wanted discrete latents, which they argue fit many modalities naturally (language is discrete, speech is usually written as a sequence of symbols, and images can often be described concisely in language), trained with low-variance gradients and used by the decoder.

Their answer was to make the posterior deterministic. A one-hot posterior with a uniform prior makes the ELBO’s KL term a constant, so it needs no estimating and, in our reading, no longer pushes the latent toward being ignored. The cost is that the nearest-neighbour step has no useful gradient, which the model handles with a straight-through estimator.

Architecture

The model has an encoder, a quantizer with its codebook, and a decoder; the prior over codes is trained afterwards (see Training).

Encoder and Codebook

The encoder maps the image $x$ to $z_e(x)$, a grid of $D$-dimensional vectors. The codebook, which the paper calls the latent embedding space, holds $K$ vectors $e_1, \dots, e_K \in \mathbb{R}^D$. Each grid vector is replaced by its nearest codebook entry (§3.1, Eq. 1–2):

$$z_q(x) = e_k, \qquad k = \arg\min_j \lVert z_e(x) - e_j \rVert_2 ,$$

applied independently at every position. For images the latents form a 2D grid, for speech a 1D sequence and for video a 3D volume (§3.1). On ImageNet at $128 \times 128 \times 3$, the grid was $32 \times 32$ with $K = 512$; storing 9-bit indices instead of 8-bit colour values reduces the size by a factor of about 42.6 (§4.2). On CIFAR-10 the latent was $8 \times 8 \times 10$ (§3.2).

Read as a VAE, the posterior is categorical and one-hot,

$$q(z = k \mid x) = \begin{cases} 1 & k = \arg\min_j \lVert z_e(x) - e_j \rVert_2 \\ 0 & \text{otherwise,} \end{cases}$$

and with a uniform prior over the $K$ values the KL term of the ELBO is the constant $\log K$ (§3.1). The quantizer is the nearest-centre rule of k-means: the codebook vectors are the centres, and their Voronoi cells partition the encoder’s output space.

Decoder

The decoder maps the grid of chosen vectors back to an image. The CIFAR-10 models used an encoder of two stride-2 $4 \times 4$ convolutions and two residual $3 \times 3$ blocks with 256 hidden units, and a decoder of two residual blocks and two stride-2 transposed convolutions (§4.1). For ImageNet the decoder was purely deconvolutional; for $84 \times 84$ DeepMind Lab frames (a $21 \times 21$ latent) the reconstructions “looked nearly identical to their originals” (§4.2). For audio the decoder was a WaveNet-like dilated convolutional network, conditioned on a speaker identity for voice conversion (§4.3).

The decoder may also be autoregressive. The authors trained a second VQ-VAE with a PixelCNN decoder on the $21 \times 21$ codes of the first, using only three latents with $K = 512$ each (27 bits) for the whole frame. Such a set-up, they write, typically breaks VAEs through posterior collapse; theirs used its latents, reconstructing room layout and nearby walls while the PixelCNN generated textures (§4.2, Fig. 5).

Training

Training has two stages. Stage one trains the encoder, codebook and decoder with the prior held uniform; stage two fits an autoregressive prior to the codes of the training set (§3.3).

The Straight-Through Gradient

The quantizer is piecewise constant: moving $z_e$ slightly leaves $z_q$ unchanged, so its derivative is zero almost everywhere and undefined on cell boundaries. Van den Oord et al. approximate the gradient “similar to” the straight-through estimator of Bengio, Léonard and Courville (2013) and copy the gradient arriving at the decoder input $z_q$ unchanged to the encoder output $z_e$ (§3.2):

$$\frac{\partial \mathcal{L}}{\partial z_e} \;:=\; \frac{\partial \mathcal{L}}{\partial z_q} .$$

Because $z_e$ and $z_q$ live in the same $D$-dimensional space, this gradient says how the encoder should move its output to lower the reconstruction loss, and a large enough move changes which codebook entry is chosen in the next forward pass (Fig. 1). They mention the subgradient through the quantizer as an alternative and report that the simple estimator worked well. In code the estimator is one line: the decoder receives $z_e + \mathrm{sg}[z_q - z_e]$, where $\mathrm{sg}$ is the stop-gradient, which equals $z_q$ in the forward pass and has the identity Jacobian with respect to $z_e$ in the backward pass; Sonnet’s implementation does exactly this.

The copied gradient is the gradient of the decoder’s loss evaluated at $z_q$, not at $z_e$, and not the gradient of the quantized loss (which is zero). Huh et al. (2023, Eq. 14) call the difference between the gradient at $z_e$ and at $z_e + \epsilon$, with $\epsilon = z_q - z_e$ the quantization error, the gradient gap: it vanishes when $\epsilon = 0$, and for a Lipschitz-continuous decoder it is bounded in proportion to the quantization error (§3.2). The example below measures all three gradients.

Codebook and Commitment Losses

The straight-through path gives the codebook no gradient from the reconstruction loss. The codebook is learned with a vector-quantization term that pulls each chosen entry toward the encoder outputs assigned to it, and a commitment term keeps the encoder’s outputs close to their entries. The paper (§3.2, Eq. 3) writes the objective for one latent as

$$\mathcal{L} = \log p(x \mid z_q(x)) + \lVert \mathrm{sg}[z_e(x)] - e \rVert_2^2 + \beta \lVert z_e(x) - \mathrm{sg}[e] \rVert_2^2 ,$$

where $e$ is the chosen entry. The first term is printed as a log-likelihood although the objective is minimized; it is to be read as the negative log-likelihood, the reconstruction loss. The decoder is trained by the first term only, the encoder by the first and third, and the codebook by the second (§3.2). With a grid of latents, the second and third terms are averaged over positions. The commitment loss exists because, the authors write, the embedding space is dimensionless and can grow arbitrarily if the codebook does not train as fast as the encoder. They found results unchanged for $\beta$ from 0.1 to 2.0 and used $\beta = 0.25$ throughout, noting that the right value depends on the scale of the reconstruction loss. VQ-VAE-2 writes the same objective with the squared error $\lVert x - D(e) \rVert_2^2$ as the reconstruction term (Razavi et al., §2.1, Eq. 2).

For one codebook entry the gradient of the second term is proportional to $e_i$ minus the mean of the encoder outputs assigned to it, so gradient descent moves each entry toward the mean of its cell, the update step of online k-means (our derivation). The k-means article covers that sequential update.

Exponential Moving-Average Updates

The NeurIPS paper notes that the dictionary could instead be updated with moving averages of $z_e(x)$, and that this was not used for its experiments (§3.2). VQ-VAE-2 uses moving averages in place of the codebook loss, “as proposed in” the VQ-VAE paper, which it cites in its arXiv version, with (Razavi et al., §2.1)

$$N_i^{(t)} = \gamma N_i^{(t-1)} + (1 - \gamma)\, n_i^{(t)}, \qquad m_i^{(t)} = \gamma m_i^{(t-1)} + (1 - \gamma) \sum_{j=1}^{n_i^{(t)}} E(x)_{i,j}^{(t)}, \qquad e_i^{(t)} = \frac{m_i^{(t)}}{N_i^{(t)}} ,$$

where $n_i^{(t)}$ counts the encoder outputs in the mini-batch assigned to entry $e_i$, $E(x)_{i,j}^{(t)}$ are those outputs, and $\gamma = 0.99$ in all their experiments. Each entry tracks a running mean of its cell, as in mini-batch k-means but with exponential forgetting. That arXiv version could not be read for this article, so the details here come from VQ-VAE-2 and the code.

Sonnet’s VectorQuantizerEMA (v2 branch, as of November 2026) adds two details. Its moving averages are debiased: the hidden averages start at zero and are divided by $1 - \gamma^t$. And the cluster sizes are Laplace-smoothed, $N_i \leftarrow (N_i + \epsilon)/(n + K\epsilon) \cdot n$ with $n = \sum_i N_i$ and $\epsilon = 10^{-5}$, so that no division is by zero. The module’s docstring gives the advantage that the codebook update no longer depends on the optimizer used for the rest of the network, and says that in most experiments the EMA version trains faster. Huh et al. (§2, Appendix A.2) show, for a single code vector, that an EMA step equals an SGD step on the codebook term with a learning rate equal to the EMA’s weight on the new value, $1 - \gamma$ in the notation above. With EMA the loss keeps only the reconstruction and commitment terms.

The Autoregressive Prior

While the autoencoder trains, the prior over codes is uniform. Afterwards the training images are encoded, and an autoregressive model is fitted to their index grids (§3.3): a PixelCNN for images, with only spatial masking because the latent has a single channel (§4.2), and a WaveNet for audio. The prior factorizes the grid in raster order,

$$p(z) = \prod_{i} p(z_i \mid z_1, \dots, z_{i-1}),$$

each factor a $K$-way softmax. Sampling draws indices one position at a time, looks up their codebook vectors and decodes them in one pass. Training the prior jointly with the autoencoder is left as future work (§3.3).

Fitting the prior after training is also a statement about the likelihood. The model’s likelihood is $\log p(x) = \log \sum_k p(x \mid z_k)\, p(z_k)$, and the paper bounds it from below by the single term $\log p(x \mid z_q(x))\, p(z_q(x))$, citing Jensen’s inequality (the bound also holds simply because the sum contains that term). Since the decoder is trained only on $z_q(x)$, the authors argue that once converged it puts no mass on other codes, so the bound is close to tight (§3.2); it follows that a better prior $p(z)$ directly improves the likelihood. Razavi et al. (§3.2) describe the fitted prior as lossless re-encoding of the latents with a better model of their distribution, which brings the bit rate closer to its entropy and makes sampled codes look like those the decoder saw in training.

Codebook Collapse

Training a vector quantizer often leaves most of the codebook unused, a failure called codebook collapse, or “index collapse” by Huh et al. (2023), in which only a small fraction of codes is used (§1). An entry that no encoder output selects receives no gradient from the codebook or commitment terms and no EMA update (the debiased variant in the example below sends the codes unused in its first update to the origin), so it stays where it is, while the encoder’s outputs drift away from the initial codebook. Huh et al. trace collapse to this drift: their commitment loss, which weights the codebook and commitment terms together, seen as a divergence averaged over the encoder outputs, trains only the selected codes, which learn a mode-seeking behaviour, and once a code is unselected it is likely to stay unselected (§3.1). On ResNet-18 trained on ImageNet-100 they saw more than 95% of codes go unselected after a few optimization steps (§3.1). Huh et al. tie the two problems together: when the encoder’s outputs drift away from the codebook, assignments become sub-optimal, the quantization error rises and with it the gradient gap (§3.2).

Two measures track collapse: the number of dead codes, used by no latent in a held-out set, and the codebook perplexity, $\exp(H(p))$ for the frequencies $p$ of the codes, which equals $K$ when every code is used equally and 1 when one code takes everything. Sonnet’s quantizers return this perplexity; Huh et al. (§5.1) write it as $2^{H(p)}$ with entropy in bits, the same number, and caution that a high perplexity also shows redundancy when the task needs fewer codes.

Remedies, in outline, as Huh et al. survey and compare them (§3, Table 1):

• EMA codebook updates (above).
• Data-dependent initialization, for example k-means on encoder outputs. Huh et al. (Appendix A.7) find the quantization error at initialization a strong indicator of later under-use.
• Codebook reset: re-running k-means every epoch, or replacing dead codes with randomly sampled encoder outputs; Huh et al. attribute the first to Łańcucki et al. (2020) and the second to Zeghidour et al. (2021) and Dhariwal et al. (2020). Among replacement policies they found a least-recently-used rule, replacing a code unused for 20 iterations, worked best.
• Stochastic quantization, sampling the code from a softmax over negative distances with a temperature that is often annealed; Huh et al. take Takida et al. (2022) as their example of this line of work, which they trace to Roy et al. (2018) and Kaiser et al. (2018).
• Huh et al.’s own changes (§4): an affine reparameterization of all code vectors with a shared mean and scale, so unselected codes move with the selected ones; alternating optimization; and a synchronized commitment loss.

In Huh et al.’s Table 1 (ResNet-18 classifying ImageNet-100, 1,024 codes of size 512), the share of active codes with a Kaiming-style initialization rose from 6.8% without a remedy to 16.3% with stochastic sampling, 99.0% with repeated k-means and 99.9% with least-recently-used replacement; with k-means initialization it was already 81.4% without a remedy. Image tokenizers that followed, such as VQGAN’s, kept the quantizer and changed the losses around it (see Variants).

Example: Codebook Use and the Straight-Through Gradient

The PyTorch example below (PyTorch 2.14.1, CPU, one thread, about 60 s) trains a small VQ-VAE on $24 \times 24$ images of one or two filled shapes (disks, squares and triangles of random size, position and intensity in $[0.4, 1]$ on a black background). The encoder maps each image to a $6 \times 6$ grid of 8-dimensional vectors, quantized against $K = 64$ codes. It trains four models for 1,000 steps each, crossing two codebook updates, the codebook loss trained by Adam and Sonnet-style EMA ($\gamma = 0.99$, debiased, Laplace-smoothed), with two initializations, which the papers leave open: Sonnet’s default uniform distribution on $[-\sqrt{3/D}, \sqrt{3/D}]$, and $K$ outputs of the untrained encoder sampled at random from those of the first 64 training images. All use $\beta = 0.25$. It prints dead codes and perplexity on 2,000 held-out images at steps 100, 300 and 1,000. The variational autoencoder article’s example measures the ELBO and interpolation of a continuous latent; this one measures codebook use and gradients.

The second half checks the straight-through gradient against finite differences in double precision on 200 held-out images, one random unit direction $u$ in latent space per image, for the reconstruction loss $L(z) = \lVert x - \mathrm{dec}(z) \rVert^2$:

• B1 compares the straight-through gradient with central differences of $L$ at $z_q$, the error divided by the gradient’s norm;
• B2 takes central differences of the quantized loss $L(\mathrm{quantize}(z))$ at $z_e$, with step $h = 10^{-5}$, and counts how many are exactly zero;
• B3 compares the straight-through gradient with the gradient of the unquantized loss $L$ at $z_e$, by cosine similarity, for all images and for the halves with the lowest and highest quantization error $\lVert z_e - z_q \rVert$.

import sys
import torch
from torch import nn
import torch.nn.functional as F

torch.set_num_threads(1)
SEED = int(sys.argv[1]) if len(sys.argv) > 1 else 0
STEPS = int(sys.argv[2]) if len(sys.argv) > 2 else 1000
S, D, K, BETA, GAMMA = 24, 8, 64, 0.25, 0.99
yy, xx = torch.meshgrid(torch.arange(S) + 0.5, torch.arange(S) + 0.5, indexing="ij")

def shapes(n, g):
    """n images 1x24x24 holding 1 or 2 filled shapes (disk, square or triangle) of random
    size, centre and intensity in [0.4, 1] on a background of 0."""
    img = torch.zeros(n, S, S)
    two = torch.rand(n, generator=g) < 0.5
    for j in range(2):
        k = torch.randint(0, 3, (n, 1, 1), generator=g)
        r = 3 + 4 * torch.rand(n, generator=g)
        cy, cx = r + (S - 2 * r) * torch.rand(2, n, generator=g)
        v = 0.4 + 0.6 * torch.rand(n, generator=g)
        dy, dx, r = yy - cy[:, None, None], xx - cx[:, None, None], r[:, None, None]
        disk = dy**2 + dx**2 <= r**2
        square = torch.maximum(dy.abs(), dx.abs()) <= 0.85 * r
        tri = (dy <= 0.7 * r) & (dy >= -r + 2 * dx.abs())
        m = torch.where(k == 0, disk, torch.where(k == 1, square, tri))
        m &= (two | (j == 0))[:, None, None]
        img = torch.where(m, v[:, None, None], img)
    return img[:, None]

class VQVAE(nn.Module):
    """24x24 image -> 6x6 grid of D-dimensional encoder outputs -> K-entry codebook -> decoder."""
    def __init__(self):
        super().__init__()
        self.enc = nn.Sequential(nn.Conv2d(1, 32, 4, 2, 1), nn.ReLU(), nn.Conv2d(32, 32, 4, 2, 1),
                                 nn.ReLU(), nn.Conv2d(32, D, 1))
        self.dec = nn.Sequential(nn.Conv2d(D, 32, 3, 1, 1), nn.ReLU(),
                                 nn.ConvTranspose2d(32, 32, 4, 2, 1), nn.ReLU(),
                                 nn.ConvTranspose2d(32, 1, 4, 2, 1))

def quantize(ze, codebook):
    """Nearest-neighbour lookup of every D-vector of ze (n, D, 6, 6) in the codebook (K, D)."""
    f = ze.permute(0, 2, 3, 1).reshape(-1, D)
    dist = (f**2).sum(1, keepdim=True) - 2 * f @ codebook.T + (codebook**2).sum(1)
    idx = dist.argmin(1)
    zq = codebook[idx].reshape(ze.shape[0], 6, 6, D).permute(0, 3, 1, 2)
    return zq, idx, f

class EMACodebook:
    """Debiased moving averages of code counts N_i and sums m_i, with Laplace smoothing of
    the counts, as in Sonnet's VectorQuantizerEMA; the codebook is set to m_i / N_i."""
    def __init__(self, eps=1e-5):
        self.eps, self.t, self.hN, self.hm = eps, 0, torch.zeros(K), torch.zeros(K, D)
    def update(self, f, idx, codebook):
        self.t += 1
        onehot = F.one_hot(idx, K).to(f.dtype)
        self.hN += (onehot.sum(0) - self.hN) * (1 - GAMMA)
        self.hm += (onehot.T @ f - self.hm) * (1 - GAMMA)
        N, m = self.hN / (1 - GAMMA**self.t), self.hm / (1 - GAMMA**self.t)
        N = (N + self.eps) / (N.sum() + K * self.eps) * N.sum()
        codebook.copy_(m / N[:, None])

def usage(model, codebook, x):
    """Dead codes (used by no latent of x) and perplexity exp(H(p)) of the code frequencies p."""
    with torch.no_grad():
        _, idx, _ = quantize(model.enc(x), codebook)
    p = torch.bincount(idx, minlength=K).double() / idx.numel()
    q = p[p > 0]
    return int((p == 0).sum()), torch.exp(-(q * q.log()).sum()).item()

x_train = shapes(20000, torch.Generator().manual_seed(SEED))
x_test = shapes(2000, torch.Generator().manual_seed(1000 + SEED))

def train(rule, init):
    torch.manual_seed(SEED)
    model = VQVAE()
    g = torch.Generator().manual_seed(SEED + 1)
    if init == "uniform":          # Sonnet's VarianceScaling(uniform): U(-sqrt(3/D), sqrt(3/D))
        cb = (2 * torch.rand(K, D, generator=g) - 1) * (3 / D) ** 0.5
    else:                          # K encoder outputs of the first training batch, at step 0
        with torch.no_grad():
            f = quantize(model.enc(x_train[:64]), torch.zeros(1, D))[2]
        cb = f[torch.randperm(len(f), generator=g)[:K]].clone()
    codebook = nn.Parameter(cb, requires_grad=(rule == "codebook loss"))
    opt = torch.optim.Adam(list(model.parameters()) + [codebook], lr=1e-3)
    ema = EMACodebook() if rule == "EMA" else None
    log = []
    for step in range(1, STEPS + 1):
        x = x_train[torch.randint(0, len(x_train), (64,), generator=g)]
        ze = model.enc(x)
        zq, idx, f = quantize(ze, codebook)
        rec = F.mse_loss(model.dec(ze + (zq - ze).detach()), x)   # straight-through
        loss = rec + BETA * F.mse_loss(ze, zq.detach())            # + commitment
        if ema is None:
            loss = loss + F.mse_loss(zq, ze.detach())              # + codebook loss
        opt.zero_grad(); loss.backward(); opt.step()
        if ema is not None:
            with torch.no_grad():
                ema.update(f, idx, codebook)
        if step in (100, 300, STEPS):
            log.append(usage(model, codebook.detach(), x_test))
    with torch.no_grad():
        mse = F.mse_loss(model.dec(quantize(model.enc(x_test), codebook)[0]), x_test).item()
    return model, codebook.detach(), log, mse

print(f"PyTorch {torch.__version__}, seed {SEED}; mean-image test MSE "
      f"{F.mse_loss(x_train.mean(0).expand_as(x_test), x_test):.4f}")
print(f"{'codebook update':15s} {'init':8s}| dead codes / perplexity at step 100, 300, {STEPS}"
      f" | test MSE")
trained = {}
for rule in ["codebook loss", "EMA"]:
    for init in ["uniform", "data"]:
        model, cb, log, mse = train(rule, init)
        trained[rule, init] = model, cb
        cells = "  ".join(f"{d:2d} / {p:5.2f}" for d, p in log)
        print(f"{rule:15s} {init:8s}| {cells} | {mse:.4f}")

# The straight-through gradient against finite differences, in float64.
print(f"{'codebook update':15s} {'init':8s}| B1 FD err: median, max | B2 FD = 0 | B3 cos: all  low err  high err")
for (rule, init), (model, cb) in trained.items():
    model, cb = model.double(), cb.double()
    g = torch.Generator().manual_seed(SEED + 2)
    L = lambda z, x: F.mse_loss(model.dec(z), x, reduction="sum")
    err1, zero2, cos3, qerr = [], [], [], []
    for x in x_test[:200].double().split(1):
        ze = model.enc(x).detach().requires_grad_(True)
        zq = quantize(ze, cb)[0].detach()
        g_ste, = torch.autograd.grad(L(ze + (zq - ze).detach(), x), ze)
        u = torch.randn(ze.shape, generator=g, dtype=torch.float64)
        u /= u.norm()
        h = 1e-5
        with torch.no_grad():
            Lq = lambda z: L(quantize(z, cb)[0], x)                 # the quantized loss
            fd1 = (L(zq + h * u, x) - L(zq - h * u, x)) / (2 * h)   # decoder loss at z_q
            fd2 = (Lq(ze + h * u) - Lq(ze - h * u)) / (2 * h)       # quantized loss at z_e
        dd = (g_ste * u).sum()
        err1.append(((fd1 - dd).abs() / g_ste.norm()).item())
        zero2.append(float(fd2 == 0))
        z = ze.detach().requires_grad_(True)
        g_true, = torch.autograd.grad(L(z, x), z)                   # unquantized loss at z_e
        cos3.append(F.cosine_similarity(g_ste.flatten(), g_true.flatten(), dim=0).item())
        qerr.append((ze - zq).norm().item())
    cos3, order = torch.tensor(cos3), torch.tensor(qerr).argsort()
    low, high = cos3[order[:100]].median(), cos3[order[100:]].median()
    print(f"{rule:15s} {init:8s}| {torch.tensor(err1).median():10.1e} {max(err1):10.1e}   | {sum(zero2) / 200:9.0%} |"
          f" {cos3.median():9.2f} {low:8.2f} {high:9.2f}")

Output (seed 0; two runs gave identical output):

PyTorch 2.14.1, seed 0; mean-image test MSE 0.0590
codebook update init    | dead codes / perplexity at step 100, 300, 1000 | test MSE
codebook loss   uniform | 53 /  2.98  47 /  5.11  44 /  6.19 | 0.0062
codebook loss   data    | 48 /  2.20  16 / 12.08   0 / 20.81 | 0.0049
EMA             uniform | 62 /  1.63  58 /  3.03  56 /  4.42 | 0.0077
EMA             data    | 48 /  3.28  41 /  9.18  23 / 14.34 | 0.0051
codebook update init    | B1 FD err: median, max | B2 FD = 0 | B3 cos: all  low err  high err
codebook loss   uniform |    9.6e-12    4.5e-05   |      100% |      0.14     0.21      0.07
codebook loss   data    |    9.4e-12    4.0e-11   |      100% |      0.32     0.45      0.25
EMA             uniform |    1.2e-11    3.1e-06   |      100% |      0.13     0.25     -0.06
EMA             data    |    8.8e-12    5.2e-05   |      100% |      0.23     0.33      0.16

Codebook use. Every model reconstructs far better than the mean image (test MSE 0.0049 to 0.0077 against 0.0590), yet codebook use differs widely. The initialization matters more than the update rule here: codes sampled from encoder outputs end with 0 and 23 dead codes (codebook loss and EMA), Sonnet’s uniform initialization with 44 and 56 and perplexities of only about 4 to 6 out of a possible 64. Two facts follow from the code’s construction, under either initialization. Under the codebook loss, a code that has never been selected has had an exactly zero gradient, so Adam has never moved it from its initial value. With the debiased EMA, whose hidden averages start at zero as in Sonnet, every code that receives no vector in the first update is set to $m_i / N_i = 0$, the origin, and stays there until it is selected. Such codes coincide, and since argmin returns the first of equal minima, at any step only the lowest-indexed of them can be selected; the next becomes selectable only once that one has been selected and moved away. The output does not show how many dead codes each mechanism accounts for, so the example does not isolate why, under data initialization, the codebook loss revived nearly every code by step 1,000 while EMA kept 23 dead. The better-used codebooks also reconstruct better.

The gradient. B1 shows that the straight-through gradient is exactly the derivative of the decoder’s loss at $z_q$: the median error is $10^{-11}$ of the gradient norm. The occasional larger maximum is, in our reading, a step of $\pm h$ that crosses a kink of the decoder’s ReLUs. B2 shows that the quantized loss has a zero derivative: no step of $10^{-5}$ changed the quantized loss, in any model. The straight-through gradient is therefore not the gradient of the quantized loss; it is a surrogate. B3 measures how far it is from the gradient at the unquantized point $z_e$: the median cosine is between 0.13 and 0.32, and in every model it is higher for the images whose latents lie closer to their codes, as Huh et al.’s bound on the gradient gap leads one to expect. In our reading, with only 64 codes and errors of this size, the surrogate is only loosely aligned with the continuous gradient, yet it was enough to train every model.

The script takes the seed as its first argument. With seeds 1 and 2 the pattern held for data initialization (codebook loss 3 and 1 dead codes, EMA 13 and 12) and for B2 (100% zero differences), and B3 was again higher in the low-error half for every model; with uniform initialization the codebook loss left fewer dead codes than EMA for seed 1 (49 against 53) and more for seed 2 (57 against 43), so the example does not rank the two rules under that initialization. B1’s median stayed near $10^{-11}$, with one maximum of $1.1 \times 10^{-3}$ (seed 2, EMA with data initialization).

Inference

Encoding and reconstruction. One encoder pass and a nearest-neighbour search over $K$ codes per position give the index grid; one decoder pass maps it back. Both are feed-forward, which VQ-VAE-2 presents as an advantage where fast encoding or decoding matters (Razavi et al., abstract).

Generation. The prior samples an index grid position by position, and the decoder renders it in one pass. Sampling the autoregressive model in the compressed latent rather than in pixels is, by Razavi et al.’s account, an order of magnitude faster (abstract) and about 30 times faster (§1), and they describe the latent as over 30 times smaller than the image. Class-conditional generation conditions the prior on the class.

Other uses. Because the decoder can be conditioned separately, VQ-VAE separates content from style in speech: latents from one speaker decoded with another speaker’s identity kept the content in the second voice (§4.3). Video generation conditioned on actions ran entirely in latent space, with each frame decoded afterwards (§4.4).

Performance

VQ-VAE (2017). On CIFAR-10, with the same architecture for each model, the VQ-VAE reached 4.67 bits per dimension, against 4.51 for a continuous VAE and 5.14 for VIMCO with discrete latents, all lower bounds after 250,000 steps at batch size 128 (§4.1). Bits per dimension is the negative log-likelihood in bits divided by the number of colour values (see the generative model article). The authors present this as the first discrete-latent model to approach continuous VAEs. ImageNet results are shown as reconstructions and samples rather than numbers (Fig. 2–3). On speech, mapping each of 128 codes to its most likely phoneme gave 49.3% accuracy on a 41-way classification, against 7.2% for always predicting the most likely phoneme (§4.3).

VQ-VAE-2 (2019) was evaluated on class-conditional ImageNet at $256 \times 256$ against BigGAN-deep:

• Negative log-likelihood. Top and bottom priors had training and validation NLL of 3.40 and 3.41, and 3.45 and 3.45; reconstruction MSE was 0.0047 (train) and 0.0050 (validation) (Table 1, units not stated). The authors read the small gaps as no overfitting and stress that the NLLs compare only priors over the same autoencoder (§5.2.1, which refers to these numbers as “Fig. 1”).
• Classification Accuracy Score, which trains an ImageNet classifier on generated samples and tests it on real test images: top-1 accuracy was 42.65% for BigGAN-deep, 54.83% for VQ-VAE-2 and 58.74% when the test images were first passed through the VQ-VAE, against 73.09% for training on real data (Table 2; top-5 65.92%, 77.59%, 80.98% and 91.47%).
• FID and Inception Score. Classifier-based rejection sampling improved both VQ-VAE-2’s Inception Score and its Fréchet Inception Distance, the FID going from about 30 to about 10. Reconstructions of real validation images themselves scored an FID of about 10 against about 2 for the originals, which the authors attribute to the Inception network’s sensitivity to slight blur; they therefore also report FID* against reconstructions (§5.3.1, which gives these approximate values in its text; Fig. 7a plots the curves).
• Precision and recall. Against BigGAN-deep, VQ-VAE-2 gave slightly lower precision and higher recall in the improved precision–recall metric over all 1,000 classes (Fig. 7b). For two classes, tench and ostrich, side-by-side samples show views and poses that BigGAN-deep’s samples at truncation 1.0 lack (Fig. 5).

An ablation in the supplementary material (Table 5) shows why the quantization bottleneck matters for the prior. On $128 \times 128$ ImageNet with a $32 \times 32$ latent, unregularized autoencoders reconstructed better (PSNR 32.41 dB with an 8-dimensional latent, up to 49.80 dB with 64), but the same PixelCNN prior needed 36,933 to 196,484 bits per image to model their latents, against 4,919 bits for VQ-VAE codes (512 codes of 64 dimensions, PSNR 27.19 dB).

Strengths

• Discrete codes with deterministic training. No sampling in the encoder and none of the variance of sampled gradients; the KL term is constant (van den Oord et al., §3.1).
• Latents that are used. The decoder receives only the codes, and the paper’s two-stage DeepMind Lab model kept using its 27-bit latent with a PixelCNN decoder (§4.2).
• A compact space for a likelihood model. The prior models tens of times fewer positions than pixels, with a softmax over $K$ codes, which made autoregressive models practical at $256 \times 256$ and above (Razavi et al., §1).
• Simple, fast encoder and decoder, both feed-forward (Razavi et al., abstract).
• Coverage. The prior is trained by maximum likelihood, which by the authors’ argument covers all modes of the data, unlike adversarial training (Razavi et al., §3.3).

Limitations

• Blurry reconstructions. With a squared-error loss, reconstructions of ImageNet looked slightly blurrier than the originals (van den Oord et al., §4.2), and Razavi et al. (§5.3.1) attribute part of VQ-VAE-2’s FID to the Inception network’s sensitivity to such blur. The authors suggested a more perceptual loss such as a GAN as future work; VQGAN took that step.
• Biased gradients and codebook collapse. The straight-through gradient is a surrogate, and codebook use depends on initialization and update rule (Huh et al., 2023; the example above).
• Slow sampling. Autoregressive sampling still visits every latent position in sequence.
• No straightforward interpolation. Razavi et al. (supplementary Appendix D) note that VQ-VAE-2 has no simple equivalent of interpolating latent vectors: a naive interpolation of two quantized vectors generally lies outside the codebook, and there is no commonly known way to interpolate between PixelCNN samples over discrete inputs.
• Two stages trained separately. The prior is fitted after the autoencoder, so it cannot shape the codes (§3.3).
• Evaluation. Razavi et al. (§6) write that measures of sample quality and diversity were in their infancy and visual inspection still necessary.

Historical Importance

The VQ-VAE (2017) showed that a VAE with discrete latents trained by nearest-neighbour quantization and straight-through gradients could approach a continuous VAE’s likelihood, and it established the two-stage recipe: compress images to a grid of tokens, then model the tokens with a powerful autoregressive network. VQ-VAE-2 (2019) took the recipe to $256 \times 256$ ImageNet and $1024 \times 1024$ faces with samples that its authors judged competitive with the best GANs, and argued that autoregressive modelling in latent space is a simple, effective objective for large generative models (§6). The quantizer, its codebook and commitment losses and its straight-through gradient were carried into VQGAN (Esser, Rombach & Ommer, 2021, Eq. 4), into token-based video models such as VideoGPT (see video generation), and, as the latent diffusion article describes, into one of the two first-stage autoencoders of latent diffusion.

Variants

VQ-VAE-2 (2019)

VQ-VAE-2 keeps the quantizer, uses EMA codebook updates and organizes the latents in a hierarchy (Razavi et al., §3). The aim is to model local information, such as texture, separately from global information, such as shape and object geometry, with a prior suited to each.

• Hierarchical latents. For $256 \times 256$ images the encoder downsamples by 4 to a $64 \times 64$ map that is quantized as the bottom latent, and further residual blocks downsample by 2 to a $32 \times 32$ top latent (§3.1, Fig. 2a). The bottom encoder is conditioned on the quantized top code, so the bottom level encodes what the top level lacks rather than repeating it (supplementary Algorithm 1). A feed-forward decoder of residual blocks and strided transposed convolutions takes all levels as input. For FFHQ at $1024 \times 1024$ a third level was added ($32 \times 32$, $64 \times 64$ and $128 \times 128$); in their three-level figure the maps are about 3,072, 768 and 192 times smaller than the image (Fig. 3). Each level had a codebook of 512 entries of dimension 64, with $\beta = 0.25$ (supplementary Table 3).
• Priors. The top prior is a gated PixelCNN with multi-head self-attention, as in PixelSNAIL, to capture long-range structure: on ImageNet 20 layers with an attention layer every five layers and 8 heads (supplementary Table 1, §A.2). The bottom prior, at four times as many positions, has no attention, for memory reasons, and is conditioned on the top codes through a deep residual conditioning stack (§3.2). The two priors are trained separately, which the supplementary material notes can be done in parallel (§A.1), and both are conditioned on the class (Fig. 2b).
• Classifier-based rejection sampling. A pretrained ImageNet classifier scores each sample by the probability of its intended class, and low scores are rejected, trading diversity for quality. It was used only for the quantitative metrics, not for the samples shown in the paper (§3.3).

The authors contrast their hierarchy with Dieleman et al.’s hierarchical VQ-VAE for music: with feed-forward decoders and a pixel MSE, they report no need for mitigations of hierarchy collapse (§4).

Later Quantized Autoencoders

VQGAN (Esser, Rombach & Ommer, 2021) keeps the VQ-VAE objective but replaces its $L^2$ reconstruction loss with a perceptual loss and adds a patch-based adversarial loss, so that images can be compressed further while keeping perceptual quality, and models the codes with a transformer instead of a PixelCNN (§3). The latent diffusion model article describes the VQ-regularized autoencoder of latent diffusion, which keeps a vector-quantization layer as the first layer of its decoder.

Implementations

• Sonnet (DeepMind, Apache License 2.0). VectorQuantizer (codebook and commitment losses) and VectorQuantizerEMA in sonnet/src/nets/vqvae.py on the v2 branch, for TensorFlow 2 (as of November 2026). Both return the quantized tensor with the straight-through trick, the loss, the codebook perplexity and the code indices; their docstrings place the commitment weight Beta in “equation 4” of the paper they link, the arXiv version; in the NeurIPS version the loss with $\beta$ is Eq. 3. VQ-VAE-2 cites the module at its earlier Sonnet v1 path, sonnet/python/modules/nets/vqvae.py (Razavi et al., §2.1).
• Elsewhere, the quantizer is a few lines in any framework, as in the example: a distance matrix, an argmin, a lookup, the stop-gradient trick and two squared-error terms.