Neural Discrete Representation Learning, annotated

How to read this page

Nothing here assumes you already know the jargon. Three things help:

• Any dotted word explains itself when you hover it, tab to it, or tap it.
• Every symbol inside an equation does the same, and each equation is followed by a table of its symbols, a sentence reading it aloud, the numbers of a tiny example, and the same example in a few lines of Python you can run.
• The pictures are live: hover or tap the parts of the redrawn architecture, drag the encoder's output around the codebook and watch which entry it snaps to, and slide the averaging knob of the codebook update.

One tiny example runs through the whole page: a codebook of three entries in a flat, two-number world, e₁ = (0, 0), e₂ = (1, 0) and e₃ = (0, 1), and an encoder output z_e = (0.9, 0.2). The multimodal lesson uses the same three entries and the same vector, and builds the snapping step in code; the autoencoders lesson builds the autoencoder and the VAE this paper starts from.

Abstract

“Our model, the Vector Quantised-Variational AutoEncoder (VQ-VAE), differs from VAEs in two key ways: the encoder network outputs discrete, rather than continuous, codes; and the prior is learnt rather than static.”van den Oord, Vinyals and Kavukcuoglu (2017), Abstract. Read the original

Everyday picture

A paint-by-numbers kit comes with a palette of numbered pots. Any small patch of a picture can be described by the number of the pot closest to its colour, and anyone holding the same palette can repaint it from the numbers alone. A VQ-VAE learns the palette. Its encoder looks at a picture and writes a grid of pot numbers; its decoder repaints the picture from those numbers. Because the numbers come from a fixed, finite list, a second model can learn to write new grids of them, one number at a time, the way a language model writes words.

What the paper claims

• A new model: a variational autoencoder whose codes are discrete: each code is snapped to the nearest entry of a learned codebook, an idea borrowed from vector quantisation.
• It trains simply: no noisy gradient estimates, and no posterior collapse, the failure where a powerful decoder learns to ignore the codes.
• It holds its own: on CIFAR10 images its likelihood is close to that of a normal VAE with continuous codes, a first for discrete codes.
• The codes are useful: with a learned prior over them, it generates images, video and speech; it swaps one speaker's voice for another's; and its speech codes line up with phonemes, though it was never told what a phoneme is.

Why it matters today

Turning a picture or a sound into a short list of whole numbers is how many systems let a transformer read and write images and audio with the same next-token machinery it uses for text. This paper is where that recipe, “an autoencoder whose code is a grid of codebook ids”, took its standard form.

1 Introduction · original

“Language is inherently discrete, similarly speech is typically represented as a sequence of symbols.”van den Oord, Vinyals and Kavukcuoglu (2017), §1

Everyday picture

When you describe a photo to a friend you do not list the brightness of every pixel. You say “a fox on snow, looking left”. That description is short, made of a few words picked from a vocabulary, and it keeps what matters while dropping the noise. The paper wants a model that learns descriptions like that for itself: discrete (picked from a list) rather than continuous (any real number), and about the big features that span many pixels, not the fine grain.

Tiny example

A continuous code for a patch might be (0.9, 0.2): two real numbers, each of which could be anything. A discrete code for the same patch is “entry 2 of the codebook”: one whole number out of K. With K = 512 entries, that is exactly log₂ 512 = 9 bits, however detailed the patch was. The code cannot waste capacity on noise, because it has only 9 bits to spend.

The problem it answers

Two training goals are common for models of raw data: maximum likelihood (give the real data high probability) and reconstruction (rebuild the input from its code). The paper notes a tension, quoting earlier work: by likelihood alone, the best models have no codes at all, just a powerful autoregressive decoder that predicts each pixel from the ones before it, such as PixelCNN. Pair such a decoder with a VAE's code and the decoder tends to ignore the code. The paper's contributions:

1. The VQ-VAE: simple, discrete, free of posterior collapse, with no variance problems in its gradients.
2. Evidence that a discrete model matches continuous ones in log-likelihood.
3. With a powerful prior, coherent samples of images, video and speech.
4. Language-like structure learned from raw speech with no labels, and speaker conversion.

Why it matters

The argument that language, speech and even images are naturally described by symbols is the same argument that later put images and audio into the vocabulary of language models. See the “Generating images and speech” section of the multimodal lesson.

2 Related work · original

“In our experiments we were unable to train using the soft-to-hard relaxation approach from scratch as the decoder was always able to invert the continuous relaxation during training, so that no actual quantisation took place.”van den Oord, Vinyals and Kavukcuoglu (2017), §2

Everyday picture

Training a network means nudging every weight a little in the direction that lowers the error, using the gradient. A discrete choice, “pick entry 2”, has no gradient: nudge the input slightly and the choice stays 2, until suddenly it jumps to 3. Earlier discrete VAEs worked around this in two ways, and both hurt:

• Estimate the gradient by sampling (NVIL, VIMCO): try random choices and see how the loss changes. It works, but the estimates are noisy, so training is slow.
• Blur the choice (Concrete, or Gumbel-softmax): replace “pick 2” by a soft mixture that slowly sharpens into a single pick. Early on the gradients are smooth but wrong; late in training they become right but noisy.

Tiny example

Say a soft version of the choice gives weights (0.1, 0.8, 0.1) to the three entries. The decoder then sees 0.1·e₁ + 0.8·e₂ + 0.1·e₃ = (0.8, 0.1), a point that is none of the entries. A clever decoder can learn to read the exact weights back from such points, so it never has to live with the snap. That is the failure the quote describes.

Where the paper sits

These estimators had mostly been tested on small datasets such as handwritten digits (MNIST), with small latent spaces (below 8 dimensions, for example). The paper uses CIFAR10, ImageNet, DeepMind Lab and the VCTK speech corpus. It also builds on a line of work that puts autoregressive models inside VAEs (as decoder or as prior) and on image compression with quantised codes.

Why it matters

The VQ-VAE's answer, in §3.2, is to keep the snap hard and simply pretend it is not there when passing gradients back. That trick, crude as it sounds, is still how most codebook models are trained.

3 VQ-VAE · original

Everyday picture

A VAE has three parts: an encoder that turns the input into a code, a prior that says which codes are likely before seeing any input, and a decoder that turns a code back into data. In the usual VAE the code is a cloud of real numbers shaped like a bell curve (see the Auto-Encoding Variational Bayes companion). In a VQ-VAE the code is a choice from a list, and the list itself is learned: an embedding table, exactly like the table of word vectors at the bottom of a language model.

Why it matters

Swap “bell curve” for “pick from a table” and three things change at once: the encoder's output becomes a whole number, the prior becomes a probability for each table entry, and the decoder's input becomes a table row. The three subsections take them in turn.

3.1 Discrete latent variables · original

“One can see this forward computation pipeline as a regular autoencoder with a particular non-linearity that maps the latents to 1-of-K embedding vectors.”van den Oord, Vinyals and Kavukcuoglu (2017), §3.1

Everyday picture

A plain autoencoder squeezes an image into a small grid of vectors and rebuilds it. A VQ-VAE adds one step in the middle: every vector in the grid is replaced by the closest vector from a fixed set of K, the codebook. What the decoder sees is always one of those K vectors, and what gets stored is just its number.

Tiny example

The codebook has K = 3 entries, each D = 2 numbers wide: e₁ = (0, 0), e₂ = (1, 0), e₃ = (0, 1). The encoder outputs z_e = (0.9, 0.2). Its straight-line distances to the three entries are √0.85 = 0.922, √0.05 = 0.224 and √1.45 = 1.204. The smallest is to e₂, so the code is k = 2, and the decoder receives z_q = e₂ = (1, 0). Close to the encoder's output, but not equal.

Reading it: read from the top down. The encoder turns the image into a grid of vectors z_e(x). Each vector goes into “Nearest entry”, which compares it with every row of the codebook on the right and outputs the number of the closest one. Those numbers are the discrete code, q(z | x). Looking the numbers up in the same codebook gives z_q(x), a grid of codebook rows, and the decoder rebuilds the image from that. The dashed path on the left is the paper's trick for training: going backwards, the gradient that arrives at z_q is handed straight to z_e, skipping the snap, which has no gradient of its own (§3.2).

The math: equation 1, the snap as a probability

A VAE's encoder outputs a posterior: a probability for every possible code, given the input. Here it is as sharp as a probability can be: all of it on the nearest entry, none anywhere else. That is a one-hot vector.

In words: “the encoder puts probability 1 on the codebook entry nearest to its output, and 0 on every other entry.”

With the numbers: distances (0.922, 0.224, 1.204); the smallest is the second, so q(z | x) = (0, 1, 0).

In Python:

import math
z_e = [0.9, 0.2]
e = [[0, 0], [1, 0], [0, 1]]
# ‖z_e(x) − e_j‖₂ for each entry j
dist = [math.sqrt(sum((a - b) ** 2 for a, b in zip(z_e, e_j))) for e_j in e]
[round(d, 3) for d in dist]  # → [0.922, 0.224, 1.204]
# argmin_j, counting entries from 1 as the paper does
k = dist.index(min(dist)) + 1
k  # → 2
# q(z = j | x): 1 at the nearest entry, 0 elsewhere
[1 if j == k else 0 for j in range(1, len(e) + 1)]  # → [0, 1, 0]

The math: equation 2, what the decoder receives

In words: “the decoder's input is the codebook entry nearest to the encoder's output.”

With the numbers: k = 2, so z_q = e₂ = (1, 0). The snap moved the vector by (0.1, −0.2), a squared distance of 0.05.

In Python:

z_e = [0.9, 0.2]
e = [[0, 0], [1, 0], [0, 1]]
# squared distances are enough to find the nearest entry
d2 = [sum((a - b) ** 2 for a, b in zip(z_e, e_j)) for e_j in e]
k = d2.index(min(d2)) + 1
# z_q(x) = e_k: a lookup, like a row of a word-embedding table
z_q = e[k - 1]
z_q  # → [1, 0]
round(d2[k - 1], 2)  # → 0.05

A consequence: the KL term is a constant

A VAE's training objective, the evidence lower bound, has a KL divergence term that measures how far the encoder's posterior strays from the prior. With a one-hot posterior and a uniform prior (every entry equally likely, 1/K each) the term is the same for every input: log K.

In words: “only the chosen entry contributes to the sum; its probability 1 against the prior's 1/K costs log K, whatever the input.”

With the numbers: K = 3 gives log 3 = 1.0986 nats. The paper's image codebook, K = 512, gives log 512 = 6.238 nats, which is 9 bits per code.

In Python:

import math
K = 3
q = [0, 1, 0]
p = [1 / K] * K
# Σ_k q log(q / p); terms with q = 0 contribute nothing
kl = sum(q_k * math.log(q_k / p_k) for q_k, p_k in zip(q, p) if q_k > 0)
round(kl, 4)  # → 1.0986
round(math.log(512), 3), math.log2(512)  # → (6.238, 9.0)

A constant cannot be lowered by training, so the paper drops it from the loss. So nothing in a VQ-VAE's loss pulls the codes towards a fixed shape, the job the KL term does in an ordinary VAE.

One detail matters for everything that follows: the code is not one number per image. For speech, images and video the encoder outputs a 1D, 2D or 3D grid of vectors, and each vector is snapped on its own. A 32 × 32 grid is 1,024 ids.

Why it matters today

This forward pass, encode, snap to the nearest row, look the row up, decode, is exactly what the multimodal lesson builds: quantize is the argmin of equation 1 and dequantize the lookup of equation 2.

3.2 Learning · original

“To make sure the encoder commits to an embedding and its output does not grow, we add a commitment loss, the third term in equation 3.”van den Oord, Vinyals and Kavukcuoglu (2017), §3.2

Everyday picture

Three parties have to learn: the encoder, the decoder and the codebook. Think of a school with a fixed set of uniforms. The decoder is a tailor who learns to make each uniform look like the child wearing it. The codebook is the uniform supplier, who every term adjusts each uniform towards the children who were issued it. The encoder is the children, who are asked to stay close to the uniform they were given rather than growing off in some direction the supplier cannot follow. Each party is trained by a different term of one loss.

The gradient that jumps the gap: the straight-through estimator

The snap has no useful gradient. The paper's fix is the straight-through estimator: on the way forward, use the snapped vector z_q; on the way back, pass the decoder's gradient to the encoder unchanged, as if the snap were not there. This works because z_e and z_q live in the same D-dimensional space, so “move the decoder's input this way” is also sensible advice for the encoder's output.

Tiny example (the gradient is illustrative): forward, z_e = (0.9, 0.2) snaps to e₂ = (1, 0). Backward, suppose the decoder reports that its loss would fall if its input moved against the gradient g = (1.2, −1.0). The encoder receives the same g. A step with learning rate 0.5 moves z_e to (0.9 − 0.6, 0.2 + 0.5) = (0.3, 0.7). Its squared distances are now 0.58, 0.98 and 0.18, so next time it snaps to e₃. That is Figure 1 (right) of the paper: the gradient can push the encoder across a border into a different entry.

z_e = [0.9, 0.2]
e = [[0, 0], [1, 0], [0, 1]]
# the gradient at z_q, copied unaltered to z_e (illustrative numbers)
g = [1.2, -1.0]
lr = 0.5
z_e = [z - lr * g_i for z, g_i in zip(z_e, g)]
[round(z, 2) for z in z_e]  # → [0.3, 0.7]
d2 = [round(sum((a - b) ** 2 for a, b in zip(z_e, e_j)), 2) for e_j in e]
d2  # → [0.58, 0.98, 0.18]
# the new nearest entry
d2.index(min(d2)) + 1  # → 3

Try it: the embedding space

Reading it: this is the paper's Figure 1 (right) made live, in the page's two-number world. The three squares are the codebook entries; the circle is the encoder's output z_e. The solid line joins z_e to the entry it snaps to, and the readout gives all three distances, the chosen id, and the two loss terms that tie z_e and its entry together. Drag the circle (or focus the picture and use the arrow keys) across the invisible border halfway between two entries and watch the id change: that border is where a straight-through gradient step can change the code. Now press Codebook step: the chosen entry moves halfway to z_e, which is what the codebook term does. Press Commitment step: z_e moves only an eighth of the way to its entry, because the commitment term is scaled by β = 0.25. Both steps use the same learning rate; the entries chase the encoder hard, and the encoder is only tethered gently.

The math: equation 3, the whole loss

In words: “rebuild the input well from the snapped code; pull each codebook entry towards the encoder outputs that chose it, without moving the encoder; and pull the encoder's output a quarter as hard towards its entry, without moving the entry.”

A slip in the printed equation: L is a loss to be made small, and the paper calls the first term “the reconstruction loss”, so it should read minus log p(x | z_q(x)): lower loss means higher probability of rebuilding x. As printed, with a plus sign, minimising L would make reconstructions worse. The worked numbers below use the minus sign the text describes.

With the numbers: suppose the decoder gives the input probability 0.74 (illustrative), so the reconstruction term is −log 0.74 = 0.30. The codebook term is the squared distance from e₂ to z_e: 0.05. The commitment term is 0.25 × 0.05 = 0.0125. The total is 0.30 + 0.05 + 0.0125 = 0.3625.

In Python:

import math
z_e = [0.9, 0.2]
e_k = [1, 0]
beta = 0.25
# −log p(x | z_q(x)), illustrative decoder probability
recon = -math.log(0.74)
round(recon, 2)  # → 0.3
# ‖sg[z_e] − e‖²: sg only changes gradients, not the value
codebook = sum((a - b) ** 2 for a, b in zip(z_e, e_k))
round(codebook, 4)  # → 0.05
# β ‖z_e − sg[e]‖²: the same distance, a quarter of the weight
commitment = beta * codebook
round(commitment, 4)  # → 0.0125
round(0.30 + codebook + commitment, 4)  # → 0.3625

The two squared distances have the same value; the stop-gradient sg[·] decides who moves. Inside sg the quantity counts as a constant, so the codebook term's gradient reaches only the entry and the commitment term's gradient reaches only the encoder:

In words: “the codebook term tells the entry to move towards the encoder output; the commitment term tells the encoder output to move towards the entry, with a quarter of the force.”

With the numbers: the entry's gradient is 2((1, 0) − (0.9, 0.2)) = (0.2, −0.4); stepping against it moves e₂ towards (0.9, 0.2). The encoder's gradient is 2 × 0.25 × ((0.9, 0.2) − (1, 0)) = (−0.05, 0.1); stepping against it moves z_e towards (1, 0), four times more slowly than the entry moves towards z_e, since the only difference between the two gradients is the factor β = 0.25.

In Python:

z_e = [0.9, 0.2]
e_k = [1, 0]
beta = 0.25
# gradient of the codebook term, reaching only the entry
grad_e = [2 * (b - a) for a, b in zip(z_e, e_k)]
[round(g, 2) for g in grad_e]  # → [0.2, -0.4]
# gradient of the commitment term, reaching only the encoder
grad_z = [2 * beta * (a - b) for a, b in zip(z_e, e_k)]
[round(g, 2) for g in grad_z]  # → [-0.05, 0.1]
# the entry is pulled four times as hard: 1 / β
round(abs(grad_e[0]) / abs(grad_z[0]), 1)  # → 4.0

Who trains on what

The codebook needs its own term because the straight-through trick sends the reconstruction gradient around the snap: the entries receive nothing from it. The paper found training robust to β, with results unchanged from 0.1 to 2.0, and used β = 0.25 throughout, adding that the best value would in general depend on the scale of the reconstruction loss. With a grid of N codes (32 × 32 for ImageNet, 8 × 8 × 10 for CIFAR10), the two distance terms are averaged over the N positions.

The model's likelihood

To compare with other models, the paper needs log p(x), the log-probability of the data. In principle it sums over every possible code; in practice the decoder, trained only on the code the encoder picks, gives almost nothing to any other code, so one term dominates:

In words: “the probability of the data is the sum, over every code, of the chance of that code times the chance the decoder rebuilds x from it; once trained, only the encoder's chosen code contributes much, so keep that one term.”

With the numbers (illustrative): three codes, each with prior 1/3, and decoder probabilities 0.01, 0.74 and 0.002. The full sum gives log(0.752 / 3) = −1.3836; the single term gives log(0.74 / 3) = −1.3997. The one term is slightly lower, as it must be, since the sum includes it: the approximation is also a lower bound, as the paper notes.

In Python:

import math
p_x_given_z = [0.01, 0.74, 0.002]
p_z = [1 / 3] * 3
# log Σ_k p(x | z_k) p(z_k)
full = math.log(sum(a * b for a, b in zip(p_x_given_z, p_z)))
round(full, 4)  # → -1.3836
# keep only the chosen code, k = 2
one_term = math.log(p_x_given_z[1] * p_z[1])
round(one_term, 4)  # → -1.3997
one_term <= full  # → True

Why it matters today

Straight-through gradients, a codebook term and a commitment term: this three-part recipe carried over, largely unchanged, to the image and audio tokenizers that followed. The paper mentions one alternative for the codebook, updating it with moving averages instead of the codebook term; its Appendix A.1 (below) spells it out.

3.3 Prior · original

“Training the prior and the VQ-VAE jointly, which could strengthen our results, is left as future research.”van den Oord, Vinyals and Kavukcuoglu (2017), §3.3

Everyday picture

The VQ-VAE alone is a compressor: give it a picture and it gives back a grid of ids and a rebuild. To invent a picture you need someone who knows which grids of ids look like real pictures. That someone is the prior. While the VQ-VAE trains, the prior is a shrug, every id equally likely. Afterwards, a second model studies the grids of ids that real training images produce and learns to write new ones, one id at a time, the way a language model writes a sentence.

Tiny example (illustrative probabilities)

A 2 × 2 grid of ids, written left to right, top to bottom:

17174035 A uniform prior over K = 512 entries gives each id probability 1/512, so the grid costs 4 × 9 = 36 bits. A learned prior can do far better: say it gives the first id 0.5, the second (given the first) 0.8, the third 0.6 and the fourth 0.9. The grid's probability is 0.5 × 0.8 × 0.6 × 0.9 = 0.216, about 2.2 bits. The learned prior knows that 17 tends to follow 17.

In words: “the probability of a whole grid of ids is the chance of the first id, times the chance of the second given the first, and so on to the last.” This is the standard way an autoregressive model factorises a sequence; the paper says the prior “can be made autoregressive” without writing it out.

With the numbers: 0.5 × 0.8 × 0.6 × 0.9 = 0.216, which is −log₂ 0.216 = 2.21 bits, against 36 bits under the uniform prior.

In Python:

import math
# p(z_i | z_<i) for the four ids, in raster order
steps = [0.5, 0.8, 0.6, 0.9]
p_z = 1
for p_i in steps:
    p_z *= p_i
round(p_z, 3)  # → 0.216
round(-math.log2(p_z), 2)  # → 2.21
# the uniform prior used while training the VQ-VAE
round(-math.log2((1 / 512) ** 4), 1)  # → 36.0

Reading it: three rows, one after another in time. In stage 1 the VQ-VAE learns to squeeze images into ids and rebuild them; its prior is uniform, so the KL term is a constant and nothing about the prior is learned. In stage 2 the VQ-VAE is frozen and used only as a translator: every training image becomes a grid of ids, and a PixelCNN (for audio, a WaveNet) learns to predict each id from the ids before it. To sample, the prior writes a new grid id by id (ancestral sampling), the codebook turns ids into vectors, and the decoder paints the image. The expensive, pixel-level detail is the decoder's job; the prior only has to get the big structure right, on a grid far smaller than the image.

Why it matters today

Replace “PixelCNN” with “transformer” and add a caption before the ids, and this is how the first wave of text-to-image models generated pictures: see the multimodal lesson's section on discrete tokens, and the latent diffusion companion for the other road, where a diffusion model works in the compressed space instead.

4 Experiments · original

Everyday picture

The experiments ask two questions. Does forcing the code to be discrete cost much, compared with ordinary continuous codes? And are the discrete codes good: do they capture the meaningful structure of images, speech and video, so that a prior over them produces coherent samples?

Why it matters

A compressor can score well by keeping everything. The interesting results here are the ones where the code is forced to throw most of the input away, and what survives turns out to be the content: the scene layout, the words spoken.

4.1 Comparison with continuous variables · original

“Our model is the first among those using discrete latent variables which challenges the performance of continuous VAEs.”van den Oord, Vinyals and Kavukcuoglu (2017), §4.1

Everyday picture

A fair race: the same encoder and decoder, the same data, three ways of making the code. A normal VAE with continuous codes, the VQ-VAE, and VIMCO, a discrete VAE trained with the sampling-based gradient estimates of §2. The score is bits per dimension: on average, how many bits the model needs to encode each colour value of an image. Lower is better.

Tiny example

A CIFAR10 image is 32 × 32 pixels with 3 colours: 3,072 numbers. At 4.67 bits per dimension, the VQ-VAE's bound says an image can be stored in 4.67 × 3,072 = 14,346 bits. The continuous VAE needs 13,855 and VIMCO 15,790.

dims = 32 * 32 * 3
dims  # → 3072
# bits per image = bits per dimension × dimensions
[round(b * dims) for b in (4.51, 4.67, 5.14)]  # → [13855, 14346, 15790]

Reading it: each bar is one model's score on CIFAR10 in bits per dimension, from §4.1 of the paper; shorter is better, and all three are bounds on the true value, not exact figures. The continuous VAE (4.51) and the VQ-VAE (4.67) are close; VIMCO (5.14) trails by more than the gap between them. The striped bar is a reference the paper quotes from other work, a deep convolutional VAE at 4.54, to show its own continuous baseline is sound. The point is not that discrete codes win: they come close, while giving a compact symbolic code the others do not.

The setup

• Encoder: two convolutions with stride 2 and a 4 × 4 window (each halves the width and height), then two residual blocks, all 256 units wide.
• Decoder: two residual blocks, then two transposed convolutions with stride 2 that double the size back.
• Training: Adam with learning rate 2 × 10⁻⁴, 250,000 steps, batches of 128. VIMCO used 50 samples per estimate.

Why it matters

Before this paper, “discrete” meant “worse”. Closing most of that gap is what made discrete codes a practical choice rather than a curiosity.

4.2 Images · original

“Images contain a lot of redundant information as most of the pixels are correlated and noisy, therefore learning models at the pixel level could be wasteful.”van den Oord, Vinyals and Kavukcuoglu (2017), §4.2

Everyday picture

A 128 × 128 photo becomes a 32 × 32 grid of ids, like reducing a mosaic of 16,384 tiles to a sketch of 1,024 labelled squares. The rebuild comes back only slightly blurrier than the original, and a PixelCNN trained on the sketches produces new ones that decode into recognisable ImageNet classes: a kit fox, a grey whale, a coral reef.

Tiny example: how much smaller?

The original has 128 × 128 × 3 colour values of 8 bits each. The code has 32 × 32 ids of 9 bits each (K = 512).

In words: “bits in the image divided by bits in the code.”

With the numbers: 393,216 bits over 9,216 bits is 42.67; the paper writes ≈ 42.6. The DeepMind Lab frames later in this section, 84 × 84 × 3 into 21 × 21 ids, shrink by the same factor, since both cut each side by 4.

In Python:

import math
image_bits = 128 * 128 * 3 * 8
# each id is one of K = 512 entries: log₂ 512 bits
code_bits = 32 * 32 * int(math.log2(512))
image_bits, code_bits  # → (393216, 9216)
round(image_bits / code_bits, 2)  # → 42.67
round((84 * 84 * 3 * 8) / (21 * 21 * 9), 2)  # → 42.67

Three more results

• ImageNet samples. The PixelCNN over the 32 × 32 codes needs only spatial masking, since each position holds one id rather than three colour channels. Its samples, decoded, show the requested classes. The reconstructions use plain squared error; the paper notes a more perceptual loss, “e.g., a GAN”, could sharpen them, and leaves that as future work.
• DeepMind Lab. Frames of 84 × 84 × 3 from a 3D game environment are coded as 21 × 21 ids and rebuilt nearly perfectly; a PixelCNN prior over the codes produces plausible new frames.
• A second VQ-VAE on top. A second VQ-VAE, with a PixelCNN as its decoder, compresses the 21 × 21 grid of ids into just 3 ids of K = 512 each: 3 × 9 = 27 bits for a whole frame, less than a single 32-bit floating-point number. A PixelCNN decoder is exactly the kind of powerful decoder that makes ordinary VAEs collapse; here the three codes are still used. The rebuilds keep the room layout and textures without copying the pixels.

import math
# the two-stage code for a whole DeepMind Lab frame
bits = 3 * int(math.log2(512))
bits, bits < 32  # → (27, True)

Why it matters

A 42-fold smaller code both speeds up the prior and frees it to model the global structure of a picture instead of the grain of individual pixels. That division of labour, a compressor for detail and a separate model for structure, is the design of most image generators since, whether the second model is autoregressive or a diffusion model.

4.3 Audio · original

“This means that the VQ-VAE has, without any form of linguistic supervision, learned a high-level abstract space that is invariant to low-level features and only encodes the content of the speech.”van den Oord, Vinyals and Kavukcuoglu (2017), §4.3

Everyday picture

Transcribe a sentence and give the transcript to a different person to read aloud. The words survive; the voice, the pitch and the rhythm do not. The speech VQ-VAE turns out to do something similar with no transcript and no labels: its codes keep what was said and drop who said it, because the decoder is told who the speaker is and so the code does not need to carry it.

Tiny example

The encoder has 6 convolutions of stride 2, so each halves the length: 2⁶ = 64 audio samples become one code. In the phoneme study, the codes run at 25 per second with a downsampling of 640, which means the audio had 25 × 640 = 16,000 samples per second. With K = 128 entries, each code is log₂ 128 = 7 bits, so a second of speech is 25 × 7 = 175 bits of code.

import math
2 ** 6  # → 64
# codes per second × samples per code = samples per second
25 * 640  # → 16000
# bits of code per second of speech, K = 128
25 * int(math.log2(128))  # → 175
# the prior's training chunks: 40,960 samples
40960 / 16000, 40960 // 128  # → (2.56, 320)

Reading it: speaker A's recording goes down the left column, through the encoder, into a sequence of ids 64 times shorter than the waveform. The decoder rebuilds audio from two inputs: those ids, and a speaker identity that enters from the right. Give it A's identity and you get a rebuild of the original, with the same words but a different waveform and altered prosody (the melody and rhythm of speech). Give it B's identity and you get the same words in B's voice. That swap only works because the ids do not carry the voice: training rewarded the code for keeping only what the decoder could not get from the speaker id.

The four findings

1. Only the content survives. On VCTK (109 speakers) with codes 64× shorter than the audio, rebuilds keep the words but not the exact waveform or prosody.
2. The prior babbles in words. On a 460-speaker corpus, with codes 128× shorter, a prior trained on chunks of 2.56 seconds (320 codes) produces samples with “clear words and part-sentences”, where the original WaveNet's unconditional samples sound like babbling.
3. Speaker conversion works, as in the diagram.
4. The codes resemble phonemes. Mapping each of the 128 codes to its most likely phoneme (of 41) classifies phonemes correctly 49.3% of the time, against 7.2% for a random code space. No phoneme labels were used in training.

Reading it: the two bars are phoneme accuracy on a 41-way choice, from §4.3 of the paper. The striped bar, 7.2%, is what you get by always guessing the most common phoneme, which is all a meaningless code allows. The solid bar, 49.3%, is the VQ-VAE's codes mapped one to one onto phonemes. The paper adds, in a footnote, that a code might mean different things after different neighbours, like pairs or triples of sounds, so a smarter mapping could score higher.

Why it matters today

Discrete audio codes are now the usual way to let a language model listen and speak. Neural audio codecs stack several codebooks, each coding what the previous one missed (residual vector quantisation); the multimodal lesson builds that stack with learn_residual_codebooks and residual_quantize, and the Whisper companion shows the other side of the ear, speech read as a spectrogram rather than as codes.

4.4 Video · original

“Therefore, VQ-VAE can be used to imagine long sequences purely in latent space without resorting to pixel space.”van den Oord, Vinyals and Kavukcuoglu (2017), §4.4

Everyday picture

A flip-book artist does not paint each page in full while planning the story. They sketch the whole sequence as stick figures first, then paint. Here the prior plans a video in codes, frame after frame, and only at the end does the decoder paint each frame.

Tiny example

Give the model 6 real frames from DeepMind Lab, then an action (“move forward”, or “move right”) for each of the next 10. The prior generates the codes z₇ to z₁₆ from the codes before them and the actions; the decoder turns each into a frame. If each frame is coded like the images above, as 21 × 21 ids, that is 4,410 ids for the 10 new frames, every one written before a single pixel is drawn. (The paper does not give the video model's grid size; the 21 × 21 here is illustrative.)

# ten generated frames, each an illustrative 21 × 21 grid of ids
10 * 21 * 21  # → 4410

What the paper reports

Its Figure 7 shows the generated frames turning or moving forward as instructed, keeping the local geometry right with no loss of visual quality. A model trained without actions gave similar results.

Why it matters today

Planning in a compact code and rendering at the end is how later video and world models keep long generations affordable. The video section of the multimodal lesson counts what video costs in tokens.

5 Conclusion · original

Everyday picture

The paper's message fits in a sentence: give an autoencoder a finite vocabulary for its codes, train through the snap by pretending it is not there, and the codes it learns are compact, meaningful, and easy for a second model to write.

What it showed

• Likelihoods on CIFAR10 almost as good as continuous VAEs.
• Coherent samples of 128 × 128 colour images, action-conditioned video, and speech with recognisable words.
• Speaker conversion, and speech codes closely related to phonemes, learned with no supervision at all.

Why it matters

Discrete, learned codes turned out to be a general adapter between raw signals and sequence models. The next section traces where they went.

Appendix A.1: codebook updates with moving averages · original

“This update is typically used in algorithms such as K-Means.”van den Oord, Vinyals and Kavukcuoglu (2017), Appendix A.1

Everyday picture

The codebook term of equation 3 nudges each entry towards the encoder outputs that chose it. If you could see all of them at once, the best place for the entry is simply their average, which is the update step of k-means clustering. Training sees only a small batch at a time, so the appendix keeps a running average instead: an exponential moving average that remembers most of the past and mixes in a little of each new batch.

Tiny example

Entry e₂ sits at (1, 0). In this batch, two encoder outputs chose it: (0.9, 0.2) and (0.7, 0.4). Their average is (0.8, 0.3), which is where k-means would put the entry. The moving average instead blends the old position with the new evidence.

The math

The codebook term, for one entry e_i and the n_i encoder outputs z_i,1, …, z_i,ni nearest to it:

In words: “the total squared distance from one entry to the outputs that chose it is smallest when the entry sits at their average.”

With the numbers: the average of (0.9, 0.2) and (0.7, 0.4) is (0.8, 0.3), and the total squared distance there is 0.02 + 0.02 = 0.04. At the old position (1, 0) it was 0.05 + 0.25 = 0.30.

In Python:

z_i = [[0.9, 0.2], [0.7, 0.4]]
n_i = len(z_i)
# e_i = (1/n_i) Σ_j z_ij
e_i = [sum(z[d] for z in z_i) / n_i for d in range(2)]
[round(v, 2) for v in e_i]  # → [0.8, 0.3]
def total(e):
    return sum(sum((a - b) ** 2 for a, b in zip(z, e)) for z in z_i)
round(total(e_i), 2), round(total([1, 0]), 2)  # → (0.04, 0.3)

The online version keeps two running totals per entry, a count N_i and a sum m_i, and divides them:

In words: “keep a fading count of how many outputs chose this entry and a fading sum of those outputs; the entry is the sum divided by the count, a smoothed average over recent batches.”

With the numbers (γ = 0.5, illustrative, so the change is visible; the paper uses 0.99): start from N = 1 and m = (1, 0), so e = (1, 0). After the batch, N = 1 × 0.5 + 2 × 0.5 = 1.5 and m = (1, 0) × 0.5 + (1.6, 0.6) × 0.5 = (1.3, 0.3), so e = (0.867, 0.2): part of the way from (1, 0) to the batch average (0.8, 0.3). With γ = 0.99 the same batch moves e only to (0.996, 0.006).

In Python:

def ema_step(N, m, batch, gamma):
    # N := N·γ + n(1 − γ);  m := m·γ + Σ z (1 − γ);  e := m / N
    n = len(batch)
    total = [sum(z[d] for z in batch) for d in range(2)]
    N = N * gamma + n * (1 - gamma)
    m = [m_d * gamma + t_d * (1 - gamma) for m_d, t_d in zip(m, total)]
    return N, m, [m_d / N for m_d in m]
batch = [[0.9, 0.2], [0.7, 0.4]]
N, m, e = ema_step(1, [1, 0], batch, gamma=0.5)
N, [round(v, 2) for v in m], [round(v, 3) for v in e]  # → (1.5, [1.3, 0.3], [0.867, 0.2])
N, m, e = ema_step(1, [1, 0], batch, gamma=0.99)
[round(v, 3) for v in e]  # → [0.996, 0.006]

Reading it: the x-axis counts batches, the y-axis is the first coordinate of one codebook entry. The entry starts at 1, and every batch its chosen outputs average about 0.8, plus noise (the batch averages are made up for this chart, from a seeded random generator). The jagged line is each batch's own average: where k-means would jump every time. The smooth line is the moving-average entry. Slide γ: at 0.5 it still jumps about with the noise; at 0.99, the paper's value, it barely wobbles but takes well over a hundred batches to arrive. A larger γ trades speed for steadiness.

Why it matters today

The paper did not use this update for its own experiments (§3.2 says so), but many later VQ models adopted it. The multimodal lesson learns its codebooks with the plain, all-at-once version: kmeans_codebook snaps every vector and moves each entry to its vectors' mean.

What changed since 2017

Everyday picture

The paint-by-numbers palette became standard equipment. The palette grew, the painters got better at hiding the seams, and the model writing the numbers changed from a PixelCNN to a transformer, but the idea of this paper, a learned codebook in the middle of an autoencoder, stayed.

Why it matters

Whenever a system says it “tokenizes” an image or a sound, there is very likely a descendant of this paper's encoder, codebook and decoder inside it, trained with some version of equation 3.

Glossary