Unit 6.2 · about 4 h
208 / 324The ELBO and reparameterized sampling
The ELBO is the only training signal in the required VAE build. A sign error in the KL term or a blocked sampler fails silently into blurry reconstructions.
- derive
- implement
You can now
- Derive the ELBO from Jensen and split it into a reconstruction term and a KL term against the prior.
- Implement the one-sample VAE loss with separate reconstruction and KL terms that match the public tests.
- Implement reparameterized sampling so gradients flow through the encoder mean and variance.
Covers
- ELBO reconstruction and KL terms
- Reparameterized sampling and gradient flow
Study, in this order
2 steps
- ELBO derivation and the reconstruction-minus-KL formUnderstanding Deep Learning - generative models
Sections 17.3 to 17.6, with Equations 17.14 to 17.24 (pp. 331 to 338)
Depth · Read Section 17.6.1 for the one-sample estimate (Eq 17.23). Leave Sections 17.8 to 17.9 and the Notes for later units.
- The reparameterization trickUnderstanding Deep Learning - generative models
Section 17.7 with Equation 17.25 and Notebook 17.2 (pp. 339 to 340)
Focus on
- Jensen turns \(\log \mathbb{E}\) into \(\mathbb{E} \log\) and creates the bound (Eqs 17.14 to 17.16). Equation 17.17 shows the gap to the true log likelihood is exactly \(D_{\mathrm{KL}}[q \,\|\, \text{posterior}]\), so the bound is tight when \(q\) equals the posterior.
- Equation 17.18 rewrites the same bound as expected log likelihood minus \(D_{\mathrm{KL}}[q \,\|\, \text{prior}]\). Section 17.6 conditions \(q\) on \(x\) to get the VAE form (Eq 17.21) with encoder networks for mean and covariance.
- The one-sample estimate (Eq 17.23) pairs one decoded log likelihood with the closed-form Gaussian KL (Eq 17.24). Keep the two terms separate in code so collapse is visible.
- Reparameterization (Eq 17.25) moves randomness into \(\epsilon\) from \(\mathcal{N}(0, I)\). Backprop passes through \(\mu\) and \(\Sigma\) and never down the stochastic branch.
Reproduce closed-book
Closed-book, derive Equation 17.18 from Equation 17.16 in four lines, then compute the Deep-ML 393 check (0.5 reconstruction, 2.5 KL, 3.0 total) with NumPy.
Traps
- Coding
stdasexp(logvar)instead ofexp(0.5 logvar). The samples then have variance \(\sigma^4\) while the KL in Eq 17.24 assumes \(\sigma^2\). - Averaging the KL over the wrong axis. The closed form sums over latent dimensions first and then averages over the batch; swapping the two changes the loss scale with latent size.
- Comparing the Equation 17.30 objective with an ordinary ELBO. With \(\beta\) above 1 it is still a lower bound on \(\log \Pr(x)\), but a looser one, so its value is not comparable with the \(\beta = 1\) bound.
Practice
Deep-ML 393: Implement VAE Loss (ELBO)
Implement vae_loss(x, x_reconstructed, mu, log_var) returning (total, reconstruction, KL); the two public tests check (0.02, 0.02, 0.0) and (3.0, 0.5, 2.5)
Solve it changed-prompt: rename the arguments, swap the batch order, and re-derive the KL from Eq 17.24 before coding. The tests check the split, not just the total.
Self-test
-
State the ELBO in reconstruction-minus-KL form and say when it is tight.
\(\mathrm{ELBO} = \mathbb{E}_q[\log \Pr(x \mid z, \phi)] - D_{\mathrm{KL}}[q \,\|\, \text{prior}]\) (Eqs 17.18 and 17.21). It equals \(\log \Pr(x \mid \phi)\) minus \(D_{\mathrm{KL}}[q \,\|\, \text{posterior}]\) (Eq 17.17), so it is tight exactly when \(q\) equals the true posterior. -
Write the reparameterized sample and explain why gradients flow.
\(z^* = \mu + \Sigma^{1/2} \epsilon\) with \(\epsilon\) from \(\mathcal{N}(0, I)\) (Eq 17.25). Randomness enters through \(\epsilon\), which has no parameters, so gradients flow through \(\mu\) and \(\Sigma\) and skip the stochastic branch. -
For
x=[[1,0]],x_hat=[[0.5,0.5]],mu=[[1,2]]andlog_var=[[0,0]], what does Deep-ML 393 return?(3.0, 0.5, 2.5). Reconstruction 0.5 is the summed squared error 0.25 + 0.25. KL 2.5 is 0.5 times (1 + 4). Total is their sum. -
Why does the VAE use one sample per data point in Equation 17.23?
The expectation over \(q\) is estimated by Monte Carlo with \(N = 1\) while the KL stays closed-form. One sample per point is unbiased and keeps the per-step cost low; more samples only reduce estimator variance.
Completion criterion
Stop when
You derive Eq 17.18 without notes and your vae_loss passes both public tests with the two terms reported separately.
Leads to
- 6.3 Posterior collapse and decoder capacity Watch the KL term collapse to zero during training.
- 6.5 Diffusion: forward corruption, reverse denoising and prediction targets The diffusion loss is another ELBO with the same reconstruction-versus-prior shape.
Comes after
- 3.7 Entropy, cross-entropy, KL divergence and mutual informationThe ELBO trades a reconstruction term against a KL regularizer.
- 6.1 Autoencoders and latent-variable modelsTurn the prior and likelihood into the ELBO and a trainable sampler.
Resources for this unit
- Understanding Deep Learning - generative models Read Chapter 14 for evaluation limits, 15 for GANs, 17 for VAE and ELBO, and 18 for diffusion. Read Chapter 16 only for flow recognition. Use notebooks 17.1-17.3 and 18.1-18.4 as mathematical and visual practice.
Also in this section
- D2L - GAN and DCGAN Use gan.md and dcgan.md for one bounded adversarial exposure. The required implementation remains a small VAE with separate reconstruction and KL terms on a two-dimensional mixture or binarized MNIST subset.
- Deep-ML candidate - VAE, GAN and diffusion problems After the platform admission check, use at most one of 393 VAE Loss, 1384 Reparameterization, 174 1D GAN or 303 Forward Diffusion as a changed-prompt drill.
- Stanford CS231n - diffusion assignment slice Use the Assignment 3 diffusion component only after the VAE gate if diffusion mechanics remain weak. It is not a second required generative implementation.
- Progressive Distillation for Fast Sampling of Diffusion Models Read Section 4 only for the velocity parameterization \(v = \alpha \epsilon - \sigma x\) and why it stays stable at zero signal-to-noise ratio. Leave Algorithms 1 and 2 and the distillation results for later.