Discrete Latents: VQ-VAE & Learned Codebooks

"I do not deal in shades of meaning. I have eight thousand words in my dictionary, and every patch of every image must choose one. It is a brutal vocabulary, but you would be amazed what you can say with it."

A Codebook With Strong Opinions About Vocabulary

Every latent in this chapter so far has been continuous: a real-valued vector, Gaussian-distributed, smooth. This section takes the opposite stance and makes the latent discrete. The motivation is partly that many things we model are inherently categorical (a phoneme, an object class, a word) and partly strategic: a discrete latent lets you put a powerful, expressive prior over the codes, a transformer or autoregressive model, and sample from it, which sidesteps the blur of a Gaussian decoder. The continuous VAE of Section 31.3 and the discrete VQ-VAE are the two great branches of the autoencoder family, and the discrete branch turned out to dominate the largest modern image generators.

1. Vector Quantization Intermediate

A VQ-VAE keeps the encoder and decoder but inserts a quantization step between them. The encoder, a convolutional network, maps an image to a grid of continuous vectors, one vector per spatial location, say a $32 \times 32$ grid of $D$-dimensional vectors for a $256 \times 256$ image. A learned codebook holds $K$ embedding vectors $e_1, \ldots, e_K$ (the dictionary, the vocabulary). Each encoder output vector $z_e(x)$ is replaced by its nearest codebook entry:

The decoder then reconstructs from the quantized grid $z_q$. The latent representation of the image is now just the grid of integer indices $k$, one per location, each pointing into the $K$-entry codebook. A $32 \times 32$ grid with a codebook of size $K = 8192$ describes the image with 1024 tokens, each a number from 0 to 8191, the image as a short paragraph in an 8192-word language. The illustration below dramatizes the moment where each patch of every image must choose one codebook entry, with no in-between option. Figure 31.6.1 shows the encode, quantize, decode pipeline and the token grid it produces.

2. The Straight-Through Estimator and the Losses Advanced

The $\arg\min$ that picks the nearest codebook entry is not differentiable: its gradient is zero almost everywhere, so backpropagation cannot pass through it to train the encoder. VQ-VAE bridges this with the straight-through estimator. On the forward pass it uses the quantized vector $z_q$; on the backward pass it pretends the quantization was the identity and copies the gradient from the decoder input straight back to the encoder output, as if $z_q = z_e$. In code this is the elegant trick z_q = z_e + (quantized - z_e).detach(): .detach() tells autograd to treat its argument as a constant with zero gradient, so the bracket contributes nothing on the backward pass. The forward value is therefore $z_e + (\text{quantized} - z_e) = \text{quantized}$, while the gradient sees only the bare $z_e$ and flows straight back into the encoder as if the quantization were not there.

That handles the encoder, but the codebook itself still needs a learning signal, and the encoder needs to be discouraged from producing outputs that drift far from any code. Two extra loss terms do this. The codebook loss $\lVert \text{sg}[z_e] - e \rVert^2$ pulls the chosen codebook entry toward the encoder output (sg is stop-gradient), and the commitment loss $\beta \lVert z_e - \text{sg}[e] \rVert^2$ pulls the encoder output toward its chosen code so the encoder "commits" to the codebook rather than oscillating. The two terms are the same distance written twice, with the stop-gradient placed on opposite sides on purpose. Each side moves only one of the two partners toward the other: the codebook loss updates the code while holding the encoder still, and the commitment loss updates the encoder while holding the code still. Splitting the distance this way lets you tune the two speeds independently through $\beta$.

The original VQ-VAE used $\beta = 0.25$, and the reason it is well below one is the mechanism just described: the codebook chases a moving target, since the encoder keeps shifting its outputs as it trains, so the codebook must be allowed to move toward the encoder faster than the encoder is pulled toward the codebook. A small $\beta$ keeps the encoder gently committed without letting it overrun the codebook's ability to keep up; push $\beta$ too high and the encoder locks onto stale codes that the codebook never catches, hurting reconstruction, which is why the value is robust enough that most implementations leave it at the default. The full objective is

Notice there is no KL term and no fixed prior during this stage. Unlike the continuous VAE, a VQ-VAE is trained purely as a reconstruction autoencoder with codebook upkeep; the generative prior over tokens is learned separately afterward (subsection 4). The next block implements the whole quantizer, including the straight-through estimator and both losses.

# The full vector quantizer: snap each encoder vector to its nearest
# codebook entry, compute the codebook and commitment losses, and apply the
# straight-through estimator so gradients reach the encoder past the argmin.
import torch
import torch.nn as nn
import torch.nn.functional as F

class VectorQuantizer(nn.Module):
    def __init__(self, num_codes=512, dim=64, beta=0.25):
        super().__init__()
        self.codebook = nn.Embedding(num_codes, dim)
        self.codebook.weight.data.uniform_(-1 / num_codes, 1 / num_codes)
        self.beta = beta

    def forward(self, z_e):                 # z_e: (B, dim, H, W)
        B, D, H, W = z_e.shape
        flat = z_e.permute(0, 2, 3, 1).reshape(-1, D)        # (B*H*W, D)
        # Nearest codebook entry by squared Euclidean distance.
        dist = (flat.pow(2).sum(1, keepdim=True)
                - 2 * flat @ self.codebook.weight.t()
                + self.codebook.weight.pow(2).sum(1))
        idx = dist.argmin(1)                                  # token indices
        quantized = self.codebook(idx).view(B, H, W, D).permute(0, 3, 1, 2)
        # Codebook loss (move codes to encoder) + commitment loss (move encoder to codes).
        codebook_loss = F.mse_loss(quantized, z_e.detach())
        commit_loss = F.mse_loss(z_e, quantized.detach())
        loss = codebook_loss + self.beta * commit_loss
        # Straight-through: forward uses `quantized`, gradient flows to z_e.
        quantized = z_e + (quantized - z_e).detach()
        return quantized, loss, idx.view(B, H, W)

vq = VectorQuantizer()
z_e = torch.randn(4, 64, 8, 8)             # encoder output: 8x8 grid of 64-dim vecs
zq, vq_loss, tokens = vq(z_e)
print("token grid:", tokens.shape)         # token grid: torch.Size([4, 8, 8])

3. The Training Math in Full Advanced

Subsection 2 stated the objective and the straight-through trick informally. Here we write the training mathematics out completely, in the notation of van den Oord, Vinyals, and Kavukcuoglu (2017), the paper that introduced the VQ-VAE, so that every gradient that flows during training is accounted for. The setup is the one already in play: the encoder produces $z_e(x)$ at each spatial location; the codebook is a set of $K$ vectors $\{e_k\}_{k=1}^{K}$; and quantization replaces $z_e(x)$ with its nearest entry,

To make the backward pass precise we introduce the stop-gradient operator $\text{sg}[\,\cdot\,]$. It is the identity on the forward pass, $\text{sg}[u] = u$, and has zero partial derivatives on the backward pass, $\nabla\,\text{sg}[u] = 0$; it is exactly PyTorch's .detach(). With it the full per-example loss is the sum of three terms,

Read the three terms by asking, for each one, which parameters can move because of it. The reconstruction term $\lVert x - \hat{x} \rVert_2^2$ trains the decoder directly, and trains the encoder through the straight-through estimator described next. The codebook term $\lVert \text{sg}[z_e(x)] - e \rVert_2^2$ has the stop-gradient on $z_e$, so only $e$ can move: it pulls the chosen code toward the encoder output, with the encoder held fixed. The commitment term $\lVert z_e(x) - \text{sg}[e] \rVert_2^2$ has the stop-gradient on $e$, so only $z_e$ can move: it pulls the encoder output toward its chosen code, with the code held fixed. The codebook and commitment terms are the same squared distance, $\lVert z_e(x) - e \rVert_2^2$, written twice with the stop-gradient on opposite arguments; that single device splits one distance into two one-directional pulls whose relative strength is set by $\beta$.

The straight-through estimator, exactly

The $\arg\min$ is piecewise constant, so $\partial z_q / \partial z_e = 0$ almost everywhere and the reconstruction gradient cannot reach the encoder through the quantizer. The straight-through estimator repairs this by defining the forward map $z_e \to z_q$ but declaring the backward map $\nabla_{z_q} \to \nabla_{z_e}$ to be the identity: whatever gradient the decoder sends to its input $z_q$ is copied unchanged to $z_e$, as though quantization were not in the path. Formally the decoder receives $z_q(x)$ on the forward pass, while during backpropagation we set

the approximation being the deliberate "lie" that $z_q = z_e$ for gradient purposes. This is exactly the identity z_q = z_e + (z_q - z_e).detach(): the forward value equals $z_q$, while the only term with a live gradient is the leading $z_e$, so $\nabla_{z_q}$ lands on $z_e$ unchanged. The commitment term is what keeps this approximation honest. By holding $z_e$ close to its chosen $e$, it keeps $z_q \approx z_e$, so copying the gradient across the gap introduces little error. The original paper sets $\beta = 0.25$ and reports the result is robust for $\beta$ anywhere in roughly $[0.1, 2.0]$, because $\beta$ only sets how hard the encoder commits, not what it represents; most implementations leave it at the default.

# The quantizer with the straight-through estimator and both loss terms,
# written to match the three-term objective term by term.
import torch
import torch.nn as nn
import torch.nn.functional as F

class STQuantizer(nn.Module):
    def __init__(self, num_codes=512, dim=64, beta=0.25):
        super().__init__()
        self.codebook = nn.Embedding(num_codes, dim)
        self.codebook.weight.data.uniform_(-1 / num_codes, 1 / num_codes)
        self.beta = beta

    def forward(self, z_e):                          # z_e: (B, dim, H, W)
        B, D, H, W = z_e.shape
        flat = z_e.permute(0, 2, 3, 1).reshape(-1, D)
        dist = (flat.pow(2).sum(1, keepdim=True)
                - 2 * flat @ self.codebook.weight.t()
                + self.codebook.weight.pow(2).sum(1))
        idx = dist.argmin(1)                          # k = argmin_j ||z_e - e_j||
        z_q = self.codebook(idx).view(B, H, W, D).permute(0, 3, 1, 2)
        codebook_loss = F.mse_loss(z_q, z_e.detach())          # ||sg[z_e] - e||^2
        commit_loss   = F.mse_loss(z_e, z_q.detach())          # ||z_e - sg[e]||^2
        loss = codebook_loss + self.beta * commit_loss
        z_q = z_e + (z_q - z_e).detach()             # straight-through estimator
        return z_q, loss, idx.view(B, H, W)

The EMA codebook update

The appendix of the original paper offers an alternative to the codebook loss term that often trains more stably. Instead of moving each code toward the encoder outputs by gradient descent on $\lVert \text{sg}[z_e] - e \rVert^2$, set each code to the running average of the encoder vectors assigned to it. Maintain, per code $k$, a count $N_k$ and an accumulated sum $m_k$, updated each batch with an exponential moving average of decay $\gamma$,

where $n_k$ is the number of encoder vectors assigned to code $k$ in the current batch and the sum runs over exactly those vectors. The code $e_k$ is then the count-normalized running mean of everything that has ever quantized to it, decayed so that recent batches matter more; the paper uses $\gamma \approx 0.99$. This is just an online k-means update on the codebook, which is why it tends to be smoother than the gradient version: it replaces the codebook loss term entirely. The commitment term stays exactly as before, still trained by gradient, because the encoder still needs to be pulled toward its codes; only the codebook's own learning signal changes. Production libraries default to this EMA scheme together with dead-code reinitialization, which is the combination that most reliably avoids codebook collapse.

4. Learning a Prior Over Tokens Advanced

A trained VQ-VAE is only an autoencoder; it cannot generate yet, because nothing tells you which token grids are plausible. The second stage fixes this by training a generative model over the token grids the encoder produces on the training set. Originally this was a PixelCNN-style autoregressive model predicting each token from those above and to the left; modern systems use a transformer, predicting tokens in raster order (autoregressive) or filling masked tokens (masked, BERT-style). To sample a new image you generate a fresh token grid from this prior, look up each token in the codebook to get the quantized grid, and decode. The two-stage structure, learn a tokenizer, then learn a prior over tokens, is the template for an enormous family of generative models. The next block shows the sampling procedure.

# Two-stage sampling: a transformer prior over token indices does the
# generating, then the VQ-VAE codebook and decoder render the grid to pixels.
# Stage 1 (a trained VQ-VAE) and Stage 2 (the `prior`) are trained separately.
import torch

@torch.no_grad()
def sample_image(prior, vqvae, grid_hw=(8, 8), start_token=0):
    H, W = grid_hw
    tokens = torch.full((1, H * W), start_token, dtype=torch.long)
    for i in range(H * W):                       # autoregressive token generation
        logits = prior(tokens)[:, i]             # predict token i
        tokens[:, i] = torch.distributions.Categorical(logits=logits).sample()
    tokens = tokens.view(1, H, W)
    quantized = vqvae.quantizer.codebook(tokens).permute(0, 3, 1, 2)  # indices -> vecs
    return vqvae.decoder(quantized)              # decode the token grid to an image
# The transformer prior does the generating; the VQ-VAE decoder only renders.
# This is the same recipe behind token-based text-to-image models.

5. VQ-VAE-2, VQGAN, and the Road to Diffusion Intermediate

Two refinements turned the VQ-VAE from a clever idea into production infrastructure. VQ-VAE-2 (2019) stacked the codebooks into a hierarchy, a top codebook for global structure and a bottom codebook for detail, the discrete cousin of the hierarchical VAE in Section 31.5, and paired it with a powerful autoregressive prior to produce sharp, large, diverse images, proving discrete latents scale. VQGAN (2021) made the autoencoder itself much better by adding a perceptual loss and an adversarial discriminator (the GAN idea of Chapter 32, used here only to sharpen the decoder) so that a small token grid could be decoded into a crisp high-resolution image. That compact, crisp tokenizer is exactly what made transformer-based image generation efficient, since the transformer only has to model a short sequence.

The line from here to the rest of Part IV is direct. VQGAN's autoencoder is the ancestor of the autoencoder that latent diffusion runs inside, and the continuous-latent AutoencoderKL of Section 31.3 and the discrete VQGAN tokenizer are the two options every latent-space generator chooses between. Latent diffusion (Chapter 33) mostly took the continuous route; autoregressive and masked-token image models took the discrete route. Both run their heavy generative model in the compressed latent space that an autoencoder from this chapter provides, which is the single most important practical lesson of Chapter 31: the autoencoder is the compression layer that makes large-scale generation affordable.

The chapter has now built three latents from the same skeleton: the plain code of Section 31.1, the samplable Gaussian code of Section 31.3, and the discrete codebook of this section. The Hands-On Lab below pulls all three into one runnable experiment so you can see, in a single figure, exactly how the constraint decides what the code becomes.

pip install torch torchvision matplotlib

import torch, torch.nn as nn, torch.nn.functional as F

device = "cuda" if torch.cuda.is_available() else "cpu"

class Encoder(nn.Module):
    """28x28x1 -> (C, 7, 7) feature grid."""
    def __init__(self, ch=32):
        super().__init__()
        self.net = nn.Sequential(
            nn.Conv2d(1, ch, 4, stride=2, padding=1), nn.ReLU(),   # 28 -> 14
            nn.Conv2d(ch, ch, 4, stride=2, padding=1), nn.ReLU())  # 14 -> 7
    def forward(self, x):
        return self.net(x)

class Decoder(nn.Module):
    """(C, 7, 7) -> 28x28x1 image in [0, 1]."""
    def __init__(self, ch=32):
        super().__init__()
        self.net = nn.Sequential(
            nn.ConvTranspose2d(ch, ch, 4, stride=2, padding=1), nn.ReLU(),  # 7 -> 14
            nn.ConvTranspose2d(ch, 1, 4, stride=2, padding=1))              # 14 -> 28
    def forward(self, z):
        return torch.sigmoid(self.net(z))

# TODO: load MNIST as tensors in [0, 1] and build a DataLoader with batch_size=128.
# Hint: torchvision.datasets.MNIST with transform=torchvision.transforms.ToTensor().
train_loader = ...

class PlainAE(nn.Module):
    def __init__(self, ch=32, zdim=16):
        super().__init__()
        self.enc, self.dec = Encoder(ch), Decoder(ch)
        self.to_z   = nn.Linear(ch * 7 * 7, zdim)
        self.from_z = nn.Linear(zdim, ch * 7 * 7)
        self.ch = ch
    def forward(self, x):
        h = self.enc(x).flatten(1)
        z = self.to_z(h)                       # the code (no constraint on its shape)
        # TODO: map z back through self.from_z, reshape to (B, ch, 7, 7), decode.
        recon = ...
        return recon, z
# Hint: loss is F.binary_cross_entropy(recon, x). Train ~10 epochs with Adam(lr=1e-3).

class VAE(nn.Module):
    def __init__(self, ch=32, zdim=16):
        super().__init__()
        self.enc, self.dec = Encoder(ch), Decoder(ch)
        self.to_mu     = nn.Linear(ch * 7 * 7, zdim)
        self.to_logvar = nn.Linear(ch * 7 * 7, zdim)
        self.from_z    = nn.Linear(zdim, ch * 7 * 7)
        self.ch = ch
    def forward(self, x):
        h = self.enc(x).flatten(1)
        mu, logvar = self.to_mu(h), self.to_logvar(h)
        # TODO: reparameterize: std = exp(0.5*logvar); z = mu + std * randn_like(std).
        z = ...
        g = self.from_z(z).view(-1, self.ch, 7, 7)
        return self.dec(g), mu, logvar

def vae_loss(recon, x, mu, logvar):
    bce = F.binary_cross_entropy(recon, x, reduction="sum") / x.size(0)
    # TODO: the closed-form KL of a diagonal Gaussian to N(0, I).
    kl = ...
    return bce + kl

class VectorQuantizer(nn.Module):
    def __init__(self, n_codes=64, dim=32, beta=0.25):
        super().__init__()
        self.codebook = nn.Embedding(n_codes, dim)
        self.codebook.weight.data.uniform_(-1/n_codes, 1/n_codes)
        self.beta = beta
    def forward(self, z):                         # z: (B, dim, H, W)
        B, C, H, W = z.shape
        flat = z.permute(0, 2, 3, 1).reshape(-1, C)         # (B*H*W, dim)
        d = (flat.pow(2).sum(1, keepdim=True)
             - 2 * flat @ self.codebook.weight.t()
             + self.codebook.weight.pow(2).sum(1))
        idx = d.argmin(1)                                    # nearest code per location
        zq = self.codebook(idx).view(B, H, W, C).permute(0, 3, 1, 2)
        # TODO: commitment + codebook loss, then straight-through: zq = z + (zq - z).detach()
        loss = ...
        zq_st = ...
        return zq_st, loss, idx

ae.eval()
with torch.no_grad():
    # TODO: collect z = ae.to_z(ae.enc(x).flatten(1)) over the test set,
    #       then sample z_rand ~ Normal(z.mean(0), z.std(0)) and decode it.
    z_rand = ...
    fake = ae.dec(ae.from_z(z_rand).view(-1, ae.ch, 7, 7))
# Display `fake`: expect blurry, non-digit blobs.

vae.eval()
with torch.no_grad():
    z = torch.randn(64, 16, device=device)                  # straight from N(0, I)
    samples = vae.dec(vae.from_z(z).view(-1, vae.ch, 7, 7))
    # TODO: encode two test images to their means mu_a, mu_b, then decode
    #       mu_a*(1-t) + mu_b*t for t in linspace(0, 1, 8).
    interp = ...
# Display `samples` (plausible digits) and `interp` (a smooth morph a -> b).

import matplotlib.pyplot as plt
# TODO: accumulate all `idx` from the VQ-VAE over the test set, then
#       used = torch.unique(all_idx).numel(); print(f"{used}/64 codes used").
#       Build a figure with: AE/VAE/VQ reconstructions, AE random (Step 5),
#       VAE samples + interpolation (Step 6), and a bar plot of code usage.
#       Save with plt.savefig("latent_showdown.png", dpi=150, bbox_inches="tight").

import torch, torch.nn as nn, torch.nn.functional as F
import torchvision, torchvision.transforms as T
import matplotlib.pyplot as plt

device = "cuda" if torch.cuda.is_available() else "cpu"
torch.manual_seed(0)

# ---- Step 1: data and shared backbone ----
ds = torchvision.datasets.MNIST("./data", train=True,  download=True, transform=T.ToTensor())
ts = torchvision.datasets.MNIST("./data", train=False, download=True, transform=T.ToTensor())
train_loader = torch.utils.data.DataLoader(ds, batch_size=128, shuffle=True)
test_loader  = torch.utils.data.DataLoader(ts, batch_size=256, shuffle=False)

class Encoder(nn.Module):
    def __init__(self, ch=32):
        super().__init__()
        self.net = nn.Sequential(
            nn.Conv2d(1, ch, 4, stride=2, padding=1), nn.ReLU(),
            nn.Conv2d(ch, ch, 4, stride=2, padding=1), nn.ReLU())
    def forward(self, x): return self.net(x)

class Decoder(nn.Module):
    def __init__(self, ch=32):
        super().__init__()
        self.net = nn.Sequential(
            nn.ConvTranspose2d(ch, ch, 4, stride=2, padding=1), nn.ReLU(),
            nn.ConvTranspose2d(ch, 1, 4, stride=2, padding=1))
    def forward(self, z): return torch.sigmoid(self.net(z))

def train(model, loss_fn, epochs=10, lr=1e-3):
    opt = torch.optim.Adam(model.parameters(), lr=lr)
    model.train()
    for ep in range(epochs):
        for x, _ in train_loader:
            x = x.to(device)
            opt.zero_grad()
            loss = loss_fn(model, x)
            loss.backward(); opt.step()
    return model

# ---- Step 2: plain autoencoder ----
class PlainAE(nn.Module):
    def __init__(self, ch=32, zdim=16):
        super().__init__()
        self.enc, self.dec = Encoder(ch), Decoder(ch)
        self.to_z, self.from_z = nn.Linear(ch*7*7, zdim), nn.Linear(zdim, ch*7*7)
        self.ch = ch
    def forward(self, x):
        z = self.to_z(self.enc(x).flatten(1))
        recon = self.dec(self.from_z(z).view(-1, self.ch, 7, 7))
        return recon, z

ae = PlainAE().to(device)
train(ae, lambda m, x: F.binary_cross_entropy(m(x)[0], x))

# ---- Step 3: VAE ----
class VAE(nn.Module):
    def __init__(self, ch=32, zdim=16):
        super().__init__()
        self.enc, self.dec = Encoder(ch), Decoder(ch)
        self.to_mu, self.to_logvar = nn.Linear(ch*7*7, zdim), nn.Linear(ch*7*7, zdim)
        self.from_z = nn.Linear(zdim, ch*7*7); self.ch = ch
    def forward(self, x):
        h = self.enc(x).flatten(1)
        mu, logvar = self.to_mu(h), self.to_logvar(h)
        z = mu + torch.exp(0.5 * logvar) * torch.randn_like(mu)   # reparameterize
        return self.dec(self.from_z(z).view(-1, self.ch, 7, 7)), mu, logvar

def vae_step(m, x):
    recon, mu, logvar = m(x)
    bce = F.binary_cross_entropy(recon, x, reduction="sum") / x.size(0)
    kl  = -0.5 * torch.sum(1 + logvar - mu.pow(2) - logvar.exp()) / x.size(0)
    return bce + kl

vae = VAE().to(device)
train(vae, vae_step)

# ---- Step 4: VQ-VAE ----
class VectorQuantizer(nn.Module):
    def __init__(self, n_codes=64, dim=32, beta=0.25):
        super().__init__()
        self.codebook = nn.Embedding(n_codes, dim)
        self.codebook.weight.data.uniform_(-1/n_codes, 1/n_codes)
        self.beta = beta
    def forward(self, z):
        B, C, H, W = z.shape
        flat = z.permute(0, 2, 3, 1).reshape(-1, C)
        d = (flat.pow(2).sum(1, keepdim=True)
             - 2 * flat @ self.codebook.weight.t()
             + self.codebook.weight.pow(2).sum(1))
        idx = d.argmin(1)
        zq = self.codebook(idx).view(B, H, W, C).permute(0, 3, 1, 2)
        loss = F.mse_loss(zq.detach(), z) * self.beta + F.mse_loss(zq, z.detach())
        zq_st = z + (zq - z).detach()                            # straight-through
        return zq_st, loss, idx

class VQVAE(nn.Module):
    def __init__(self, ch=32, n_codes=64):
        super().__init__()
        self.enc, self.dec = Encoder(ch), Decoder(ch)
        self.vq = VectorQuantizer(n_codes, ch)
    def forward(self, x):
        zq, vq_loss, idx = self.vq(self.enc(x))
        return self.dec(zq), vq_loss, idx

def vq_step(m, x):
    recon, vq_loss, _ = m(x)
    return F.binary_cross_entropy(recon, x) + vq_loss

vq = VQVAE().to(device)
train(vq, vq_step)

# ---- Step 5: the plain AE cannot be sampled ----
ae.eval(); xb, _ = next(iter(test_loader)); xb = xb.to(device)
with torch.no_grad():
    codes = torch.cat([ae.to_z(ae.enc(x.to(device)).flatten(1)) for x, _ in test_loader])
    z_rand = torch.randn(16, 16, device=device) * codes.std(0) + codes.mean(0)
    ae_fake = ae.dec(ae.from_z(z_rand).view(-1, ae.ch, 7, 7)).cpu()

# ---- Step 6: VAE sampling and interpolation ----
vae.eval()
with torch.no_grad():
    vae_samples = vae.dec(vae.from_z(torch.randn(16, 16, device=device))
                          .view(-1, vae.ch, 7, 7)).cpu()
    xa, xb2 = xb[0:1], xb[1:2]
    mu_a = vae.to_mu(vae.enc(xa).flatten(1)); mu_b = vae.to_mu(vae.enc(xb2).flatten(1))
    interp = torch.cat([vae.dec(vae.from_z(mu_a*(1-t) + mu_b*t).view(-1, vae.ch, 7, 7))
                        for t in torch.linspace(0, 1, 8)]).cpu()

# ---- Step 7: codebook usage and the composite figure ----
vq.eval()
with torch.no_grad():
    all_idx = torch.cat([vq(x.to(device))[2] for x, _ in test_loader])
    usage = torch.bincount(all_idx, minlength=64).cpu()
    vq_recon = vq(xb[:8])[0].cpu()
    ae_recon = ae(xb[:8])[0].cpu(); vae_recon = vae(xb[:8])[0].cpu()
print(f"VQ-VAE codebook usage: {(usage > 0).sum().item()}/64 codes used")

def row(ax, imgs, title):
    ax.imshow(torch.cat([i[0] for i in imgs], dim=1), cmap="gray"); ax.set_title(title); ax.axis("off")

fig, axs = plt.subplots(6, 1, figsize=(9, 11))
row(axs[0], ae_recon,  "AE reconstructions")
row(axs[1], vae_recon, "VAE reconstructions")
row(axs[2], ae_fake,   "AE: decoded random codes (garbage, the fatal flaw)")
row(axs[3], vae_samples, "VAE: samples from the prior (plausible digits)")
row(axs[4], interp,    "VAE: interpolation between two digits (smooth)")
axs[5].bar(range(64), usage); axs[5].set_title("VQ-VAE codebook usage"); axs[5].set_xlabel("code index")
plt.tight_layout(); plt.savefig("latent_showdown.png", dpi=150, bbox_inches="tight")
print("saved latent_showdown.png")

6. Further Reading Intermediate