Variational Autoencoders: ELBO, Reparameterization, and KL Divergence

🎯 Learning Objectives

• Understand the motivation for variational inference in latent variable models
• Derive the ELBO (Evidence Lower Bound) from KL divergence
• Implement the reparameterization trick for differentiable sampling
• Compare VAE with AE and understand the regularization effect of KL

📋 Prerequisites

• Bayesian inference basics
• KL divergence and entropy
• Autoencoder architecture

---

1. 📖 Core Content

1.1 The Generative Modeling Problem

$$p(x) = \int p(x|z) p(z) dz$$

1.2 The ELBO Derivation

$$\log p(x) = \log \int p(x|z) p(z) dz$$ $$\log p(x) = D_{KL}(q_\phi(z|x) || p(z|x)) + \mathcal{L}(x; \theta, \phi)$$ $$\log p(x) \geq \mathcal{L}(x; \theta, \phi) = \mathbb{E}_{q_\phi(z|x)}[\log p_\theta(x|z)] - D_{KL}(q_\phi(z|x) || p(z))$$

• $\mathbb{E}_{q_\phi(z|x)}[\log p_\theta(x|z)]$: Reconstruction loss (log-likelihood of data given latent)
• $D_{KL}(q_\phi(z|x) || p(z))$: KL regularization (encourages latent distribution to match prior)

1.3 The Reparameterization Trick

$$z = \mu + \sigma \odot \epsilon, \quad \epsilon \sim \mathcal{N}(0, I)$$

python
Run Inline
# runnable
import torch
import torch.nn as nn
import torch.nn.functional as F
class VAE(nn.Module):
    def __init__(self, input_dim=784, latent_dim=20, hidden_dim=400):
        super().__init__()
        # Encoder
        self.encoder = nn.Sequential(
            nn.Linear(input_dim, hidden_dim),
            nn.ReLU()
        )
        self.mu_layer = nn.Linear(hidden_dim, latent_dim)
        self.logvar_layer = nn.Linear(hidden_dim, latent_dim)
        # Decoder
        self.decoder = nn.Sequential(
            nn.Linear(latent_dim, hidden_dim),
            nn.ReLU(),
            nn.Linear(hidden_dim, input_dim),
            nn.Sigmoid()
        )
    def encode(self, x):
        h = self.encoder(x)
        return self.mu_layer(h), self.logvar_layer(h)
    def reparameterize(self, mu, logvar):
        """Reparameterization trick: z = μ + σ * ε"""
        if self.training:
            std = torch.exp(0.5 * logvar)
            eps = torch.randn_like(std)
            return mu + eps * std
        return mu  # Deterministic at inference
    def decode(self, z):
        return self.decoder(z)
    def forward(self, x):
        mu, logvar = self.encode(x)
        z = self.reparameterize(mu, logvar)
        recon = self.decode(z)
        return recon, mu, logvar
def vae_loss(recon, x, mu, logvar):
    """VAE loss = reconstruction + KL divergence"""
    # Reconstruction loss (binary cross-entropy)
    recon_loss = F.binary_cross_entropy(recon, x, reduction='sum')
    # KL divergence: KL(q(z|x) || p(z)) where p(z) = N(0, I)
    # Formula: -0.5 * sum(1 + log(σ²) - μ² - σ²)
    kl_loss = -0.5 * torch.sum(1 + logvar - mu.pow(2) - logvar.exp())
    return recon_loss + kl_loss

1.4 Interpreting the KL Term

1. Posterior collapse prevention: Prevents the encoder from mapping different inputs to distant latent codes
2. Latent space regularization: Encourages the latent space to be smooth and continuous

$$KL(\mathcal{N}(\mu, \sigma^2) || \mathcal{N}(0, 1)) = -\frac{1}{2}(1 + \log \sigma^2 - \mu^2 - \sigma^2)$$

1.5 Beta-VAE

$$\mathcal{L}_{\beta-VAE} = \mathbb{E}[\log p_\theta(x|z)] - \beta \cdot KL(q_\phi(z|x) || p(z))$$

---

📝 Practice Questions

Q1

<strong>Q1</strong>: Show that the ELBO is a lower bound on the log-likelihood.

Start with log-likelihood: $\log p(x) = \log \int p(x,z) dz$ $= \log \int q(z|x) \frac{p(x,z)}{q(z|x)} dz$ $\geq \int q(z|x) \log \frac{p(x,z)}{q(z|x)} dz$ (by Jensen's inequality: log(E[X]) ≥ E[log(X)]) $= \int q(z|x) \log \frac{p(x|z)p(z)}{q(z|x)} dz$ $= \int q(z|x) \log p(x|z) dz + \int q(z|x) \log \frac{p(z)}{q(z|x)} dz$ $= \mathbb{E}_{q(z|x)}[\log p(x|z)] - KL(q(z|x) || p(z))$

This lower bound is the ELBO. Maximizing it improves the model's log-likelihood. Q2

<strong>Q2

<strong>Q2</strong>: In the reparameterization trick, why can't we directly sample z ~ N(μ, σ²) and backpropagate?

Direct sampling: z = sample(N(μ, σ²))

The gradient ∂z/∂μ and ∂z/∂σ doesn't exist — sampling is a stochastic operation with no defined gradient. The "sampling node" breaks the computation graph.

Reparameterization: z = μ + σ·ε, ε ~ N(0,1)

Now z is a deterministic function of μ, σ, and ε. Gradients ∂z/∂μ = 1 and ∂z/∂σ = ε exist and can be computed. The randomness comes from ε (external), which doesn't need gradients.

This seemingly simple trick is what makes VAEs trainable with SGD. Q3

<strong>Q3

<strong>Q3

<strong>Q3

<strong>Q3</strong>: What is posterior collapse in VAEs, and how can it be detected?

Posterior collapse: The KL term dominates the loss, and the encoder learns to map all inputs to the prior distribution: $q(z|x) \approx p(z) = N(0, I)$.

Symptoms:

1. Ignored latent code: The decoder ignores z (produces same output regardless of input)
2. KL divergence approaches 0: μ → 0, σ² → 1 for all inputs
3. Reconstruction is blurry/mean: Decoder averages over all possibilities
4. No latent structure: Interpolating in latent space doesn't change output

Causes: Powerful decoder (e.g., expressive autoregressive), weak latent, high β value.

Fixes: Reduce β, use KL annealing, use a weaker decoder, increase latent dimensionality. Q4

<strong>Q4

<strong>Q4

<strong>Q4

<strong>Q4

<strong>Q4</strong>: Compare VAEs and GANs on sample quality, latent space structure, and training stability.

VAEs excel at structured latent spaces (useful for interpolation, semi-supervised learning) while GANs excel at sample quality. Diffusion models (covered in Week 4-5) combine benefits of both.

</details> * * * ## 🔗 Cross-References - **Next**: [Diffusion Models](/notes/04-degree-electives-bsda5002-genai-foundations-week04-04-diffusion-models) - **Previous**: [GANs](/notes/04-degree-electives-bsda5002-genai-foundations-week02-02-gans) - **Video**: BSDA5002 Week 4-5 transcripts [Join Discord](https://discord.gg/gE2m4Qrdqv) [Previous**GANs**](/notes/04-degree-electives-bsda5002-genai-foundations-week02-02-gans)[Next**Diffusion Models (DDPM)**](/notes/04-degree-electives-bsda5002-genai-foundations-week04-04-diffusion-models)