Understand VAE / VQ-VAE

Chapter 1: Autoencoders — The Bottleneck

A plain autoencoder trains two neural networks end-to-end: the encoder f and decoder g. The loss is simply the reconstruction error: how different is g(f(x)) from x? The bottleneck — a narrow latent layer — forces compression.

It works. But there's a problem: the latent space is messy. Points are scattered unpredictably. If you pick a random point in latent space and decode it, you get garbage. The autoencoder memorizes efficient codes but doesn't organize them.

Concretely, the encoder is a stack of convolutional layers that progressively downsample. An image [batch, 3, 256, 256] passes through conv → downsample → conv → downsample until the spatial dimensions shrink and the channel count grows: [batch, 256, 32, 32]. A final linear layer (or global average pooling) flattens this to the latent vector [batch, latent_dim]. The decoder mirrors this process: upsample → conv → upsample until you're back to the original resolution.

python
# Minimal encoder architecture
class Encoder(nn.Module):
    def forward(self, x):        # x: [B, 3, 256, 256]
        x = self.conv1(x)         # [B, 64, 128, 128]
        x = self.conv2(x)         # [B, 128, 64, 64]
        x = self.conv3(x)         # [B, 256, 32, 32]
        x = x.flatten(1)         # [B, 256*32*32]
        return self.fc(x)         # [B, latent_dim]

To make this concrete: train an autoencoder on face images. Encode two faces and get z₁ and z₂. Decode the midpoint (z₁ + z₂) / 2 — you might get a coherent blend, or you might get static noise. The autoencoder never promised that midpoints would be meaningful. It only learned to reconstruct, not to organize. The latent space is full of "dead zones" where decoded output is garbage.

The VAE's fix is elegant: force every encoded point to look like it came from a standard Gaussian. Then sampling is trivial — just draw z ~ N(0,1) and decode. But this constraint comes at a cost: the encoder can't use arbitrary regions of latent space anymore, which limits reconstruction fidelity. This tension defines the entire field.

Chapter 2: The Variational Trick

The key insight of the VAE: instead of encoding x to a single point z, encode it to a distribution — specifically, a Gaussian with mean μ and variance σ². Then sample z from that distribution. This forces the latent space to be smooth.

In code, the encoder doesn't output one vector. It outputs two: μ (the mean) and log σ² (the log-variance). Why log-variance instead of variance? Because variance must be positive, and log-variance is unconstrained — the network can output any real number. To get the standard deviation: std = exp(0.5 * log_var).

python
# Encoder outputs TWO heads, not one
h = encoder_backbone(x)           # [B, 256, 32, 32] → flatten → [B, hidden]
mu      = self.fc_mu(h)           # [B, latent_dim]
log_var = self.fc_logvar(h)       # [B, latent_dim]

# Reparameterization trick
std = torch.exp(0.5 * log_var)   # [B, latent_dim]
eps = torch.randn_like(std)       # [B, latent_dim] ~ N(0,1)
z   = mu + std * eps              # [B, latent_dim] — differentiable!

But wait: sampling is not differentiable. How do we backpropagate through a random number? The reparameterization trick: instead of z ~ N(μ, σ²), write z = μ + σ · ε where ε ~ N(0,1). Now the randomness (ε) is external, and gradients flow through μ and σ.

Why does this matter so much? If we sampled z directly from N(μ, σ²), the sampling operation has no gradient with respect to μ or σ. Backprop hits a wall. But z = μ + σ · ε is a deterministic function of μ, σ, and the random noise ε. Gradients flow through μ and σ just fine — ε is treated as a constant input. The stochasticity is still there (each forward pass samples a different ε), but the computation graph is smooth.

Let's trace the gradient path to make this visceral. During backprop, we need ∂L/∂μ and ∂L/∂log_var. With reparameterization: ∂z/∂μ = 1 and ∂z/∂σ = ε. Both are simple, well-defined numbers. Without reparameterization, z is drawn from N(μ, σ²) and ∂z/∂μ doesn't exist — sampling is a black box. The trick turns a stochastic bottleneck into a deterministic computation with a random input.

Chapter 3: ELBO — The Training Objective

We want to maximize the likelihood of our data: p(x). But computing p(x) directly requires integrating over all possible latent codes z, which is intractable. The integral p(x) = ∫ p(x|z) p(z) dz sums over every possible latent code. With a 64-dimensional z, that's integration over a 64-dimensional space — computationally hopeless.

Instead, we optimize a lower bound on log p(x): the Evidence Lower Bound (ELBO). The gap between log p(x) and the ELBO is exactly KL(q(z|x) || p(z|x)) — the KL divergence between our approximate posterior and the true posterior. Since KL is always ≥ 0, the ELBO is always ≤ log p(x). Maximizing the ELBO simultaneously tightens the bound and improves the model.

The ELBO has two terms: a reconstruction term (how well can we decode z back to x?) and a KL divergence term (how close is our learned distribution q(z|x) to the prior p(z) = N(0,1)?). The reconstruction term wants good decoding; the KL term wants an organized latent space.

Let's make this concrete with numbers. Suppose for one image, the encoder outputs μ = 3.0, σ = 0.2 for some latent dimension. The KL term for that dimension is ½(3.0² + 0.04 − log 0.04 − 1) = ½(9 + 0.04 + 3.22 − 1) = 5.63. That's a big penalty — the KL is screaming "your mean is way off from zero!" If instead μ = 0.1, σ = 0.9, the KL is ½(0.01 + 0.81 + 0.21 − 1) = 0.015. Almost free. The reconstruction term pushes back: "but I need μ = 3.0 to encode this cat!" The balance between these two forces IS the VAE.

Chapter 4: Training a VAE

In practice, the KL term has a closed-form solution for Gaussians. For each latent dimension j, the KL divergence is: ½(μ_j² + σ_j² − log σ_j² − 1). This is cheap to compute and differentiable. The total KL sums over all latent dimensions.

The big practical knob is β: a weight on the KL term. With β=1, you get the standard VAE (ELBO). With β>1, you get β-VAE, which forces more disentanglement at the cost of blurrier reconstructions. With β<1, you get sharper images but messier latent space.

Why do VAEs produce blurry images? The reconstruction loss is usually MSE: ||x − x̂||². MSE penalizes pixel-level differences, so the decoder hedges its bets — when uncertain, it outputs the average, which is blurry. Modern VAEs fix this with perceptual loss (compare features from a pre-trained VGG network instead of raw pixels) or adversarial loss (add a discriminator that penalizes blurriness). Stable Diffusion's VAE uses both, which is why its reconstructions are sharp.

Let's trace through a concrete example. Suppose latent_dim = 4 and the encoder outputs:

Total KL = 5.87. With β = 1 and a reconstruction loss of, say, 150, the KL is only ~4% of the total — hence the model mostly focuses on reconstruction. With β = 10, KL becomes 58.7, nearly half the loss — now the model is strongly pushed toward N(0,1).

Here's the full training loop in code. Notice how each loss component maps to a specific part of the model:

python
# Forward pass
mu, log_var = encoder(x)                 # two heads
std = torch.exp(0.5 * log_var)
z = mu + std * torch.randn_like(std)     # reparam trick
x_hat = decoder(z)                       # reconstruction

# Reconstruction loss (applies to decoder)
recon_loss = F.mse_loss(x_hat, x, reduction='sum')

# KL loss (applies to encoder's mu and log_var heads)
kl_loss = -0.5 * torch.sum(1 + log_var - mu.pow(2) - log_var.exp())

# Total
loss = recon_loss + beta * kl_loss

Chapter 5: VQ-VAE — Discrete Codes

What if latent codes were discrete instead of continuous? VQ-VAE replaces the Gaussian latent space with a codebook: a table of K learned vectors, each of dimension D. In code, the codebook is just an nn.Embedding(K, D) — a matrix of shape [K, D]. Typical values: K = 8192 entries, D = 256 dimensions.

The encoder outputs a continuous feature map [batch, D, H, W] — say [1, 256, 32, 32]. For each of the 32×32 = 1,024 spatial locations, we find the nearest codebook entry and replace it. That's 1,024 independent nearest-neighbor lookups. The decoder only ever sees codebook entries, not the raw encoder output.

The loss has three parts, and each applies to different parameters:

python
# VQ-VAE quantization (per spatial location)
z_e = encoder(x)                    # [B, D, 32, 32]
z_flat = z_e.permute(0,2,3,1).reshape(-1, D)  # [B*1024, D]
dists = torch.cdist(z_flat, codebook.weight)    # [B*1024, K]
indices = dists.argmin(dim=1)                   # [B*1024]
z_q = codebook(indices)                         # [B*1024, D]

# Straight-through: forward=discrete, backward=continuous
z_q_st = z_e + (z_q.reshape_as(z_e) - z_e).detach()

python
def vq_straight_through(z_e, codebook):
    # Step 1: compute distances [B, K]
    dists = torch.cdist(z_e, codebook)  # Euclidean distance

    # Step 2: find nearest codebook entry
    indices = dists.argmin(dim=1)  # [B]

    # Step 3: look up codebook vectors
    z_q = codebook[indices]  # [B, D]

    # Step 4: straight-through estimator
    # Forward: value is z_q (discrete codebook vector)
    # Backward: gradient goes to z_e (as if no quantization)
    z_q_st = z_e + (z_q - z_e).detach()

    return z_q_st, indices

Chapter 6: Codebook Learning

The codebook is only useful if all entries are active. A common failure: codebook collapse — the encoder only uses a handful of entries while the rest go "dead." This wastes representational capacity.

Exponential Moving Average (EMA) updates are a popular fix. Instead of updating codebook entries with gradient descent, track the running average of all encoder outputs assigned to each entry. This is faster and more stable.

In practice, γ = 0.99 works well. For a codebook of K = 8192, D = 256, the EMA update processes all assigned vectors per batch in one vectorized operation. If entry k was assigned to 50 encoder outputs this batch, its new value is 0.99 × old + 0.01 × mean(those 50 vectors). Entries assigned to zero vectors don't update at all — this is exactly how they die.

python
# EMA codebook update
for k in range(K):
    mask = (indices == k)           # which encoder outputs chose this entry
    if mask.sum() == 0:             # dead code! no one chose it
        continue
    avg = z_flat[mask].mean(dim=0) # mean of assigned vectors
    codebook[k] = 0.99 * codebook[k] + 0.01 * avg

Dead code revival works by tracking usage counts per codebook entry. When an entry's count drops below a threshold (say, 1 assignment in the last 100 batches), it's declared dead. The revival strategy: replace the dead entry with the encoder output of a randomly chosen training example. This teleports the dead entry into a region of latent space that the encoder actually uses, giving it a fresh chance to attract assignments. After revival, the EMA update takes over and refines its position.

Chapter 7: FSQ — Finite Scalar Quantization

VQ-VAE's codebook is elegant but fragile: you need to manage dead codes, tune commitment loss, and balance EMA rates. FSQ takes a radically simpler approach: instead of learning a codebook, just round each scalar to a small set of levels.

If each of d dimensions has L levels, you get L^d possible codes — an implicit codebook. With d=6 and L=5, that's 5⁶=15,625 codes. No codebook parameters. No collapse. No commitment loss. Just rounding.

The implementation is almost embarrassingly simple. The encoder outputs a d-dimensional vector. Apply tanh to bound it to [−1, 1]. Scale by L. Round to the nearest integer. That's quantization. The straight-through estimator handles gradients: forward uses the rounded value, backward pretends the rounding didn't happen.

python
# FSQ: 4 lines of code for quantization
z = encoder(x)                       # [B, d] continuous
z_bounded = torch.tanh(z)            # [B, d] in [-1, 1]
z_scaled = z_bounded * (L - 1) / 2  # scale to [-L/2, L/2]
z_q = z_scaled.round()               # snap to integers
z_q = z_scaled + (z_q - z_scaled).detach()  # straight-through

Chapter 8: Image / Video Tokenizers

Modern generative models don't operate on raw pixels. They first tokenize images into compact representations using a VAE or VQ-VAE, then model the distribution of those representations using a transformer or diffusion model. This is the architecture behind Stable Diffusion, DALL-E, and MAGVIT.

The tokenizer compresses a 256×256 image to, say, a 32×32 grid of codebook indices. That's 1,024 tokens instead of 196,608 pixel values — a 192× compression. A transformer can then model these tokens autoregressively, just like words.

Each system uses a slightly different flavor. Stable Diffusion uses a continuous latent space (KL-regularized, not VQ), so the "tokens" are 4-channel feature maps at 64×64. DALL-E 1 used a proper discrete VQ-VAE codebook — each of its 32×32 spatial positions becomes a single integer index, giving 1,024 tokens that a GPT-style transformer can generate one by one. MAGVIT-2 pushes codebook size to 2¹⁸ = 262,144 entries using lookup-free quantization, achieving near-perfect reconstruction.

Think of the tokenizer as a lossy camera: everything downstream sees the world through this lens. If the tokenizer drops fine text, subtle textures, or thin structures, no generative model can hallucinate them back faithfully. This is why MAGVIT-2 pushed codebook size to 262K entries — more codes = finer-grained visual vocabulary = less information lost at the tokenization boundary.

For video, the challenge compounds: you need temporal consistency between frames. A video tokenizer like Cosmos's causal VQ-VAE processes frames in sequence, ensuring that the latent representation of frame t depends on frames 0...t-1. Without this, independently tokenized frames produce flickering artifacts when decoded. The spatial compression might be 8× per axis (256→32), and the temporal compression 4× (32 frames → 8 latent frames), giving a combined compression of 8×8×4 = 256×.

Chapter 9: VAEs in the Wild

VAEs and VQ-VAEs aren't just academic exercises — they're load-bearing infrastructure in the biggest generative AI systems. Stable Diffusion's latent space? A VAE. DALL-E's image tokens? A VQ-VAE. Every video model you've seen? Some flavor of temporal VQ-VAE.

Why does this 48× compression matter so much? Diffusion models need to run the U-Net (or transformer) hundreds of times per image. On raw pixels at 512×512×3 = 786K dimensions, each U-Net pass would cost ~100× more FLOPs than at 64×64×4 = 16K dimensions. The VAE encoder runs once, diffusion runs ~50 steps in the tiny latent space, then the VAE decoder runs once. It's the VAE that makes latent diffusion feasible on consumer GPUs.

Stable Diffusion's VAE is trained separately from the diffusion model. It uses a KL penalty (like a standard VAE) but tuned very lightly — the KL weight is small (around 1e-6), so the latent space is barely regularized. The focus is reconstruction quality: the VAE must encode and decode images with minimal perceptual loss, because every artifact in the VAE gets inherited by every generated image.

The training recipe for Stable Diffusion's VAE includes three loss components: (1) L1 pixel reconstruction, (2) perceptual loss using VGG features (compares high-level structure, not pixel values), and (3) adversarial loss from a PatchGAN discriminator (penalizes blurriness locally). The KL term is almost an afterthought — just enough to keep the latent distribution vaguely Gaussian so the diffusion model has a reasonable starting distribution.

python
# Stable Diffusion VAE loss (simplified)
loss = (1.0  * l1_loss(x, x_hat)        # pixel reconstruction
      + 1.0  * perceptual_loss(x, x_hat)  # VGG feature matching
      + 0.5  * adversarial_loss(x_hat)    # PatchGAN discriminator
      + 1e-6 * kl_loss(mu, log_var))     # tiny KL regularization