ELBO objective and KL term

Overview

The Evidence Lower BOund (ELBO) is the central optimization objective in variational inference and the training loss for Variational Autoencoders (VAEs). It provides a tractable way to approximate intractable posterior distributions by maximizing a lower bound on the log-likelihood, while simultaneously including a KL divergence term that regularizes the approximate posterior.

Derivation from First Principles

Step 1: Start with the intractable log-likelihood

We want to maximize $\log p_\theta(x)$ for observed data $x$, where the generative model is: $p_\theta(x) = \int p_\theta(x|z) p(z) dz$

Why this is hard: For most choices of $p_\theta(x|z)$ (like a neural network decoder), this integral has no closed form and is exponentially expensive in the dimension of $z$.

Step 2: Introduce the approximate posterior

Since we can't compute $p_\theta(z|x)$ directly, we introduce an approximate posterior $q_\phi(z|x)$ (the encoder network in VAEs). This is our variational distribution—we'll optimize $\phi$ to make $q_\phi(z|x)$ close to $p_\theta(z|x)$.

Step 3: Derive the ELBO using Jensen's inequality

Start with the log-likelihood: $\log p_\theta(x) = \log \int p_\theta(x, z) dz$

Multiply and divide by $q_\phi(z|x)$: $= \log \int \frac{q_\phi(z|x)}{q_\phi(z|x)} p_\theta(x, z) dz = \log \mathbb{E}_{q_\phi(z|x)} \left[ \frac{p_\theta(x, z)}{q_\phi(z|x)} \right]$

Key insight: $\log(\cdot)$ is a concave function, so by Jensen's inequality (which for concave functions gives $f(\mathbb{E}[X]) \geq \mathbb{E}[f(X)]$): $\log \mathbb{E}_{q} [f(z)] \geq \mathbb{E}_{q} [\log f(z)]$

Applying this: $\log p_\theta(x) \geq \mathbb{E}_{q_\phi(z|x)} \left[ \log \frac{p_\theta(x, z)}{q_\phi(z|x)} \right]$

Why this step? Jensen's inequality lets us move the $\log$ inside the expectation at the cost of introducing an inequality. Because $\log$ is concave, the inequality points the right way ($\log$ of an average $\geq$ average of $\log$s), giving us a valid lower bound. This makes the expression computable via Monte Carlo sampling.

Step 4: Rewrite using conditional and prior

Using $p_\theta(x, z) = p_\theta(x|z) p(z)$: $\mathcal{L} = \mathbb{E}_{q_\phi(z|x)} [\log p_\theta(x|z) + \log p(z) - \log q_\phi(z|x)]$

Rearrange: $= \mathbb{E}_{q_\phi(z|x)} [\log p_\theta(x|z)] + \mathbb{E}_{q_\phi(z|x)} \left[ \log \frac{p(z)}{q_\phi(z|x)} \right]$

Why this form?

• The first term rewards the model for reconstructing $x$ from latent codes sampled from the encoder.
• The second term penalizes the encoder distribution for deviating from the prior $p(z)$, preventing it from collapsing or overfitting.

Understanding the KL Term

Why we need it:

1. Prevents posterior collapse: Without the KL term, the encoder could learn to map each $x$ to a deterministic, wildly different $z$, making the latent space discontinuous and unusable for generation.
2. Enables sampling: By keeping $q_\phi(z|x) \approx p(z)$, we can generate new samples by drawing $z \sim p(z)$ and decoding, even for $x$ we've never seen.
3. Regularization: Acts like a complexity penalty—simpler, more "prior-like" encodings are preferred.

Mathematical Properties

Key properties:

• $D_{KL}(q \| p) \geq 0$ always (Gibbs' inequality)
• $D_{KL}(q \| p) = 0$ iff $q = p$ almost everywhere
• Not symmetric: $D_{KL}(q \| p) \neq D_{KL}(p \| q)$ in general

Why this direction? We use $D_{KL}(q \| p)$ not $D_{KL}(p \| q)$ because:

• $q$ is the distribution we control (encoder)
• Expectations are taken w.r.t. $q$, which we can sample from
• This is called reverse KL or exclusive KL—it makes $q$ underestimate the support of $p$ (zero-forcing behavior)

The Gap: ELBO vs True Log-Likelihood

Step 5: Quantifying the gap

Start from Bayes' rule: $p_\theta(z|x) = \frac{p_\theta(x|z) p(z)}{p_\theta(x)}$

Taking logs: $\log p_\theta(x) = \log p_\theta(x|z) + \log p(z) - \log p_\theta(z|x)$

Take expectation over $q_\phi(z|x)$: $\log p_\theta(x) = \mathbb{E}_{q_\phi} [\log p_\theta(x|z)] + \mathbb{E}_{q_\phi} [\log p(z)] - \mathbb{E}_{q_\phi} [\log p_\theta(z|x)]$

Add and subtract $\mathbb{E}_{q_\phi} [\log q_\phi(z|x)]$: $= \underbrace{\mathbb{E}_{q_\phi} [\log p_\theta(x|z)] - D_{KL}(q_\phi(z|x) \| p(z))}_{\mathcal{L}(\theta, \phi; x)} + \underbrace{D_{KL}(q_\phi(z|x) \| p_\theta(z|x))}_{\text{Gap}}$

Why this matters:

• The gap between ELBO and true log-likelihood is exactly the KL divergence between our approximation and the true posterior
• As we maximize ELBO, we simultaneously improve the generative model ($\theta$) and the approximation quality ($\phi$)
• When $q_\phi(z|x) = p_\theta(z|x)$, ELBO equals log-likelihood

Worked Examples

ELBO: $\mathcal{L} = \mathbb{E}_{z \sim q_\phi} [-\frac{1}{2} \|x - \mu_\theta(z)\|^2] - D_{KL}(q_\phi(z|x) \| \mathcal{N}(0, I))$

Why this step? The reconstruction term becomes negative squared error because $\log \mathcal{N}(x; \mu, I) = -\frac{1}{2}\|x - \mu\|^2 + \text{const}$.

Closed-form KL for diagonal Gaussians: $D_{KL}(q_\phi \| p) = \frac{1}{2} \sum_{j=1}^{d} \left( \mu_j^2 + \sigma_j^2 - \log \sigma_j^2 -1 \right)$

Derivation of this formula: Starting from definition: $D_{KL} = \mathbb{E}_q [\log q(z) - \log p(z)]$ $= \mathbb{E}_q \left[ -\frac{1}{2} \sum_j \frac{(z_j - \mu_j)^2}{\sigma_j^2} - \frac{1}{2}\log \sigma_j^2 + \frac{1}{2} \sum_j z_j^2 \right] + \text{const}$

Using $\mathbb{E}_q[z_j^2] = \mu_j^2 + \sigma_j^2$ and $\mathbb{E}_q[(z_j - \mu_j)^2] = \sigma_j^2$: $= \frac{1}{2} \sum_j \left( -1 - \log \sigma_j^2 + \mu_j^2 + \sigma_j^2 \right)$

Why this matters: This closed form makes VAE training efficient—no Monte Carlo estimation needed for the KL term.

Step 1: Reconstruction term $\log p_\theta(x|z) = -\frac{1}{2} \|(0.2, 0.8) - (0.19, 0.79)\|^2 = -\frac{1}{2}(0.01^2 + 0.01^2) = -0.0001$

Why this step? We evaluate how well the decoder reconstructed the original $x$ from the sampled latent code.

Step 2: KL term (using $-\log \sigma_j^2 = -2\log\sigma_j$, so $-\log(0.8^2) \approx 0.446$ and $-\log(1.2^2) \approx -0.365$) $D_{KL} = \frac{1}{2} [(0.5^2 + 0.8^2 - \log(0.8^2) - 1) + ((-0.3)^2 + 1.2^2 - \log(1.2^2) - 1)]$ $= \frac{1}{2} [(0.25 + 0.64 + 0.446 - 1) + (0.09 + 1.44 - 0.365 - 1)]$ $= \frac{1}{2} [0.336 + 0.165] = 0.2505$

ELBO: $\mathcal{L} = -0.0001 - 0.2505 = -0.2506$

Why this step? We penalize the encoder for producing a distribution that deviates from the standard normal prior. Note the crucial detail: the KL formula uses $\log \sigma_j^2 = 2\log\sigma_j$, not $\log \sigma_j$—forgetting the factor of 2 is a common error.

What happens:

• $\beta < 1$: Prioritizes reconstruction, KL can grow large → better reconstructions, entangled latents
• $\beta = 1$: Standard ELBO, balanced trade-off
• $\beta > 1$: Prioritizes matching prior → more disentangled latents, worse reconstructions

Why this works: The KL term's strength controls the information bottleneck. Higher $\beta$ forces the encoder to use fewer, more independent dimensions of $z$, leading to disentanglement.

Common Mistakes and Misconceptions

The truth: The KL term isn't ad-hoc regularization; it arises naturally from the variational inference framework. It's the gap between your approximate posterior and the true posterior. Removing it doesn't give you a "less regularized" VAE—it breaks the mathematical justification for why you're optimizing this objective at all.

How to fix it: Remember the derivation: ELBO + KL gap = log-likelihood. The KL term is part of the bound, not an add-on.

The nuance: Maximizing ELBO w.r.t. $\theta$ does improve $\log p_\theta(x)$. But when we also optimize $\phi$, we're tightening the bound by reducing the KL gap. The log-likelihood only increases if the ELBO increase exceeds the gap reduction. In practice, we jointly optimize, so both happen together.

How to fix it: ELBO is a lower bound for fixed $\phi$. When $\phi$ changes, the bound itself moves. The guarantee is: $\mathcal{L} \leq \log p(x)$ always, but improving $\mathcal{L}$ improves the model.

The truth: KL is asymmetric. $D_{KL}(q \| p)$ (reverse KL, used in ELBO) is zero-forcing: it makes $q(z) \approx 0$ wherever $p(z) \approx 0$, leading to underestimation of the true posterior's support. $D_{KL}(p \| q)$ (forward KL) is zero-avoiding: it makes $q(z) > 0$ wherever $p(z) > 0$, leading to overestimation of support.

Why it matters for VAEs: Reverse KL can cause posterior collapse where the encoder ignores certain modes of the true posterior. Forward KL would make the approximate posterior overly broad and diffuse.

How to fix it: Remember the direction: we want expectations w.r.t. $q$ (what we can sample), so we use $D_{KL}(q \| p)$.

The truth: ELBO is a lower bound on log-probability, which can be negative (probabilities can be < 1). What matters is that ELBO increases during training, not its absolute value. A typical ELBO value might be -1000 or -10000 depending on data dimensionality.

How to fix it: Track ELBO improvement, not absolute value. Compare ELBO across architectures on the same data for meaningful comparisons.

Active Recall Practice

So you build a guesser machine that looks at picture and guesses what secret code created it. Now you play a game:

1. Show your guesser a picture → it guesses a code
2. Feed that code to the magic box → see if you get the same picture back
3. You get points if the picture comes out right

But there's a catch! Your guesser can't just make up any random codes—it has to use codes that are similar to the ones the magic box naturally uses (like numbers between 0 and 1, not crazy numbers like 999).

ELBO is your score in this game: You get points for recreating pictures correctly (reconstruction term), but you lose points if your guesser uses weird, un-natural codes (KL term). Maximizing ELBO means getting better at both guessing codes and recreating pictures while keeping your guesses reasonable.

For KL direction: "Queen Prefers Priors" → $D_{KL}(\mathbf{Q}_\phi \| \mathbf{P})$, we measure how far our Queen (encoder) is from the Prior.

Connections

• Variational Inference: ELBO is the fundamental objective in VI; VAEs are a deep learning application
• KL Divergence: Deep dive into properties, forward vs reverse, f-divergences
• VAE Architecture: How ELBO is implemented in neural networks with reparameterization trick
• Expectation Maximization: ELBO appears in EM algorithm for latent variable models
• Information Bottleneck: β-VAE connects to IB principle via weighted KL term
• Posterior Collapse: What happens when KL → 0 and the model ignores latent variables
• Importance Weighted Autoencoders: Tighter bounds than ELBO using multiple samples
• Normalizing Flows: Alternative to diagonal Gaussian encoder for more expressive $q_\phi$

Summary

The ELBO objective $\mathcal{L} = \mathbb{E}_{q_\phi(z|x)} [\log p_\theta(x|z)] - D_{KL}(q_\phi(z|x) \| p(z))$ is a tractable lower bound on the intractable log-likelihood $\log p_\theta(x)$. It decomposes into:

1. Reconstruction term: Reward accurate decoding
2. KL regularization: Keep encoder close to prior

The gap between ELBO and log-likelihood is exactly $D_{KL}(q_\phi(z|x) \| p_\theta(z|x))$—the approximation error of the variational posterior. Maximizing ELBO jointly improves the generative model and reduces the approximation gap. The KL term is not ad-hoc regularization but a fundamental consequence of the variational framework that enables sampling, prevents collapse, and provides a principled training objective for VAEs.

#flashcards/ai-ml

Concept Map

Hinglish (regional understanding)

Hinglish (regional understanding)