Generative Models
VAEs - Learning Probabilistic Latent Representations
VAEs combine autoencoders with variational inference to learn smooth, continuous latent spaces. By encoding to distributions instead of points, they enable principled generation, meaningful interpolation, and disentangled representations - all with a stable training objective.
- Key point 1 - ELBO loss balances reconstruction quality with latent space regularization
- Key point 2 - Reparameterization trick enables backpropagation through stochastic sampling
- Key point 3 - Beta-VAE and VQ-VAE extend the framework for disentanglement and discrete latents
"In the latent space, every point tells a story."
Variational Autoencoders - Deep Dive
What You'll Learn
- 1Derive the evidence lower bound and identify what each of its two terms controls
- 2Implement the reparameterization trick and explain why direct sampling blocks gradient flow
- 3Tune the beta weight to trade reconstruction quality against disentanglement
- 4Compare VAEs with GANs and choose the right model for interpolation, likelihoods, or sharp samples
VAEs are generative models that learn a smooth latent space by combining autoencoders with variational inference. Unlike GANs, they provide a principled probabilistic framework for generation.
From Autoencoder to VAE
How this diagram works: This diagram shows the VAE architecture, which differs from a standard autoencoder by encoding inputs to a probability distribution rather than a fixed point. The encoder (green) maps input to parameters of a Gaussian distribution - a mean and variance . Instead of sampling directly (which is non-differentiable), the reparameterization trick samples where , making the process differentiable for backpropagation. The decoder (red) then reconstructs the input from this sampled latent vector. The loss combines reconstruction quality with a KL divergence term that regularizes the latent space toward a standard Gaussian prior, ensuring it remains smooth and continuous - enabling meaningful interpolation and generation by sampling from the prior.
Evidence Lower Bound (ELBO)
For Bernoulli pixel values the expectation collapses to a per-pixel binary cross entropy summed over all 784 input dimensions:
For a single latent dimension with and the bracket is , so that dimension alone contributes
Rearranging that identity isolates the bound itself, with the approximation gap written as a non-negative term:
Reparameterization Trick
Beta-VAE
A plain VAE is exactly the case, and Higgins et al. used to disentangle the factors on dSprites:
VQ-VAE (Vector Quantized VAE)
Quantization is a nearest-neighbour lookup in the codebook, and the straight-through estimator copies the decoder gradient back to the encoder:
PyTorch Implementation
VAE vs GAN
| Aspect | VAE | GAN |
|---|---|---|
| Training | Stable (single objective) | Unstable (minimax game) |
| Latent space | Smooth, continuous | May have gaps |
| Generation | Sample from prior | Sample from prior |
| Quality | Often blurry | Often sharper |
| Mode coverage | Good (covers all modes) | May miss modes |
| Likelihood | Can compute ELBO | Cannot compute |
| Disentanglement | Possible with -VAE | Requires architecture changes |
Practice Exercises
-
Train VAE on MNIST: Generate digits and visualize the latent space with t-SNE.
-
Beta-VAE experiment: Vary and observe the effect on disentanglement vs. reconstruction quality.
-
Conditional VAE: Add class labels to the encoder and decoder. Generate specific digits.
-
VQ-VAE: Implement vector quantization and compare with continuous VAE.
What to Learn Next
-> GANs Deep Dive Master generative adversarial networks with minimax training and Nash equilibrium.
-> Diffusion Models Deep Dive Explore the math behind gradual noising and denoising for image generation.
-> Self-Supervised Learning Learn useful representations from unlabeled data without manual annotation.
-> DL Systems Design Master distributed training, monitoring, and production deployment of deep learning models.
-> Model Compression Make deep learning models fast and efficient for production deployment.
-> Graph Neural Networks Learn from graph-structured data with message passing and attention mechanisms.
Common Mistakes
Sampling z directly from N(mu, sigma^2) inside the training loop
Fix: Draw epsilon ~ N(0, I) and set z = mu + sigma * epsilon so gradients reach mu and sigma
Letting reconstruction and KL use different normalizations
Fix: Normalize reconstruction per pixel and KL per latent dimension, otherwise beta changes meaning with input size
Reading blurry samples as a decoder bug
Fix: Blur is what an averaged likelihood encourages; move to a VAE-GAN loss or discrete VQ latents when you need sharpness
Best Practices
- Anneal beta from 0 to 1 over the first 20-40 epochs instead of starting at 1
- Keep at least 0.05 nats per latent dimension with free bits so the KL term cannot collapse
- Evaluate with ELBO on a held-out set, and draw generation samples from the prior rather than from encoded training data
Key Takeaways
- log p_theta(x) = ELBO + KL(q_phi(z|x) || p_theta(z|x)), so maximizing the bound maximizes likelihood while closing the gap to the true posterior
- The reparameterization z = mu + sigma * epsilon turns sampling into a differentiable function of mu and sigma
- The closed-form Gaussian KL is -0.5 * sum(1 + log sigma^2 - mu^2 - sigma^2), zero only when q matches N(0, I)
- beta weights the KL term: 1 is the standard VAE, above 1 buys disentanglement at some reconstruction cost, VQ-VAE removes continuous latents entirely
Frequently Asked Questions
Why are VAE samples blurrier than GAN samples?
The decoder is trained with a likelihood whose optimum for ambiguous inputs is the conditional mean, so a pixel that could be sharp in one of two configurations gets the average of both. MSE and binary cross entropy both have this averaging behavior, and the ELBO adds no term that rewards confident, high-frequency detail. When sharpness matters, replace the likelihood with an adversarial critic (VAE-GAN), quantize the latent (VQ-VAE), or let a diffusion head model the residual detail.
What actually stops the encoder from collapsing?
Two forces act at once: the reconstruction term wants to encode everything about , and the KL term pulls it back toward . If the decoder is powerful enough to ignore , the KL wins outright and drops to zero per dimension, which is posterior collapse. Log the per-dimension KL during training to catch it early, then fix it with KL annealing or free bits.
Do I have to use a Gaussian prior?
No. The framework only needs a prior you can sample from and an encoder whose posterior you can evaluate; closed-form KL is what makes the diagonal Gaussian convenient rather than required. VampPrior, mixtures, and normalizing flows all work, and they give the latent space structure that a single isotropic Gaussian cannot represent.
Interview Questions
Q1: Derive the KL term for a diagonal Gaussian posterior against a standard normal prior.
KL = 0.5 * sum_j (mu_j^2 + sigma_j^2 - 1 - log sigma_j^2). For mu = 0.5 and sigma^2 = 4 the per-dimension value is 0.932 nats, and it is zero only at mu = 0 and sigma^2 = 1, which is exactly the constraint the prior imposes.
Q2: Why does the reparameterization trick matter for training?
Sampling z ~ N(mu, sigma^2) places a non-differentiable random operation between encoder and decoder. Rewriting it as z = mu + sigma * epsilon with epsilon ~ N(0, 1) moves randomness off the parameter path, giving dL/dmu = dL/dz and dL/dsigma = (dL/dz) * epsilon, so standard backprop works.
Q3: How do you detect and fix posterior collapse?
Log KL per latent dimension each epoch; if it falls to near zero while reconstruction stays strong, the decoder is ignoring z. Anneal beta from 0 to 1 over 20-40 epochs, add free bits so each dimension keeps at least 0.05 nats, or weaken the decoder.