Variational Autoencoders (VAE) theory

#ai-ml/generative-models #unsupervised-learning #probabilistic-models #latent-variables

Overview

Variational Autoencoders (VAEs) are generative models that learn to encode data into a latent space and decode it back, while ensuring the latent space follows a known distribution (typically Gaussian). Unlike standard autoencoders, VAEs are probabilistic and enable sampling of new data.

Core Intuition

Why do we need the probabilistic formulation?

Deterministic autoencoders can "memorize" and create disconnected regions in latent space
We want to generate new samples by sampling from latent space
Probabilistic interpretation gives us a principled way to balance reconstruction and regularization

Mathematical Foundation

The Generative Story

Why Gaussian prior? It's mathematically convenient, has infinite support, and the central limit theorem suggests many natural processes converge to it.

The Problem: Intractable Marginal Likelihood

We want to maximize: $p_\theta(\mathbf{x}) = \int p_\theta(\mathbf{x}|\mathbf{z}) p(\mathbf{z}) d\mathbf{z}$

Why is this hard? For each data point $\mathbf{x}$ , we'd need to integrate over ALL possible latent codes $\mathbf{z}$ to see which ones could have generated it. With high-dimensional $\mathbf{z}$ , this integral is computationally intractable.

The Solution: Variational Inference

Introduce an inference network (encoder) $q_\phi(\mathbf{z}|\mathbf{x})$ that approximates the true posterior $p_\theta(\mathbf{z}|\mathbf{x})$ .

Deriving the ELBO from scratch:

Start with log-likelihood: $\log p_\theta(\mathbf{x}) = \log \int p_\theta(\mathbf{x}|\mathbf{z}) p(\mathbf{z}) d\mathbf{z}$

Multiply inside by $\frac{q_\phi(\mathbf{z}|\mathbf{x})}{q_\phi(\mathbf{z}|\mathbf{x})} = 1$ : $\log p_\theta(\mathbf{x}) = \log \int q_\phi(\mathbf{z}|\mathbf{x}) \frac{p_\theta(\mathbf{x}|\mathbf{z}) p(\mathbf{z})}{q_\phi(\mathbf{z}|\mathbf{x})} d\mathbf{z}$

This is expectation over $q_\phi$ : $\log p_\theta(\mathbf{x}) = \log \mathbb{E}_{q_\phi(\mathbf{z}|\mathbf{x})} \left[ \frac{p_\theta(\mathbf{x}|\mathbf{z}) p(\mathbf{z})}{q_\phi(\mathbf{z}|\mathbf{x})} \right]$

Why this step? Jensen's inequality states $\log \mathbb{E}[X] \geq \mathbb{E}[\log X]$ for any random variable $X$ (log is concave). Apply it:

$\log p_\theta(\mathbf{x}) \geq \mathbb{E}_{q_\phi(\mathbf{z}|\mathbf{x})} \left[ \log \frac{p_\theta(\mathbf{x}|\mathbf{z}) p(\mathbf{z})}{q_\phi(\mathbf{z}|\mathbf{x})} \right]$

Split the logarithm: $\log p_\theta(\mathbf{x}) \geq \mathbb{E}_{q_\phi(\mathbf{z}|\mathbf{x})} \left[ \log p_\theta(\mathbf{x}|\mathbf{z}) \right] + \mathbb{E}_{q_\phi(\mathbf{z}|\mathbf{x})} \left[ \log \frac{p(\mathbf{z})}{q_\phi(\mathbf{z}|\mathbf{x})} \right]$

The second term is negative KL divergence: $\mathbb{E}_{q_\phi(\mathbf{z}|\mathbf{x})} \left[ \log \frac{p(\mathbf{z})}{q_\phi(\mathbf{z}|\mathbf{x})} \right] = -D_{KL}(q_\phi(\mathbf{z}|\mathbf{x}) \| p(\mathbf{z}))$

Why "lower bound"? The inequality from Jensen guarantees $\log p_\theta(\mathbf{x}) \geq \mathcal{L}$ . Maximizing the ELBO pushes up a lower bound on the log-likelihood.

The Gap: KL Between Approximate and True Posterior

We can show: $\log p_\theta(\mathbf{x}) = \mathcal{L}(\theta, \phi; \mathbf{x}) + D_{KL}(q_\phi(\mathbf{z}|\mathbf{x}) \| p_\theta(\mathbf{z}|\mathbf{x}))$

Why? Start from: $D_{KL}(q_\phi \| p_\theta(\mathbf{z}|\mathbf{x})) = \mathbb{E}_{q_\phi} \left[ \log \frac{q_\phi(\mathbf{z}|\mathbf{x})}{p_\theta(\mathbf{z}|\mathbf{x})} \right]$

Use Bayes rule: $p_\theta(\mathbf{z}|\mathbf{x}) = \frac{p_\theta(\mathbf{x}|\mathbf{z}) p(\mathbf{z})}{p_\theta(\mathbf{x})}$ :

$D_{KL}(q_\phi \| p_\theta(\mathbf{z}|\mathbf{x})) = \mathbb{E}_{q_\phi} \left[ \log q_\phi(\mathbf{z}|\mathbf{x}) - \log p_\theta(\mathbf{x}|\mathbf{z}) - \log p(\mathbf{z}) + \log p_\theta(\mathbf{x}) \right]$

Since $\log p_\theta(\mathbf{x})$ doesn't depend on $\mathbf{z}$ : $D_{KL}(q_\phi \| p_\theta(\mathbf{z}|\mathbf{x})) = \log p_\theta(\mathbf{x}) - \mathcal{L}(\theta, \phi; \mathbf{x})$

Since KL divergence is always $\geq 0$ , this confirms the ELBO is a lower bound.

The Reparameterization Trick

Problem: We need to compute gradients through $\mathbb{E}_{q_\phi(\mathbf{z}|\mathbf{x})}[\cdot]$ , but sampling is not differentiable.

Assume $q_\phi(\mathbf{z}|\mathbf{x}) = \mathcal{N}(\boldsymbol{\mu}_\phi(\mathbf{x}), \boldsymbol{\sigma}_\phi^2(\mathbf{x}))$ .

Standard sampling (non-differentiable): $\mathbf{z} \sim \mathcal{N}(\boldsymbol{\mu}_\phi(\mathbf{x}), \boldsymbol{\sigma}_\phi^2(\mathbf{x}))$

Intuition: Think of it like separating the "decision" (mean and variance, which depend on parameters) from the "randomness" (epsilon, which doesn't). Now gradients can flow through the decision part.

Practical VAE Architecture

Why predict log variance? Ensures positivity ( $\sigma^2 = \exp(\log \sigma^2) > 0$ ) and numerical stability.

The Loss Function in Practice

For a dataset $\mathcal{D} = \{\mathbf{x}^{(i)}\}_{i=1}^N$ :

$\mathcal{L}(\theta, \phi) = \frac{1}{N} \sum_{i=1}^N \left[ \mathbb{E}_{q_\phi(\mathbf{z}|\mathbf{x}^{(i)})} [\log p_\theta(\mathbf{x}^{(i)}|\mathbf{z})] - D_{KL}(q_\phi(\mathbf{z}|\mathbf{x}^{(i)}) \| p(\mathbf{z})) \right]$

Monte Carlo estimate (typically 1 sample per data point): $\mathcal{L} \approx \frac{1}{N} \sum_{i=1}^N \left[ \log p_\theta(\mathbf{x}^{(i)}|\mathbf{z}^{(i)}) - D_{KL}(q_\phi(\mathbf{z}|\mathbf{x}^{(i)}) \| \mathcal{N}(\mathbf{0}, \mathbf{I})) \right]$

where $\mathbf{z}^{(i)} = \boldsymbol{\mu}_\phi(\mathbf{x}^{(i)}) + \boldsymbol{\sigma}_\phi(\mathbf{x}^{(i)}) \odot \boldsymbol{\epsilon}^{(i)}$ .

KL Divergence for Gaussians (Closed Form)

For $q = \mathcal{N}(\boldsymbol{\mu}, \boldsymbol{\sigma}^2\mathbf{I})$ and $p = \mathcal{N}(\mathbf{0}, \mathbf{I})$ :

Derivation: $D_{KL}(q \| p) = \mathbb{E}_q[\log q(\mathbf{z})] - \mathbb{E}_q[\log p(\mathbf{z})]$

For multivariate Gaussian (dimension $d$ ): $\log q(\mathbf{z}) = -\frac{d}{2}\log(2\pi) - \sum_{j=1}^d \log \sigma_j - \frac{1}{2}\sum_{j=1}^d \frac{(z_j - \mu_j)^2}{\sigma_j^2}$

$\log p(\mathbf{z}) = -\frac{d}{2}\log(2\pi) - \frac{1}{2}\sum_{j=1}^d z_j^2$

Taking expectations over $q$ (using $\mathbb{E}_q[z_j] = \mu_j$ , $\mathbb{E}_q[z_j^2] = \mu_j^2 + \sigma_j^2$ ):

$\mathbb{E}_q[\log q(\mathbf{z})] = -\frac{d}{2}\log(2\pi) - \sum_{j=1}^d \log \sigma_j - \frac{d}{2}$

$\mathbb{E}_q[\log p(\mathbf{z})] = -\frac{d}{2}\log(2\pi) - \frac{1}{2}\sum_{j=1}^d (\mu_j^2 + \sigma_j^2)$

Subtract: $D_{KL}(q \| p) = -\sum_{j=1}^d \log \sigma_j - \frac{d}{2} + \frac{1}{2}\sum_{j=1}^d (\mu_j^2 + \sigma_j^2)$

$= \frac{1}{2}\sum_{j=1}^d \left( \mu_j^2 + \sigma_j^2 - \log \sigma_j^2 - 1 \right)$

Worked Examples

Worked example Example 1: Computing ELBO for MNIST

Setup: Binary MNIST images ( $28 \times 28$ ), latent dimension $d=2$ .

Encoder outputs: $\boldsymbol{\mu} = [0.5, -0.3]$ , $\log \boldsymbol{\sigma}^2 = [-1.0, -0.5]$
Sample: $\boldsymbol{\epsilon} = [0.2, -0.4]$
Decoder outputs Bernoulli logits for each pixel

Step 1: Compute $\mathbf{z}$ via reparameterization $\sigma_1 = \exp(-1.0/2) = \exp(-0.5) \approx 0.606$ $\sigma_2 = \exp(-0.5/2) = \exp(-0.25) \approx 0.779$ $z_1 = 0.5 + 0.606\times 0.2 = 0.621$ $z_2 = -0.3 + 0.779 \times (-0.4) = -0.612$

Why these steps? We're applying the reparameterization trick to get a differentiable sample.

Step 2: Compute KL term $D_{KL} = \frac{1}{2}[(0.5^2 + 0.606^2 - (-1.0) - 1) + (-0.3)^2 + 0.779^2 - (-0.5) - 1)]$ $= \frac{1}{2}[(0.25 + 0.367 + 1.0 - 1) + (0.09 + 0.607 + 0.5 - 1)]$ $= \frac{1}{2}[0.617 + 0.197] = 0.407$

Why? This measures how much the encoder distribution deviates from standard normal.

Step 3: Compute reconstruction term Decoder produces logits $\ell_{ij}$ for each pixel. If true pixel is 1: $\log p_\theta(x_{ij}=1|\mathbf{z}) = \log \sigma(\ell_{ij}) = -\log(1 + e^{-\ell_{ij}})$

Sum over all 784 pixels to get total reconstruction log-likelihood.

Step 4: ELBO = Reconstruction - KL Suppose reconstruction term = -150 (typical for MNIST): $\mathcal{L} = -150 - 0.407 = -150.407$

Why negative? We're computing log-likelihood (log of probabilities < 1). Higher (less negative) is better.

Worked example Example 2: Understanding the KL Penalty

Scenario: Encoder learns $\boldsymbol{\mu} = [5, 5]$ , $\boldsymbol{\sigma}^2 = [0.01, 0.01]$ (tight cluster far from origin). Compute KL: $D_{KL} = \frac{1}{2}[(25 + 0.01 - \log 0.01 - 1) + (25 + 0.01 - \log 0.01 - 1)]$ $= \frac{1}{2}[(25 + 0.01 + 4.61 - 1) + (25 + 0.01 + 4.61 - 1)]$ $= \frac{1}{2}[28.62 + 28] = 28.62$

Why is this bad? Huge KL penalty! The encoder is placing codes far from the prior, creating gaps in latent space.

Better solution: Encoder learns $\boldsymbol{\mu} = [0.1, 0.1]$ , $\boldsymbol{\sigma}^2 = [1.0, 1.0]$ : $D_{KL} = \frac{1}{2}[(0.01 + 1 - 0 - 1) + (0.01 + 1 - 0 - 1)]$ $= \frac{1}{2}[0.01 + 0.01] = 0.01$

Why this step? The KL term forces the latent space to be organized around the origin with unit variance, enabling sampling.

Common Mistakes

Common mistake Mistake 1: Forgetting the Reparameterization Trick

Wrong approach:

z = np.random.normal(mu, sigma)  # Non-differentiable!
reconstruction = decoder(z)
loss.backward()  # Gradients don't flow to encoder

Why it feels right: "I'm sampling from the distribution the encoder outputs."

Why it's wrong: The sampling operation breaks the gradient flow. $\frac{\partial \mathbf{z}}{\partial \phi} = 0$ because sampling is a stochastic operation.

Correct approach:

epsilon = torch.randn_like(sigma)
z = mu + sigma * epsilon  # Differentiable!
reconstruction = decoder(z)
loss.backward()  # Gradients flow through mu and sigma

The fix: Randomness in $\epsilon$ (independent of parameters), deterministic transformation with $\mu$ and $\sigma$ (parameters we can differentiate).

Common mistake Mistake 2: Using Standard Autoencoder Loss Only

Wrong approach:

loss = reconstruction_loss(decoder(encoder(x)), x)

Why it feels right: "I just want good reconstructions!"

Why it's wrong: Without KL regularization:

Encoder can map different inputs to far-apart, arbitrary codes
Latent space has holes and discontinuities
Can't generate new samples (sampling from prior won't decode to realistic data)
Overfitting: memorizes training data

Correct approach:

reconstruction_loss = ...
kl_loss = 0.5 * torch.sum(mu**2 + sigma**2 - torch.log(sigma**2) - 1)
loss = reconstruction_loss + kl_loss

The fix: KL term forces encoder to use the prior's structure, enabling generation.

Common mistake Mistake 3: KL Vanishing (Posterior Collapse)

Problem: During training, KL term goes to zero, encoder ignores input, decoder learns to generate "average" images from prior.

Why it happens:

Decoder is too powerful (large capacity)
Can model $p(\mathbf{x})$ directly from prior without neding information in $\mathbf{z}$
Encoder takes shortcut: outputs prior distribution $\mathcal{N}(\mathbf{0}, \mathbf{I})$ for all inputs Why it feels right initially: "Low KL means good match to prior!"

Why it's wrong: Zero mutual information between $\mathbf{x}$ and $\mathbf{z}$ . The latent codes are useless.

Fixes:

KL anealing: Start with $\beta \ll 1$ in $\beta \cdot D_{KL}$ , gradually increase to 1
Free bits: Only penalize KL if it's below a threshold (allow some bits of information)
Weaker decoder: Limit capacity to force reliance on $\mathbf{z}$

Intuition: We want to find the sweet spot where decoder needs $\mathbf{z}$ but $\mathbf{z}$ is still close to prior.

Active Recall Checks

Recall Explain VAEs to a 12-Year-Old

Imagine you have a magic box that can compress photos of cats into secret codes (just a few numbers), and another magic box that can turn those codes back into cat photos.

A regular magic box might be mesy—some codes are far apart, some codes don't work at all. If you try a random code, you might get garbage.

A VAE is a smarter magic box that:

Makes sure all the secret codes follow a pattern (like a bell curve)
Makes sure nearby codes give similar cats
Has no empty spots—every code you try gives you a cat!

The trick: While learning to compress and uncompress photos, the boxes also learn to organize codes nicely. Now you can:

Give any random code → get a new cat photo that looks real
Take two cat photos, mix their codes, get a cat that's in-between The "variational" part means we're using probability and smart math to make sure the organization happens automatically!

Memory Aids

Reparameterization: "Separate the random from the learnable"

Random: $\epsilon$ (standard normal, no parameters)
Learnable: $\mu$ and $\sigma$ (encoder outputs)

KL Divergence Terms: "My Variance Logically Matters"

Mean squared: $\mu^2$ (penalty for mean away from zero)
Variance: $\sigma^2$ (penalty for too much spread)
Log variance: $-\log \sigma^2$ (penalty for too little spread)
Minus one: $-1$ (normalization)

Connections

Autoencoders: VAEs extend autoencoders with probabilistic formulation
Kullback-Leibler Divergence: Core to measuring distribution mismatch
Bayesian Inference: VAEs use variational inference for approximate posteriors
Latent Variable Models: VAEs are latent variable models with amortized inference
Evidence Lower Bound (ELBO): Fundamental objective function
Reparameterization Trick: Enables gradient-based learning
Generative Adversarial Networks (GANs): Alternative generative model approach
Normalizing Flows: Another probabilistic generative model family
Posterior Collapse: Common training pathology in VAEs
Disentangled Representations: Beta-VAE and other variants for interpretability
Conditional VAE: Extension for conditional generation
Importance-Weighted Autoencoders: Tighter ELBO bounds

#flashcards/ai-ml

What is the fundamental difference between autoencoders and VAEs? :: VAEs are probabilistic (encode/decode distributions not points), have a regularized latent space (KL term enforces prior), and enable generation by sampling from the prior distribution.

What does ELBO stand for and what are its two components?

Evidence Lower Bound. Components: (1) Reconstruction term - expected log-likelihood of data given latent code, (2) KL divergence term - regularization that keeps posterior close to prior.

Why do we need the reparameterization trick?

Standard sampling from encoder distribution is non-differentiable, blocking gradient flow to encoder parameters. Reparameterization separates random component (ε) from learnable parameters (μ, σ), making the sampling process differentiable.

Write the reparameterization trick formula :: z = μ_φ(x) + σ_φ(x) ⊙ ε, where ε ~ N(0, I). The randomness is in ε (independent of φ), allowing gradients to flow through μ and σ.

What is posterior collapse and why does it happen? :: When KL term goes to zero, encoder ignores input and outputs the prior distribution for all inputs. Happens when decoder is too powerful and can model p(x) from prior alone without needing information in z.

Write the closed-form KL divergence between diagonal Gaussian and standard normal

D_KL(N(μ, σ²I) || N(0, I)) = (1/2) Σⱼ (μⱼ² + σ - log σⱼ² - 1)

Why is the ELBO called a "lower bound"?

Due to Jensen's inequality (log E[X] ≥ E[log X]), the ELBO is guaranteed to be ≤ log p_θ(x). Maximizing ELBO pushes up a lower bound on the data log-likelihood.

What does the μ² term in the KL divergence penalize?

Means that are far from zero (the prior mean). Forces encoder to organize codes around the origin of latent space.

What does the -log σ² term in the KL divergence encourage?

Encourages having enough variance (penalty for too little spread). Balances with σ² term which penalizes too much variance.

How do you generate new samples from a trained VAE?

(1) Sample z ~ N(0, I) from prior, (2) Pass z through decoder network p_θ(x|z), (3) Sample x from the output distribution (e.g., Bernoulli for binary images).

What is the relationship between ELBO and the KL to true posterior?

log p_θ(x) = ELBO + D_KL(q_φ(z|x) || p_θ(z|x). Since KL ≥ 0, ELBO ≤ log p_θ(x). Tighter ELBO means better posterior approximation.

Why do we typically predict log σ² instead of σ² directly?

Ensures positivity (σ² = exp(log σ²) > 0 always) and numerical stability (avoids very small values during optimization).

What is KL annealing and why is it used?

Gradually increasing the KL weight from β < 1 to β = 1 during training. Helps prevent posterior collapse by allowing reconstruction to be learned first before enforcing prior constraint.

In VAE objective, what happens if you remove the KL term entirely?

Model becomes standard autoencoder: latent space can be arbitrary/discontinuous, codes far apart, holes in space, cannot generate by sampling from prior, tends to overfit.

What is the generative story of a VAE?

(1) Sample latent code z ~ p(z) = N(0, I), (2) Generate data x ~ p_θ(x|z) via decoder network. Goal:

Concept Map

Hinglish (regional understanding)

Intuition Hinglish mein samjho

Hinglish (regional understanding)

Intuition Hinglish mein samjho

Dekho, VAE ki asli baat samajhne ke liye ek image ko "meaning space" mein compress karne ka socho — jahan similar images paas-paas hoti hain. Simple autoencoder is space mein gaps aur weird clusters bana deta hai, jisse tum us space se random point uthakar valid image nahi bana sakte. VAE isko solve karta hai do tareeke se: pehla, woh compressed codes ko ek normal (Gaussian) distribution follow karne pe majboor karta hai taaki space smooth ho aur gaps na rahein; doosra, model ko probabilistic banata hai — matlab single point nahi, balki ek distribution output karta hai. Isse tum latent space se koi bhi point sample karke usse ek realistic naya image generate kar sakte ho. Yahi generative power hai jo VAE ko special banati hai.

Ab maths side pe — hum chahte hain ki $p_\theta(\mathbf{x})$ maximize ho, lekin uske liye humein saare possible latent codes $\mathbf{z}$ pe integrate karna padta hai, jo high dimensions mein computationally impossible hai (intractable). Isliye ek smart trick lagate hain: ek encoder network $q_\phi(\mathbf{z}|\mathbf{x})$ introduce karte hain jo true posterior ko approximate karta hai. Fir Jensen's inequality use karke hum ek lower bound nikalte hain jise ELBO kehte hain. ELBO ke do parts hote hain — pehla reconstruction quality (image kitni achhi wapas banti hai) aur doosra KL divergence (latent space kitna organized hai, prior ke kitna kareeb hai). In dono ke beech balance banana hi VAE training ka core hai.

Yeh cheez isliye important hai kyunki VAE tumhe ek principled, mathematically sound tareeka deta hai naya data generate karne ka — na ki sirf memorize karke copy paste karna. Real world mein iska use hota hai image generation, anomaly detection, aur data compression jaise applications mein. Aur sabse badi baat, ELBO derivation samajhne se tumhe variational inference ka foundation clear ho jaata hai, jo modern deep generative models — jaise diffusion models tak — mein baar-baar aata hai. Toh yeh concept ek building block hai jispe aage bahut kuch khada hai.

Overview

Core Intuition

Mathematical Foundation

The Generative Story

The Problem: Intractable Marginal Likelihood

The Solution: Variational Inference

The Gap: KL Between Approximate and True Posterior

The Reparameterization Trick

Practical VAE Architecture

The Loss Function in Practice

KL Divergence for Gaussians (Closed Form)

Worked Examples

Common Mistakes

Active Recall Checks

Memory Aids

Connections

Concept Map

Hinglish (regional understanding)

Hinglish (regional understanding)

Go deeper — visual, from zero

Test yourself — Generative Models

Connections