KL Divergence

Learning Objectives

By the end of this section, you will be able to:

Where You'll Apply This: Training variational autoencoders, implementing knowledge distillation, understanding why cross-entropy loss works, policy optimization in RL, measuring model calibration, and comparing probability distributions in any context.

The Big Picture

How do you measure how "different" two probability distributions are? This fundamental question has a precise answer: Kullback-Leibler divergence (KL divergence), also called relative entropy. It quantifies the information lost when one distribution is used to approximate another.

Historical Context

KL divergence has a rich history connecting statistics, information theory, and physics.

Mathematical Definition

The Kullback-Leibler divergence from distribution Q to distribution P (also written as "the KL divergence of Q from P" or "D_KL(P || Q)") is defined as:

Intuitive Meaning

Let's unpack what the KL divergence formula actually measures:

Interactive: KL Divergence Explorer

Explore how KL divergence changes as you adjust two Gaussian distributions. Notice how the divergence increases when the distributions differ more in mean or variance.

Key Properties

KL divergence has several important mathematical properties that distinguish it from true distance metrics:

Asymmetry: Why Order Matters

The asymmetry of KL divergence is not a bug - it's a feature. The two directions measure fundamentally different things:

Interactive: Asymmetry Demonstration

Adjust the distributions to see how D_KL(P || Q) and D_KL(Q || P) differ. Notice that when distributions are very different, the difference can be substantial.

Relationship to Entropy and Cross-Entropy

KL divergence connects beautifully to the concepts we learned in previous sections. This relationship is fundamental to understanding why we use cross-entropy as a loss function:

Expanding the definitions:

Forward vs Reverse KL

The choice between minimizing D_KL(P || Q) (forward KL) vs D_KL(Q || P) (reverse KL) leads to fundamentally different behaviors. This distinction is crucial for understanding variational inference.

Interactive: Mode-Covering vs Mode-Seeking

Watch how a Gaussian approximation Q behaves when fitted to a bimodal target P using forward vs reverse KL. This demonstrates why variational inference can sometimes miss modes of the true posterior.

AI/ML Applications

KL divergence appears throughout modern machine learning. Understanding it deeply will make you a more effective practitioner.

KL in Variational Autoencoders

The Variational Autoencoder (VAE) objective function contains a KL divergence term that acts as a regularizer on the latent space:

The KL term pushes the approximate posterior q(z|x) toward the prior p(z) (typically a standard normal). This ensures:

• Smooth latent space: Similar inputs map to nearby latent representations
• Valid sampling: Random samples from p(z) produce meaningful outputs
• Interpolation: Walking through latent space produces coherent transitions

Interactive: VAE ELBO Visualization

Explore how the encoder's output distribution affects the KL penalty term. See how deviating from the prior increases the regularization cost.

Real-World Examples

Python Implementation

Let's implement KL divergence calculation in Python. We'll handle numerical edge cases carefully.

For P = [0.7, 0.3] and Q = [0.5, 0.5]: 0.7*log(0.7/0.5) + 0.3*log(0.3/0.5) ≈ 0.082 nats

1import numpy as np
2
3# KL divergence for discrete distributions
4def kl_divergence(p, q, epsilon=1e-10):
5    """
6    Compute KL divergence D_KL(P || Q).
7
8    Parameters:
9        p: Array-like, true distribution
10        q: Array-like, approximating distribution
11        epsilon: Small value to prevent log(0)
12
13    Returns:
14        KL divergence in nats (natural log units)
15    """
16    # Convert to numpy arrays
17    p = np.array(p, dtype=np.float64)
18    q = np.array(q, dtype=np.float64)
19
20    # Normalize to ensure valid probabilities
21    p = p / p.sum()
22    q = q / q.sum()
23
24    # Check for zero in q where p is nonzero
25    zero_q_nonzero_p = (q < epsilon) & (p > epsilon)
26    if zero_q_nonzero_p.any():
27        print("Warning: Q has zeros where P is nonzero - KL may be inaccurate")
28        q = np.clip(q, epsilon, 1.0)
29
30    # Compute KL: sum(p * log(p/q)) where p > 0
31    mask = p > epsilon
32    kl = np.sum(p[mask] * (np.log(p[mask]) - np.log(q[mask])))
33    return kl

1import numpy as np
2from scipy.stats import entropy as scipy_entropy
3
4def kl_divergence(p, q, epsilon=1e-10):
5    """Compute KL divergence D_KL(P || Q) in nats."""
6    p = np.array(p, dtype=np.float64)
7    q = np.array(q, dtype=np.float64)
8    p, q = p / p.sum(), q / q.sum()
9    q = np.clip(q, epsilon, 1.0)
10    mask = p > epsilon
11    return np.sum(p[mask] * (np.log(p[mask]) - np.log(q[mask])))
12
13def kl_gaussian(mu1, sigma1, mu2, sigma2):
14    """Closed-form KL for univariate Gaussians."""
15    return (np.log(sigma2 / sigma1) +
16            (sigma1**2 + (mu1 - mu2)**2) / (2 * sigma2**2) - 0.5)
17
18# ============================================
19# Example 1: Discrete distributions
20# ============================================
21print("=== Discrete KL Divergence ===")
22p = [0.4, 0.3, 0.2, 0.1]  # True distribution
23q = [0.25, 0.25, 0.25, 0.25]  # Uniform approximation
24
25kl_pq = kl_divergence(p, q)
26kl_qp = kl_divergence(q, p)
27print(f"D_KL(P || Q) = {kl_pq:.4f} nats")
28print(f"D_KL(Q || P) = {kl_qp:.4f} nats")
29print(f"Asymmetry: |difference| = {abs(kl_pq - kl_qp):.4f}")
30
31# Compare with SciPy
32print(f"\nSciPy verification: {scipy_entropy(p, q):.4f} nats")
33
34# ============================================
35# Example 2: Gaussian distributions
36# ============================================
37print("\n=== Gaussian KL Divergence ===")
38# Prior: N(0, 1), Posterior: N(1.5, 0.7)
39mu1, sigma1 = 1.5, 0.7  # Encoder output
40mu2, sigma2 = 0, 1       # Prior
41
42kl = kl_gaussian(mu1, sigma1, mu2, sigma2)
43print(f"VAE KL term: D_KL(N({mu1},{sigma1}²) || N({mu2},{sigma2}²)) = {kl:.4f} nats")
44
45# Closed-form VAE KL (for standard normal prior)
46vae_kl = -0.5 * (1 + 2*np.log(sigma1) - mu1**2 - sigma1**2)
47print(f"VAE formula: {vae_kl:.4f} nats")
48
49# ============================================
50# Example 3: Cross-entropy = Entropy + KL
51# ============================================
52print("\n=== Relationship Verification ===")
53p = [0.7, 0.2, 0.1]
54q = [0.5, 0.3, 0.2]
55
56# Compute each term
57H_p = -np.sum(np.array(p) * np.log(np.array(p)))
58H_pq = -np.sum(np.array(p) * np.log(np.array(q)))
59D_kl = kl_divergence(p, q)
60
61print(f"Entropy H(P) = {H_p:.4f} nats")
62print(f"Cross-Entropy H(P,Q) = {H_pq:.4f} nats")
63print(f"KL Divergence D_KL(P||Q) = {D_kl:.4f} nats")
64print(f"Verification: H(P,Q) - H(P) = {H_pq - H_p:.4f} nats")
65
66# ============================================
67# Example 4: Knowledge distillation loss
68# ============================================
69print("\n=== Knowledge Distillation ===")
70teacher_logits = np.array([2.0, 1.0, 0.5, -0.5])
71student_logits = np.array([1.5, 1.2, 0.3, -0.3])
72temperature = 2.0
73
74def softmax(x, T=1.0):
75    exp_x = np.exp(x / T)
76    return exp_x / exp_x.sum()
77
78teacher_soft = softmax(teacher_logits, temperature)
79student_soft = softmax(student_logits, temperature)
80
81distillation_loss = kl_divergence(teacher_soft, student_soft)
82print(f"Teacher probs: {np.round(teacher_soft, 3)}")
83print(f"Student probs: {np.round(student_soft, 3)}")
84print(f"Distillation loss (T={temperature}): {distillation_loss:.4f} nats")

Knowledge Check

Test your understanding of KL divergence with this interactive quiz.

Summary

Looking Ahead: In the next section, we'll explore mutual information, which uses KL divergence to measure how much information two random variables share. This is fundamental to representation learning and feature selection in ML.