MIT 6.7960 L15: Variational Autoencoders (VAE) — ELBO, Reparameterization Trick, and Latent Representations

🌏 中文版

Source: based on MIT 6.7960 Fall 2024 OCW (corresponds to OCW Lec 15). Videos, slides, and assignments are all open on MIT OCW. This lecture is taught by Phillip Isola.

The problem L14 left behind: p(x) is intractable

L14 covered the three big families of generative models (likelihood, autoregressive, latent-variable). Latent-variable models p(x) = ∫ p(x|z) p(z) dz look elegant, but the integral over z is intractable — you cannot directly maximize log p(x). This lecture is about how to get around that.

ELBO: replace log p(x) with an optimizable lower bound

Introduce an approximate posterior q(z|x) (learned by a neural net) and use Jensen's inequality:

log p(x) = log ∫ p(x|z) p(z) dz
         = log ∫ q(z|x) [p(x|z) p(z) / q(z|x)] dz
         ≥ E_{q(z|x)}[log p(x|z)] − KL(q(z|x) ‖ p(z))
        ≡ ELBO

The right side is the Evidence Lower Bound (ELBO): the first term is "expected log-likelihood of reconstructing x from z", the second is "KL between the approximate posterior and the prior." Maximizing ELBO = simultaneously (1) make q estimate well, and (2) make p(x|z) reconstruct x from z.

Write these two terms as a loss:

loss = E_q[−log p(x|z)] + KL(q(z|x) ‖ p(z))
     = recon_loss + kl_loss

p(z) is usually standard Normal N(0, I); q(z|x) is taken as diagonal Gaussian N(μ(x), diag σ²(x)), and the KL has a closed form.

Reparameterization trick: make sampling differentiable

The first ELBO term E_q[log p(x|z)] involves z ~ q(z|x) = N(μ, σ²) — a stochastic node that gradients cannot pass through. The reparameterization trick rewrites it as

z = μ(x) + σ(x) ⊙ ε,    ε ~ N(0, I)

Randomness moves from z to ε, and z becomes a deterministic function of μ, σ — gradients can now flow from the decoder back to the encoder through z. This "wire from N(0,1)" trick is the key to training VAEs with SGD.

A minimal VAE training skeleton

import torch, torch.nn as nn

class VAE(nn.Module):
    def __init__(self, z_dim=16):
        super().__init__()
        self.enc = nn.Sequential(nn.Linear(784, 256), nn.ReLU())
        self.mu  = nn.Linear(256, z_dim)
        self.lv  = nn.Linear(256, z_dim)  # log-variance
        self.dec = nn.Sequential(nn.Linear(z_dim, 256), nn.ReLU(),
                                 nn.Linear(256, 784), nn.Sigmoid())
    def forward(self, x):
        h = self.enc(x.view(-1, 784))
        mu, lv = self.mu(h), self.lv(h)
        z = mu + (0.5*lv).exp() * torch.randn_like(mu)  # reparameterize
        return self.dec(z), mu, lv

def loss(x, xh, mu, lv):
    recon = ((x.view(-1,784) - xh)**2).sum(-1).mean()
    kl = (-0.5 * (1 + lv - mu**2 - lv.exp())).sum(-1).mean()  # closed-form for N(0,I)
    return recon + kl, recon, kl

loss.backward() updates encoder and decoder together; to sample, decode a random z ~ N(0, I).

Practical pitfalls

• Posterior collapse: the decoder becomes so strong that the KL collapses to 0 and z is ignored. Common fixes: KL warm-up, free bits.
• Blurry outputs: VAE's likelihood at the pixel level (often Gaussian / MSE) hurts FID scores, but its latent structure is clean and great for interpolation / editing — that's the trade-off.
• β-VAE: weight KL by β > 1 to compress the latent → more disentangled but worse reconstruction.
• VQ-VAE (van den Oord et al.): quantize the latent to a discrete codebook, avoid posterior collapse; DALL·E and Stable Diffusion's latent diffusion build on this.
• Hierarchical / NVAE: multiple layers of latent, each contributing to the ELBO; more stable convergence, higher resolution.

Where VAE sits: comparison with L14's three families

VAE is the "cheapest of the likelihood family with the cleanest latent structure" compromise — and the Latent Diffusion / Stable Diffusion stack stands on its shoulders.

References

• MIT 6.7960 OCW (Fall 2024): course home
• Kingma & Welling, Auto-Encoding Variational Bayes: arXiv:1312.6114
• Doersch, Tutorial on Variational Autoencoders: arXiv:1606.05908
• Rezende et al., Stochastic Backpropagation and Approximate Inference in Deep Generative Models: arXiv:1401.4082