Variational Autoencoders

Table of Contents

1. Learning Objectives
2. Prerequisites
3. Notation
4. Core Intuition
5. The Generative Model
6. The ELBO for VAEs
7. The Encoder and Amortized Inference
8. Closed-Form KL Divergence
9. The Reparameterization Trick
10. The Complete VAE Objective
11. Posterior Collapse
12. β-VAE and Disentanglement
13. VQ-VAE and Discrete Latents
14. Connection to Diffusion Models
15. Worked Examples
16. Connection to the Broader Curriculum
17. Common Pitfalls and Misconceptions
18. Research Perspective
19. Summary of Takeaways
20. Exercises

Learning Objectives

After reading this chapter, you should be able to:

1. Specify the VAE generative model $p_\theta(\mathbf{x} \mid \mathbf{z}) p(\mathbf{z})$ and explain why $p_\theta(\mathbf{x})$ is intractable.
2. Derive the VAE objective as the ELBO from Variational Inference.
3. Prove the closed-form KL between diagonal Gaussian $q_\phi$ and standard normal prior.
4. Apply the reparameterization trick and explain why it enables Backpropagation through stochastic nodes.
5. Analyze posterior collapse and mitigation strategies.
6. State the $\beta$-VAE objective and VQ-VAE discrete latent formulation.

Prerequisites

• Variational Inference — ELBO, mean-field, reparameterization
• Multivariate Gaussian — KL between Gaussians
• Backpropagation — gradient computation

Notation

• $\mathbf{x}, \mathbf{z}$ — Observed data and latent variables
• $q_\phi(\mathbf{z} \mid \mathbf{x})$ — Encoder (approximate posterior)
• $p_\theta(\mathbf{x} \mid \mathbf{z})$ — Decoder (likelihood)
• $p(\mathbf{z})$ — Prior (typically $\mathcal{N}(\mathbf{0}, \mathbf{I})$)
• $\mathcal{L}_{\mathrm{ELBO}}$ — Variational lower bound

Core Intuition

A Variational Autoencoder (VAE) learns a generative model: sample latent code $\mathbf{z}$, decode to data $\mathbf{x}$. It simultaneously learns an encoder that maps $\mathbf{x}$ to a distribution over $\mathbf{z}$ — amortized approximate inference.

Training maximizes a lower bound on $\log p_\theta(\mathbf{x})$ (the ELBO) because the true marginal is intractable. The VAE unites deep learning with Bayesian inference: neural networks parameterize both $p_\theta(\mathbf{x} \mid \mathbf{z})$ (decoder) and $q_\phi(\mathbf{z} \mid \mathbf{x})$ (encoder).

VAEs underpin latent diffusion models (Stable Diffusion operates in a VAE latent space) and remain foundational for representation learning.

Series context. Chapter 9 (Generative Models), Volume II.

The Generative Model

Definition 1 (VAE Generative Process).

Definition 2 (Marginal Likelihood).

Proposition 1. For neural network decoder $p_\theta$, integral (2) is intractable — no closed form for $\log p_\theta(\mathbf{x})$ or its gradient.

The ELBO for VAEs

From Variational Inference, introduce $q_\phi(\mathbf{z} \mid \mathbf{x})$:

Theorem 1 (VAE ELBO).

Identity (ELBO Gap).

Maximizing (3) minimizes the gap — tightest when $q_\phi = p_\theta(\mathbf{z} \mid \mathbf{x})$.

The Encoder and Amortized Inference

Definition 3 (Amortized Encoder).

where $\boldsymbol{\mu}_\phi, \log \boldsymbol{\sigma}_\phi^2$ are neural network outputs.

Definition 4 (Amortization). One network $q_\phi$ serves all data points — vs. per-example variational optimization in classical VI.

Closed-Form KL Divergence

Theorem 2 (Gaussian KL). For $q = \mathcal{N}(\boldsymbol{\mu}, \text{diag}(\boldsymbol{\sigma}^2))$, $p = \mathcal{N}(\mathbf{0}, \mathbf{I})$:

Proof. Apply multivariate Gaussian KL formula with $\boldsymbol{\Sigma}_1 = \text{diag}(\boldsymbol{\sigma}^2)$, $\boldsymbol{\Sigma}_2 = \mathbf{I}$, $\boldsymbol{\mu}_2 = \mathbf{0}$:

Since $|\boldsymbol{\Sigma}_1| = \prod_j \sigma_j^2$: (6) follows. $\blacksquare$

The Reparameterization Trick

Problem: $\nabla_\phi \mathbb{E}_{q_\phi}[f(\mathbf{z})]$ — distribution depends on $\phi$.

Theorem 3 (Reparameterization).

Proof. See Variational Inference, Theorem 5. $\blacksquare$

Monte Carlo estimate with $L=1$ sample enables Backpropagation.

The Complete VAE Objective

Definition 5 (VAE Loss).

Decoder likelihoods:

• Bernoulli (images): reconstruction = binary cross-entropy
• Gaussian: reconstruction = MSE (equivalent to $\log p$ up to constants)

Posterior Collapse

Definition 6 (Posterior Collapse). Encoder learns $q_\phi(\mathbf{z} \mid \mathbf{x}) \approx p(\mathbf{z}) = \mathcal{N}(\mathbf{0}, \mathbf{I})$ — latents carry no information about $\mathbf{x}$.

Cause: Powerful autoregressive decoder models $p(\mathbf{x})$ without $\mathbf{z}$; KL term pushes $q \to p(\mathbf{z})$.

Mitigations:

• KL annealing: ramp $\beta$ from 0 to 1
• Free bits: $\max(D_{\text{KL}}, \lambda)$
• $\beta$-VAE with $\beta < 1$

β-VAE and Disentanglement

Definition 7 ($\beta$-VAE).

• $\beta = 1$: standard ELBO
• $\beta > 1$: stronger regularization, more disentangled latents
• $\beta < 1$: better reconstruction, risk of collapse

VQ-VAE and Discrete Latents

Definition 8 (Vector Quantization). Codebook $\{\mathbf{e}_k\}_{k=1}^K$. Encoder output $\mathbf{z}_e$ quantized:

Loss:

Foundation of latent diffusion (Forward Diffusion in VAE latent space).

Connection to Diffusion Models

Stable Diffusion: VQ-VAE encoder maps images to latent $\mathbf{z}$; diffusion process operates on $\mathbf{z}$ rather than pixels — reducing dimensionality and compute. See Reverse Process.

Worked Examples

Example 1: KL for $\boldsymbol{\mu} = \mathbf{0}$, $\boldsymbol{\sigma} = \mathbf{1}$

$D_{\text{KL}} = 0$ — encoder matches prior.

Example 2: Single-Dimensional

$\mu = 2$, $\sigma^2 = 0.5$: $D_{\text{KL}} = \frac{1}{2}(0.5 + 4 - 1 - \log 0.5) \approx 2.1$.

Connection to the Broader Curriculum

• ELBO derivation — Variational Inference
• Gradients — Backpropagation
• Diffusion — Forward Diffusion
• Low-rank latents — SVD

Common Pitfalls and Misconceptions

Pitfall 1: Confusing ELBO with true log-likelihood.

Pitfall 2: Ignoring posterior collapse with powerful decoders.

Pitfall 3: Using $L=1$ Monte Carlo without noting variance.

Pitfall 4: $\beta > 1$ invalidates ELBO as lower bound (still useful).

Research Perspective

VAE (Kingma & Welling, 2014; Rezende et al., 2014). $\beta$-VAE (Higgins et al., 2017). VQ-VAE (van den Oord et al., 2017). Modern: hierarchical VAEs, NVAE, connection to diffusion and flow matching.

Summary of Takeaways

• Encoder $q_\phi(\mathbf{z} \mid \mathbf{x})$ — Amortized posterior
• Decoder $p_\theta(\mathbf{x} \mid \mathbf{z})$ — Generative model
• KL term (6) — Regularize toward prior
• Reparameterization (8) — Differentiable sampling
• ELBO gap (4) — Approximation quality

Next: Forward Diffusion

Exercises

Exercise 1. Derive (6) from scratch.

Exercise 2. Prove (9) for scalar $z$.

Exercise 3. Explain posterior collapse with a autoregressive decoder.

Exercise 4. How does $\beta$-VAE affect the ELBO as a valid bound?

Exercise 5. Derive VQ-VAE gradient through straight-through estimator.

Exercise 6. Compare VAE latent dimension to LoRA rank.

Exercise 7. When is $L=1$ Monte Carlo sufficient?

Exercise 8. Connect VAE to Bayesian Linear Regression with linear decoder.