Lesson overview | Previous part | Lesson overview

KL Divergence: Appendix K: Worked Problems with Full Solutions to Appendix M: Further Reading and References

Appendix K: Worked Problems with Full Solutions

K.1 Computing KL Between Two Categorical Distributions

Problem. A language model assigns the following probabilities to the next token in "The capital of France is ___":

Compute (a) $D_{\mathrm{KL}}(p\|q)$, (b) cross-entropy $H(p,q)$, (c) entropy $H(p)$.

Solution.

(a) $D_{\mathrm{KL}}(p\|q)$:

$$= 0.92\ln\frac{0.92}{0.70} + 0.03\ln\frac{0.03}{0.10} + 0.02\ln\frac{0.02}{0.10} + 0.02\ln\frac{0.02}{0.05} + 0.01\ln\frac{0.01}{0.05}$$ $$= 0.92(0.274) + 0.03(-1.204) + 0.02(-1.609) + 0.02(-0.916) + 0.01(-1.609)$$ $$= 0.252 - 0.036 - 0.032 - 0.018 - 0.016 = 0.150 \text{ nats}$$

(b) Cross-entropy $H(p,q) = -\sum p\log q$:

$$= -(0.92\ln 0.70 + 0.03\ln 0.10 + 0.02\ln 0.10 + 0.02\ln 0.05 + 0.01\ln 0.05)$$ $$= -(0.92(-0.357) + 0.05(-2.303) + 0.02(-2.303) + 0.03(-2.996))$$ $$= 0.328 + 0.115 + 0.046 + 0.090 = 0.579 \text{ nats}$$

(c) Entropy $H(p) = -\sum p\log p$:

$$= -(0.92\ln 0.92 + 0.03\ln 0.03 + 0.02\ln 0.02 + 0.02\ln 0.02 + 0.01\ln 0.01)$$ $$= -(0.92(-0.083) + 0.03(-3.507) + 0.04(-3.912) + 0.01(-4.605))$$ $$= 0.076 + 0.105 + 0.156 + 0.046 = 0.383 \text{ nats}$$

Verification: $H(p,q) = H(p) + D_{\mathrm{KL}}(p\|q) \stackrel{?}{=} 0.383 + 0.150 = 0.533$ nats. (Small discrepancy from rounding; exact values match.)

Interpretation: The model wastes 0.150 nats per prediction compared to an oracle. The perplexity of the oracle is $e^{0.383} \approx 1.47$; the model's perplexity is $e^{0.579} \approx 1.78$ - 21% higher, entirely due to the KL gap.

K.2 The ELBO Calculation for a Toy VAE

Setup. One-dimensional VAE: prior $p(z) = \mathcal{N}(0,1)$, likelihood $p_\theta(x|z) = \mathcal{N}(z, 0.25)$ (decoder with $\sigma^2 = 0.25$), encoder $q_\phi(z|x) = \mathcal{N}(\mu_\phi(x), \sigma_\phi^2)$.

For $x = 1.5$, $\mu_\phi = 1.2$, $\sigma_\phi^2 = 0.4$:

KL term: $D_{\mathrm{KL}}(\mathcal{N}(1.2, 0.4)\|\mathcal{N}(0,1))$:

$$= \frac{1}{2}(1.2^2 + 0.4 - \ln 0.4 - 1) = \frac{1}{2}(1.44 + 0.4 + 0.916 - 1) = \frac{1}{2}(1.756) = 0.878 \text{ nats}$$

Reconstruction term: $\mathbb{E}_{q_\phi}[\log p_\theta(x|z)]$ where $x = 1.5$, $z \sim \mathcal{N}(1.2, 0.4)$:

$$= \mathbb{E}_{z\sim\mathcal{N}(1.2,0.4)}\left[-\frac{1}{2}\ln(2\pi \cdot 0.25) - \frac{(1.5-z)^2}{2 \cdot 0.25}\right]$$ $$= -\frac{1}{2}\ln(0.5\pi) - \frac{1}{0.5}\mathbb{E}[(1.5-z)^2]$$ $$= -\frac{1}{2}(0.452) - 2[(1.5-1.2)^2 + 0.4]$$ $$= -0.226 - 2[0.09 + 0.4] = -0.226 - 0.980 = -1.206 \text{ nats}$$

ELBO: $\mathcal{L} = -1.206 - 0.878 = -2.084$ nats.

Interpretation: The true log-evidence satisfies $\log p_\theta(1.5) \ge -2.084$. Better encoder parameters (closer $\mu_\phi$ to posterior mean) and tighter $\sigma_\phi^2$ would close the ELBO gap.

K.3 Verifying Data Processing Inequality Numerically

Consider input $X \in \{0,1,2\}$ with $p = (0.5, 0.3, 0.2)$ and $q = (0.2, 0.3, 0.5)$.

A stochastic function $T: \{0,1,2\} \to \{A, B\}$ with transition matrix:

$$T(A|0) = 0.8, T(B|0) = 0.2; \quad T(A|1) = 0.5, T(B|1) = 0.5; \quad T(A|2) = 0.1, T(B|2) = 0.9$$

Original KL:

$$D_{\mathrm{KL}}(p\|q) = 0.5\ln(2.5) + 0.3\ln(1) + 0.2\ln(0.4) = 0.5(0.916) + 0 + 0.2(-0.916) = 0.458 - 0.183 = 0.275 \text{ nats}$$

Induced distributions on $\{A,B\}$:

$$p_T(A) = 0.5(0.8) + 0.3(0.5) + 0.2(0.1) = 0.4 + 0.15 + 0.02 = 0.57$$ $$q_T(A) = 0.2(0.8) + 0.3(0.5) + 0.5(0.1) = 0.16 + 0.15 + 0.05 = 0.36$$

KL after processing:

$$D_{\mathrm{KL}}(p_T\|q_T) = 0.57\ln\frac{0.57}{0.36} + 0.43\ln\frac{0.43}{0.64} = 0.57(0.458) + 0.43(-0.399) = 0.261 - 0.172 = 0.089 \text{ nats}$$

Verification: $0.089 \le 0.275$ PASS - the data processing inequality holds. The stochastic transformation $T$ reduced the KL from 0.275 to 0.089 nats, losing information about whether data came from $p$ or $q$.

---

Appendix L: Connections to Optimization Theory

L.1 KL Divergence and Mirror Descent

Mirror descent is a generalization of gradient descent where the update uses a Bregman divergence instead of Euclidean distance:

$$\mathbf{x}^{(t+1)} = \arg\min_{\mathbf{x}} \left\{\eta \langle \nabla f(\mathbf{x}^{(t)}), \mathbf{x} \rangle + D_\phi(\mathbf{x}, \mathbf{x}^{(t)})\right\}$$

When $\phi(\mathbf{p}) = \sum_k p_k \ln p_k$ (negative entropy) and $\mathbf{x} = \mathbf{p}$ is a probability vector, the Bregman divergence $D_\phi(\mathbf{p}, \mathbf{q}) = D_{\mathrm{KL}}(p\|q)$ and mirror descent becomes the multiplicative weights update or exponentiated gradient:

$$p_k^{(t+1)} \propto p_k^{(t)} \exp(-\eta g_k^{(t)})$$

where $\mathbf{g}^{(t)} = \nabla f(\mathbf{p}^{(t)})$. This is exactly how softmax policies in RL are updated, and how online learning over the probability simplex works (regret analysis via Bregman projections).

Natural gradient as mirror descent. The natural gradient descent update:

$$\boldsymbol{\theta}^{(t+1)} = \arg\min_{\boldsymbol{\theta}} \left\{\eta \langle \nabla_{\boldsymbol{\theta}}\mathcal{L}, \boldsymbol{\theta} \rangle + D_{\mathrm{KL}}(p_{\boldsymbol{\theta}^{(t)}}\|p_{\boldsymbol{\theta}})\right\}$$

is mirror descent using KL divergence as the Bregman proximity term. The solution is $\boldsymbol{\theta}^{(t+1)} = \boldsymbol{\theta}^{(t)} - \eta F^{-1}\nabla\mathcal{L}$ where $F$ is the Fisher matrix - reproducing the natural gradient formula.

L.2 KL-Penalized Optimization and Soft Constraints

The RLHF objective $\max_\pi \mathbb{E}[r] - \beta D_{\mathrm{KL}}(\pi\|\pi_{\mathrm{ref}})$ is a classic penalized optimization problem. The KL term acts as a regularizer that keeps $\pi$ near $\pi_{\mathrm{ref}}$.

Lagrangian form. The hard-constrained version:

$$\max_\pi \mathbb{E}[r(x,y)] \quad \text{s.t.} \quad D_{\mathrm{KL}}(\pi\|\pi_{\mathrm{ref}}) \le \epsilon$$

is equivalent to the penalized form at the $\beta$ that is the Lagrange multiplier for the constraint. As $\epsilon \to \infty$, $\beta \to 0$ (no constraint); as $\epsilon \to 0$, $\beta \to \infty$ (stay at reference). The optimal $\beta$ depends on the trade-off between reward and divergence.

TRPO (Schulman et al., 2015). Trust Region Policy Optimization solves:

$$\max_\pi \mathcal{L}_{\pi_{\mathrm{old}}}(\pi) \quad \text{s.t.} \quad \bar{D}_{\mathrm{KL}}(\pi_{\mathrm{old}}, \pi) \le \delta$$

where $\mathcal{L}$ is the surrogate advantage and $\bar{D}_{\mathrm{KL}}$ is the maximum KL over states. Each TRPO step solves a constrained optimization problem with a KL trust region - computationally expensive. PPO approximates this with the clip objective.

L.3 The Frank-Wolfe Algorithm and KL

The Frank-Wolfe algorithm solves $\min_{p \in \mathcal{C}} f(p)$ (for a convex function $f$ on a convex constraint set $\mathcal{C}$) by linearizing $f$ at each step. When $f = D_{\mathrm{KL}}(\cdot\|q)$ and $\mathcal{C}$ is the probability simplex, Frank-Wolfe recovers the multiplicative weights algorithm. When $f = D_{\mathrm{KL}}(p\|\cdot)$ and $\mathcal{C}$ is a parametric exponential family, Frank-Wolfe corresponds to the M-step of EM.

This connection shows that common ML algorithms - EM, multiplicative weights, mirror descent, natural gradient - are all special cases of convex optimization using KL divergence as the geometry.

---

Appendix M: Further Reading and References

M.1 Primary Sources

M.2 Key ML Papers Using KL Divergence

M.3 Numerical Tools

---

M.4 Quick Diagnostic Checklist

Before using KL divergence in any ML project, verify:

• Which direction? Identify which distribution is $p$ (truth/reference) and which is $q$ (approximation). The first argument $p$ is where the expectation is taken.
• Support check. Does $p(x) > 0$ wherever $q(x) = 0$? If so, $D_{\mathrm{KL}}(p\|q) = +\infty$. Add smoothing if needed.
• Units. Is your implementation using natural log (nats) or log base 2 (bits)? Be consistent.
• Forward or reverse? Is this an MLE/distillation problem (forward KL) or variational inference (reverse KL)? The choice determines whether the result is mode-seeking or mean-seeking.
• Numerical stability. Are you computing $\log(p/q)$ directly (unstable) or using log-softmax (stable)?
• Sample size for estimation. KL estimators from finite samples have high variance for large $d$ - consider using MMD or Wasserstein for high-dimensional distributions.
• KL coefficient. If using KL as a regularizer, is $\beta$ tuned? Too small -> reward hacking; too large -> no alignment gain.
• ELBO monitoring. In VAE training, are both the reconstruction term and KL term logged separately? Monitoring their ratio over training diagnoses posterior collapse early.

---

M.5 Conceptual Summary: One Page

KL divergence $D_{\mathrm{KL}}(p\|q) = \mathbb{E}_p[\log(p/q)]$ measures:

• The expected excess code length when coding $p$-data with a $q$-code
• The average log-likelihood ratio evidence that data came from $p$ not $q$
• The information gain from using $p$ instead of $q$ as a model

Three things to always remember:

1. Direction matters. $D_{\mathrm{KL}}(p\|q) \ne D_{\mathrm{KL}}(q\|p)$. Forward KL -> mass-covering. Reverse KL -> mode-seeking.

2. Non-negativity is everything. $D_{\mathrm{KL}} \ge 0$ with equality iff $p = q$. This single fact proves: entropy is maximized at uniform, the ELBO is a lower bound, MLE is consistent.

3. It's everywhere in AI. MLE, VAE, RLHF, distillation, diffusion, VI - all are special cases of optimizing a KL divergence in a specific direction. Understanding which direction, and what approximations are made, is the key to understanding modern deep learning at a principled level.

---

Section 09-02 KL Divergence - 1966+ lines, 13 main sections, 13 appendices (A-M) Prerequisites: Section 09-01 Entropy | Follows to: Section 09-03 Mutual Information