45 · Variational autoencoders (VAEs)
A VAE turns the bottleneck into a probability distribution. By encoding to a mean and variance, regularizing it toward N(0,1), and sampling with the reparameterization trick, the latent space becomes smooth and sampleable — a true generator.
A VAE encodes each input to a DISTRIBUTION (a mean and variance) instead of a point, regularizes that distribution toward N(0,1) with a KL term, and trains through the randomness via the reparameterization trick — making the latent space smooth enough to SAMPLE new data from.
Without this:
The VAE is the first true deep generative model and the encoder/decoder that Stable Diffusion runs inside. Its KL + reconstruction loss is the template for the ELBO that recurs across generative modeling.
Last lesson ended with a problem: a plain autoencoder's latent space has holes, so you can't sample from it to generate. The variational autoencoder (VAE) fixes this with one elegant change — make the bottleneck probabilistic.
Instead of the encoder mapping x to a single point z, it maps x to a distribution over latent space: a mean vector μ and a (log-)variance logσ². To get a code you sample z ∼ N(μ, σ²), then decode z as before. Encoding to a fuzzy cloud rather than a sharp point means nearby codes must decode to similar outputs — the latent space becomes continuous.
The loss has two terms:
- Reconstruction —
‖x − x̂‖²(or a cross-entropy), exactly as before: the decoded sample should look like the input. - KL divergence —
KL(N(μ,σ²) ‖ N(0,1)), a regularizer that pulls every encoded distribution toward the standard normal. This packs all the per-example clouds into one shared, well-shaped region around the origin, with no gaps. Now a code drawn fromN(0,1)lands somewhere the decoder understands — so you can generate by samplingz ∼ N(0,1)and decoding.
Together these two terms are (the negative of) the ELBO — the evidence lower bound — the quantity VAEs maximize. Reconstruction says "be faithful"; KL says "stay organized and sampleable". They pull against each other, and the balance is the whole game.
One obstacle remains: you cannot backpropagate through a random sample. The fix is the reparameterization trick: write z = μ + σ · ε where ε ∼ N(0,1) is drawn outside the network. Now the randomness lives in ε (no gradient needed there), while μ and σ are deterministic functions of the network — so gradients flow straight through them. The cell below shows this trick reproducing a target distribution exactly, and computes the KL term in closed form.
Python (in browser)
Python runs entirely in your browser via Pyodide (~6 MB on first Run, cached after).
A VAE in PyTorch: encoder outputs mu and logvar, the reparameterization trick samples z = mu + sigma*eps, and the loss sums reconstruction + KL (the negative ELBO); generate by decoding z ~ N(0,1).
What is the role of the KL-divergence term in a VAE's loss?
- A VAE encodes to a distribution (μ, logσ²) instead of a point, samples z via the reparameterization trick z = μ + σ·ε, and decodes — making the latent space continuous.
- Its loss = reconstruction + KL(N(μ,σ²)‖N(0,1)) (the negative ELBO): reconstruction enforces faithfulness, KL packs codes into a sampleable N(0,1) region so you can generate by decoding z ~ N(0,1).
- The reparameterization trick moves the randomness into an external ε ~ N(0,1) so gradients flow through μ and σ; the recon/KL tension is why VAE samples tend to be blurry.
VAEs generate and interpolate images, molecules, and audio, and compress data into structured latents; Stable Diffusion's first and last steps are literally a VAE encoder and decoder around the diffusion process.
If you remove it: You'd lack the bridge from 'autoencoder that compresses' to 'model that generates', and the ELBO (reconstruction + KL) that reappears throughout probabilistic generative modeling.