VAE, reparameterization & latent diffusion
Lesson 9 promised “latent diffusion: ~80 LOC, 10× faster” and moved on. This is the 80 lines. The “Latent” in Latent Diffusion is a VAE, and the trick that makes a VAE trainable is the same reparameterization you met in lesson 3.
Why diffuse in latent space at all
A 512×512×3 image is ~786k numbers. Running a UNet/DiT denoiser at that resolution for ~25–1000 steps is brutally expensive, and most of those pixels are perceptually redundant (smooth gradients, texture). The Latent Diffusion insight (Rombach et al. 2022, the “LDM” that became Stable Diffusion): compress the image once with an autoencoder, diffuse in the small latent, decode once at the end. SD’s autoencoder maps 512×512×3 → 64×64×4 — a 48× reduction in elements. The denoiser’s per-step cost drops with it; that is the entire “10× faster” (lesson 9) and most of what interviewers mean by SD’s “core optimization.”
But which autoencoder? The choice is the AE/VAE/VQ-VAE distinction interviewers ask about — and it matters because we are about to run a generative process in this latent.
AE vs VAE vs VQ-VAE: it’s about the latent’s distribution
All three are encoder → bottleneck → decoder trained to reconstruct. They differ in what distribution the latent code is forced into — and that determines whether you can generate by sampling the latent.
| Model | Latent | Reconstruction | Can you sample the latent directly? |
|---|---|---|---|
| AE | arbitrary, whatever minimizes recon | excellent | No — the encoder carves out unknown pockets; random z decodes to garbage. A compressor, not a generator. |
| VAE | regularized toward N(0,I) by a KL term | slightly blurrier | Yes — that’s the point. Sample z∼N(0,I), decode. |
| VQ-VAE (lesson 11) | nearest entry in a discrete codebook | sharp | Not from a fixed prior — you train a second model (PixelCNN / transformer / MaskGIT) over the discrete codes. |
Said in one line: AE learns any latent distribution; VAE forces it to be Gaussian; VQ-VAE forces it onto a discrete codebook. The design question is always “what makes the latent easy to put a generator on top of.”
The VAE objective, and why it needs a trick
A VAE wants to maximize log p(x) with a latent-variable model p(x) = ∫ p(x|z) p(z) dz. Same intractable integral as DDPM (lesson 3), same fix: a variational lower bound, with an encoder qφ(z|x) as proposal.
Don’t read this as two clumps of symbols — read it as a tug-of-war over the latent space, which is the whole point of a VAE. The first term, reconstruction, says “encode x to z, decode it back, and the result had better still look like x” — it wants the encoder to hand each input a private, well-separated code so nothing gets confused. The second term, the KL regularizer, pulls every per-sample posterior qφ(z|x) back toward one shared prior N(0,I) — it wants all those codes packed into a single tidy Gaussian blob with no empty gaps. Those two wants fight: recon wants codes spread apart, KL wants them collapsed together. The winner of that fight is the shape of the latent space, and a latent shaped like N(0,I) is exactly the one you can sample from — draw a random z, decode, get a new sample. (DDPM’s ELBO in lesson 3 is this same decomposition unrolled over a 1000-step chain. A diffusion model is a deep hierarchical VAE with a fixed Gaussian encoder.)
Claim: the VAE bound is a single tradeoff — “reconstruct well” vs. “keep the latent a clean Gaussian” — and where you land on that tradeoff is the only thing that decides whether you can generate.
- What we actually want. A space of codes z such that decoding a random one gives a realistic x. If the codes are scattered into unknown pockets with gaps between them, a random draw lands in a gap and decodes to garbage. So the goal is not just compression — it’s a latent with no holes.
- The recon term’s wish. Left alone, “decode z back to x” rewards giving every input its own far-flung code — maximum separation, easiest reconstruction. That is exactly the holey, un-samplable latent we don’t want.
- The KL term’s counter-wish. “Stay close to N(0,I)” pushes every code back toward the same blob, filling the gaps — but pushed too hard it crushes distinct inputs on top of each other, and reconstructions blur.
- The dial. The relative strength of these two terms is a knob. Strong recon, weak KL ⇒ an autoencoder: crisp but not samplable. Strong KL ⇒ samplable but blurry. A usable generative VAE sits in between.
Central point. The latent space is the deliverable; the two terms are just the forces that shape it, and balancing them is what turns a compressor into a generator.
The problem: the reconstruction term is an expectation over z ∼ qφ(z|x), and the encoder parameters φ sit inside the sampling distribution. You cannot backprop through “draw a random sample.” Sampling is not differentiable.
The reparameterization trick
Interviewers ask this directly. The move: push the randomness out of the parameters into an auxiliary variable. Instead of sampling z ∼ N(μφ(x), σφ(x)²) directly, write
The point of writing it this way: the randomness and the parameters now live in different places. All the dice-rolling is in ε, which has no parameters at all; the only parameter-bearing parts, μφ and σφ, are ordinary deterministic functions that just scale and slide that fixed draw. Because nothing random sits between φ and z anymore, the gradient ∂z/∂φ flows straight through — backprop never has to differentiate “roll a die,” only “stretch and shift a number I already have.” This is the identical maneuver as DDPM’s xt = √ᾱt x₀ + √(1−ᾱt) ε (lesson 2/3): separate “a deterministic affine map” from “a standard normal draw,” so the affine part is trainable. Same trick, two places.
Claim: you can’t backprop through a coin flip, so you move the coin flip out of the way and only differentiate the deterministic dressing around it.
- The blocker. The loss needs an average over z ∼ qφ(z|x), but the parameters φ we want gradients for are baked into the thing we’re sampling. “Draw a random sample” has no derivative w.r.t. the knobs that set the distribution — there’s no smooth handle to push on.
- The relocation. Pull the one fixed random ingredient, ε ∼ N(0,I), completely outside the parameters. Draw it first, then let the network only transform it: z = μφ + σφ ⊙ ε.
- Why that unblocks it. Now the path from φ to z is pure arithmetic — a multiply and an add — with the randomness frozen as a constant alongside. Arithmetic has derivatives; the gradient walks right through.
Central point. Keep the randomness as a parameter-free input you sample once, and the only thing left to differentiate is the deterministic stretch-and-shift — which is trivial.
The autoencoder Stable Diffusion actually uses
SD’s first stage is not a vanilla VAE. It is a VAE backbone trained with two extra losses (a VQGAN-style autoencoder), because plain VAE reconstructions are too blurry for a perceptual product:
- Reconstruction (L1/L2) — pixel fidelity.
- Perceptual / LPIPS — match deep features, not pixels, so edges and texture survive.
- Adversarial (a patch discriminator, lesson 21) — pushes decoded patches to look real, killing VAE blur.
- A very weak KL (or a VQ layer in the VQ variant) — just enough to keep the latent smooth and bounded so the diffusion model has a well-behaved space to work in, but not so strong it blurs reconstructions.
That last point is the subtlety most summaries gloss: SD’s autoencoder is barely a generative VAE at all — the KL is tiny, sampling its prior directly would look bad. Its only job is to give the diffusion model (the real generator) a compact, smooth, near-Gaussian-scaled latent. The latents are even rescaled by a constant so their variance ≈ 1, matching the noise schedule diffusion expects.
So the full Latent Diffusion training recipe (lesson 9’s ~80 LOC): (1) train the autoencoder on images, freeze it; (2) encode the dataset to latents; (3) run the entire DDPM/FM training of lessons 2–8 on the latents instead of pixels; (4) at sample time, diffuse a latent and decode once. Every earlier lesson applies unchanged — only the data tensor got smaller.
Interactive · reparameterized sampling & the KL knob
Left: a 1D encoder posterior N(μ, σ²). Watch the reparameterized draw z = μ + σε — the gray noise ε is fixed; μ, σ stretch and shift it (the differentiable part). Right: crank the KL weight β and watch the encoded latents of three classes either spread out (low β → sharp recon, gappy latent → bad for sampling) or collapse onto N(0,I) (high β → samplable, but classes overlap and recon blurs). SD lives at the far-left of this dial; a generative VAE lives near the middle.