KL Divergence, Cross-Entropy
& Machine Learning

Key idea: an ideal code gives symbol x a codeword of length −log p(x). Your model of the world is your code.

Units: log base 2 gives bits. Natural log gives nats. 1 nat = 1.4427 bits. ML libraries use nats.

from math import log2

def entropy(p):
    return -sum(pi * log2(pi) for pi in p if pi > 0)

def cross_entropy(p, q):
    return -sum(pi * log2(qi) for pi, qi in zip(p, q) if pi > 0)

def kl(p, q):
    return sum(pi * log2(pi / qi) for pi, qi in zip(p, q) if pi > 0)

p = [0.5, 0.25, 0.125, 0.125]   # true source
q = [0.25, 0.25, 0.25, 0.25]    # model: uniform
print(f"H(p)   = {entropy(p):.3f} bits")
print(f"H(p,q) = {cross_entropy(p, q):.3f} bits")
print(f"KL     = {kl(p, q):.3f} bits")
assert abs(cross_entropy(p, q) - (entropy(p) + kl(p, q))) < 1e-12

H(p)   = 1.750 bits
H(p,q) = 2.000 bits
KL     = 0.250 bits

The if pi > 0 guard uses the convention 0 log 0 = 0. It is the limit of t log t as t → 0.

It is a divergence, not a distance. It fails symmetry and the triangle inequality. Its square root is not a metric either.

import random
from math import log2
def kl(p, q):
    return sum(pi * log2(pi / qi) for pi, qi in zip(p, q) if pi > 0)

def rand_dist(k):
    w = [random.random() for _ in range(k)]
    s = sum(w)
    return [x / s for x in w]

random.seed(0)
worst = min(kl(rand_dist(5), rand_dist(5)) for _ in range(100_000))
print(f"smallest KL over 100k random pairs: {worst:.5f}")

smallest KL over 100k random pairs: 0.00076

Corollary: H(p) ≤ log |X|. Take q uniform: KL(p‖u) = log|X| − H(p) ≥ 0.

The assert in the cross-entropy snippet checks this identity numerically. Try it yourself on a few random pairs.

True weather p and a forecaster's model q over sun, cloud, rain. All logs base 2.

1. Step 1. Code lengths: −log₂ q = 2, 2, 1 bits under the model.
2. Step 2. Weight each length by the true p: that is H(p,q).
3. Step 3. Check the identity: 1.75 = 1.5 + 0.25. It holds.

Here KL(q‖p) is also 0.25, because q is p with two labels swapped. That is luck. The next slide shows the usual case.

from math import log2
def kl(p, q):
    return sum(pi * log2(pi / qi) for pi, qi in zip(p, q) if pi > 0)

p = [0.9, 0.1]
q = [0.5, 0.5]
print(f"KL(p||q) = {kl(p, q):.4f}")
print(f"KL(q||p) = {kl(q, p):.4f}")
print(f"KL(p||p) = {kl(p, p):.4f}")

KL(p||q) = 0.5310
KL(q||p) = 0.7370
KL(p||p) = 0.0000

Rule of thumb: KL punishes the model hard where the first argument has mass but the second does not.

Forward KL is "zero-avoiding": q must be > 0 wherever p is. Reverse KL is "zero-forcing": q must be 0 wherever p is.

Check it: 10 coin flips, 7 heads. Scan θ on a grid.

from math import log
flips = "HHTHHHTHHT"                  # data: 7 heads, 3 tails
p_hat = flips.count("H") / len(flips) # empirical distribution

def nll(theta):                       # average negative log-likelihood
    return -sum(log(theta if c == "H" else 1 - theta) for c in flips) / len(flips)

def kl_emp(theta):                    # KL(p_hat || Bernoulli(theta))
    return (p_hat * log(p_hat / theta)
            + (1 - p_hat) * log((1 - p_hat) / (1 - theta)))

grid = [i / 100 for i in range(1, 100)]
print("argmin NLL :", min(grid, key=nll))
print("argmin KL  :", min(grid, key=kl_emp))
h = -(p_hat * log(p_hat) + (1 - p_hat) * log(1 - p_hat))
print(all(abs(nll(t) - (h + kl_emp(t))) < 1e-12 for t in grid))

argmin NLL : 0.7
argmin KL  : 0.7
True

The last line confirms NLL(θ) = H(p̂) + KL(p̂‖q_θ) for all 99 grid values. The closed form is θ̂ = 7/10.

As N → ∞, p̂ → p. So MLE is a sample estimate of "minimize KL from the truth". This is the forward, mode-covering direction.

The gradient entries sum to 0 because softmax only cares about differences of logits. Adding a constant to every z_i changes nothing.

Finite-difference check (central differences, h = 10⁻⁶):

from math import exp, log

def softmax(z):
    m = max(z)
    e = [exp(v - m) for v in z]
    s = sum(e)
    return [v / s for v in e]

def ce_loss(z, y):              # y = index of the true class
    return -log(softmax(z)[y])

def ce_grad(z, y):              # analytic: p - onehot(y)
    p = softmax(z)
    return [pi - (1 if i == y else 0) for i, pi in enumerate(p)]

z, y, h = [2.0, -1.0, 0.5], 0, 1e-6
numeric = []
for i in range(len(z)):
    zp = z[:]; zp[i] += h
    zm = z[:]; zm[i] -= h
    numeric.append((ce_loss(zp, y) - ce_loss(zm, y)) / (2 * h))
print("analytic:", [round(g, 6) for g in ce_grad(z, y)])
print("numeric :", [round(g, 6) for g in numeric])

analytic: [-0.214403, 0.039113, 0.17529]
numeric : [-0.214403, 0.039113, 0.17529]

import random
from math import exp, log

random.seed(1)
# 2-D points: class 1 if x + y > 1 (with a little label noise)
data = []
for _ in range(200):
    x, y = random.random(), random.random()
    label = 1 if x + y > 1 else 0
    if random.random() < 0.05:
        label = 1 - label
    data.append(((x, y), label))

def sigmoid(t):
    return 1 / (1 + exp(-t))

w1 = w2 = b = 0.0
lr = 1.0
for epoch in range(2001):
    g1 = g2 = gb = loss = 0.0
    for (x, y), t in data:
        p = sigmoid(w1 * x + w2 * y + b)
        loss -= t * log(p) + (1 - t) * log(1 - p)   # binary cross-entropy
        g1 += (p - t) * x; g2 += (p - t) * y; gb += p - t
    n = len(data)
    w1 -= lr * g1 / n; w2 -= lr * g2 / n; b -= lr * gb / n
    if epoch % 500 == 0:
        print(f"epoch {epoch:4d}  BCE = {loss / n:.4f}")

acc = sum((sigmoid(w1*x + w2*y + b) > 0.5) == t for (x, y), t in data) / len(data)
print(f"accuracy = {acc:.2f}")

epoch    0  BCE = 0.6931
epoch  500  BCE = 0.3724
epoch 1000  BCE = 0.3668
epoch 1500  BCE = 0.3661
epoch 2000  BCE = 0.3660
accuracy = 0.88

• Start: all weights 0, so p = 0.5 and loss = ln 2 = 0.6931 nats.
• The loss is convex in (w, b), so gradient descent finds the global min.
• The floor is not 0. We flipped 5% of labels on purpose, so some cost is H(p) of the noise.
• Accuracy is 0.88. Those 5% noise flips plus boundary points cap it.

MSE is still right for regression with Gaussian noise. That is MLE too, just with a different model.

from math import log
def ce(target, pred):
    return -sum(t * log(p) for t, p in zip(target, pred) if t > 0)

def smooth(k, y, eps):
    return [(1 - eps) + eps / k if i == y else eps / k for i in range(k)]

confident = [0.998, 0.001, 0.001]
calibrated = [0.90, 0.05, 0.05]
for eps in (0.0, 0.1):
    t = smooth(3, 0, eps)
    print(f"eps={eps}: loss(confident)={ce(t, confident):.3f}"
          f"  loss(calibrated)={ce(t, calibrated):.3f}")

eps=0.0: loss(confident)=0.002  loss(calibrated)=0.105
eps=0.1: loss(confident)=0.462  loss(calibrated)=0.298

With hard labels, the 99.8%-confident model wins. With ε = 0.1, the calmer 90% model wins. The optimum is now q = y^LS, not a spike.

from collections import Counter
from math import log2

train = "the cat sat on the mat . the dog sat on the log . the cat saw the dog .".split()
test = "the dog sat on the mat .".split()
vocab = set(train)
V = len(vocab)
bigrams = Counter(zip(train, train[1:]))
unigrams = Counter(train[:-1])

def prob(prev, word):                 # add-one (Laplace) smoothing
    return (bigrams[(prev, word)] + 1) / (unigrams[prev] + V)

bits = -sum(log2(prob(a, b)) for a, b in zip(test, test[1:]))
n = len(test) - 1
print(f"vocab={V}  cross-entropy={bits / n:.3f} bits/token")
print(f"perplexity={2 ** (bits / n):.2f}  (uniform guess would be {V})")

vocab=9  cross-entropy=2.293 bits/token
perplexity=4.90  (uniform guess would be 9)

Add-one smoothing keeps every bigram probability above 0. Without it, any unseen pair in a test sentence, such as ("cat", "log"), would get probability 0 and infinite perplexity.

Seeing the umbrella removes about half of our 0.93 bits of doubt about the weather.

from math import log2
# joint distribution of (weather, umbrella)
joint = {("rain", "yes"): 0.30, ("rain", "no"): 0.05,
         ("sun",  "yes"): 0.05, ("sun",  "no"): 0.60}
px, py = {}, {}
for (x, y), p in joint.items():
    px[x] = px.get(x, 0) + p
    py[y] = py.get(y, 0) + p

mi = sum(p * log2(p / (px[x] * py[y])) for (x, y), p in joint.items())
hx = -sum(p * log2(p) for p in px.values())
print(f"I(X;Y) = {mi:.4f} bits   H(X) = {hx:.4f} bits")

I(X;Y) = 0.4727 bits   H(X) = 0.9341 bits

Pinsker: TV(p,q) ≤ √( KL(p‖q) / 2 ) in nats. Small KL forces the distributions to be close.

from math import log2
def kl(p, q):
    return sum(pi * log2(pi / qi) for pi, qi in zip(p, q) if pi > 0)

def js(p, q):
    m = [(a + b) / 2 for a, b in zip(p, q)]
    return 0.5 * kl(p, m) + 0.5 * kl(q, m)

p = [1.0, 0.0]
q = [0.0, 1.0]
print(f"JS(p,q) = {js(p, q):.4f} bits   JS(q,p) = {js(q, p):.4f}")
try:
    kl(p, q)
except ZeroDivisionError:
    print("KL(p||q) = infinity (q is 0 where p is not)")

JS(p,q) = 1.0000 bits   JS(q,p) = 1.0000
KL(p||q) = infinity (q is 0 where p is not)

With the VAE prior N(0, 1) it becomes

KL = ½( σ² + μ² − 1 ) − log σ

For a d-dim diagonal Gaussian, sum this over the d coordinates (additivity).

import random
from math import log, pi, exp, sqrt

def kl_gauss(mu, sigma):          # KL( N(mu, sigma^2) || N(0, 1) )
    return 0.5 * (sigma**2 + mu**2 - 1) - log(sigma)

def logpdf(x, mu, s):
    return -0.5 * log(2 * pi * s * s) - (x - mu) ** 2 / (2 * s * s)

random.seed(0)
mu, sigma = 1.0, 0.5
xs = [random.gauss(mu, sigma) for _ in range(200_000)]
mc = sum(logpdf(x, mu, sigma) - logpdf(x, 0, 1) for x in xs) / len(xs)
print(f"closed form : {kl_gauss(mu, sigma):.4f} nats")
print(f"Monte Carlo : {mc:.4f} nats")

closed form : 0.8181 nats
Monte Carlo : 0.8206 nats

Monte Carlo averages log q(z) − log p(z) over samples from q. It agrees to 2 decimals with 200k samples. The closed form has no noise, which is why VAEs use it.

from math import exp, log

ref = {"helpful": 0.40, "verbose": 0.30, "flattery": 0.20, "rude": 0.10}
reward = {"helpful": 1.0, "verbose": 0.6, "flattery": 1.5, "rude": -2.0}

def optimal_policy(beta):           # argmax E[r] - beta * KL(pi || ref)
    w = {y: ref[y] * exp(reward[y] / beta) for y in ref}
    z = sum(w.values())
    return {y: v / z for y, v in w.items()}

for beta in (10.0, 1.0, 0.1):
    pi = optimal_policy(beta)
    kl = sum(p * log(p / ref[y]) for y, p in pi.items())
    print(f"beta={beta:<4}", {y: round(p, 2) for y, p in pi.items()},
          f"KL={kl:.3f}")

beta=10.0 {'helpful': 0.41, 'verbose': 0.3, 'flattery': 0.22, 'rude': 0.08} KL=0.004
beta=1.0  {'helpful': 0.43, 'verbose': 0.21, 'flattery': 0.35, 'rude': 0.01} KL=0.141
beta=0.1  {'helpful': 0.01, 'verbose': 0.0, 'flattery': 0.99, 'rude': 0.0} KL=1.528

The reward model overrates "flattery". With large β the policy barely moves. With β = 0.1 it puts 99% on flattery. That is reward hacking in four lines.

This is reverse KL (π first), so it is mode-seeking. PPO usually adds it as a per-token penalty.

from math import exp, log
def softmax(z, T=1.0):
    m = max(z)
    e = [exp((v - m) / T) for v in z]
    s = sum(e)
    return [v / s for v in e]

teacher_logits = [6.0, 2.5, 2.0, -1.0]      # cat, dog, fox, car
for T in (1, 4):
    print(f"T={T}:", [round(p, 3) for p in softmax(teacher_logits, T)])

def distill_loss(student, teacher, T):
    p, q = softmax(teacher, T), softmax(student, T)
    return T * T * sum(pi * log(pi / qi) for pi, qi in zip(p, q))

print(f"KD loss = {distill_loss([4.0, 1.0, 1.5, 0.0], teacher_logits, 4):.4f}")

T=1: [0.953, 0.029, 0.017, 0.001]
T=4: [0.511, 0.213, 0.188, 0.089]
KD loss = 0.4710

At T = 1 the three wrong classes share under 5%, so they barely teach anything. At T = 4, dog and fox get about 20% each while car gets 9%. The student learns what looks like a cat.

from math import exp, log

def logsumexp(z):
    m = max(z)
    return m + log(sum(exp(v - m) for v in z))

z = [1000.0, 999.0, 998.0]
try:
    naive = log(sum(exp(v) for v in z))
except OverflowError:
    naive = "OverflowError"
print("naive :", naive)
print("stable:", round(logsumexp(z), 6))
# log-softmax, then cross-entropy for class 0, never forming tiny probs
print("CE    :", round(logsumexp(z) - z[0], 6))

naive : OverflowError
stable: 1000.407606
CE    : 0.407606

Stay in log-space as long as you can. Exponentiate only at the very end, if ever.

1. A classifier gives the true class probability 0.25. What is the cross-entropy loss in bits? In nats?
2. True label p = (1, 0), model q = (0.9, 0.1). Find KL(p‖q) and KL(q‖p) in nats.
3. A language model has cross-entropy 3 bits per token. What is its perplexity?
4. Logits z = (2, 1, 0), true class 0. Find softmax, the loss in nats, and ∇_z of the loss.
5. Find KL( N(0, 2²) ‖ N(0, 1) ) in nats.
6. Your 3-class loss is stuck at 0.30 nats with ε = 0.1 smoothing. Is the model broken?

Try each one before you look to the right.

1. A model is a code. Codeword length = −log q(x).
2. Cross-entropy is the average length you pay with the wrong code.
3. KL is the extra cost. Gibbs says it is ≥ 0, and 0 only when you are right.
4. KL is asymmetric. Forward covers modes. Reverse picks one.
5. MLE = min KL from the empirical distribution. So the classifier loss is cross-entropy.
6. Softmax + CE gives the clean gradient q − y.
7. MI, JS, ELBO, distillation and RLHF are all KL with different arguments.