KL Divergence Intermediate

Kullback-Leibler (KL) divergence measures how one probability distribution differs from a reference distribution. It is one of the most widely used concepts from information theory in modern AI, appearing in VAEs, RLHF, knowledge distillation, and more.

The KL Divergence Formula

For discrete distributions P (true) and Q (approximate):

D_KL(P || Q) = SUM[ P(x) * log( P(x) / Q(x) ) ]  for all x

# Equivalently:
D_KL(P || Q) = H(P, Q) - H(P)
# where H(P,Q) is cross-entropy and H(P) is entropy of P

Computing KL Divergence in Python

import numpy as np
from scipy.stats import entropy

def kl_divergence(p, q):
    """Compute KL(P || Q) in bits."""
    p, q = np.array(p), np.array(q)
    mask = p > 0
    return np.sum(p[mask] * np.log2(p[mask] / q[mask]))

P = [0.4, 0.3, 0.2, 0.1]
Q = [0.25, 0.25, 0.25, 0.25]  # uniform

print(f"KL(P || Q) = {kl_divergence(P, Q):.4f} bits")
print(f"KL(Q || P) = {kl_divergence(Q, P):.4f} bits")
# Note: these are different! KL is asymmetric

# Using scipy
print(entropy(P, Q, base=2))  # KL(P || Q)

Forward vs. Reverse KL

The asymmetry of KL divergence gives rise to two distinct optimization behaviors:

AI Applications of KL Divergence

Variational Autoencoders (VAEs)

The VAE loss function includes a KL divergence term that regularizes the learned latent distribution to be close to a standard Gaussian:

VAE Loss = Reconstruction Loss + beta * D_KL(q(z|x) || p(z))

RLHF (Reinforcement Learning from Human Feedback)

In RLHF, a KL penalty prevents the fine-tuned model from diverging too far from the base model:

Reward = Human_Reward(response) - beta * D_KL(pi_new || pi_base)

Knowledge Distillation

When training a smaller "student" model to mimic a larger "teacher" model, KL divergence measures how well the student's output distribution matches the teacher's soft predictions.