Cross-entropy and KL divergence

How the loss steers training

A model doesn't output chances directly. First it gives each word a raw score, called a logit. Softmax, from the softmax lesson, turns those scores into chances. Training then nudges each score so the loss goes down.

To nudge well, training needs each score's slope. The slope says how much the loss changes when that score goes up a little. Write $z_i$ for the score of word $i$. Write $L$ for the loss. When softmax feeds cross-entropy, the slope has a famously simple form:

The curly d, $\partial$, is said "partial". It means a slope with every other score held still.

In plain words: each score's slope is the predicted chance minus the target. The target is 1 for the right word and 0 for every other word.

Use the same prediction as before, with "great" as the right word:

• "great": 0.6 − 1 = −0.4. The slope is negative, so raising this score lowers the loss. Training pushes it up.
• "good": 0.25 − 0 = 0.25. The slope is positive, so training pushes this score down.
• "okay" and "bad": 0.1 and 0.05. They get pushed down a little.

Each push is scaled by a small step size, so scores move a little at a time. A wrong word the model gave 0.3 would have slope 0.3. The surer the model was of a wrong word, the harder that word's score is pushed down.

q = [0.6, 0.25, 0.1, 0.05]   # predicted: great, good, okay, bad
p = [1, 0, 0, 0]             # target: "great" was right
slopes = [q_i - p_i for q_i, p_i in zip(q, p)]
print([round(s, 2) for s in slopes])

It prints [-0.4, 0.25, 0.1, 0.05]. The slopes add up to 0, because both lists add up to 1. Every training step of a language model does this at every position in every piece of text.

KL divergence: the extra cost

Back to the town. Forecaster B averaged 1.204 nats of surprise. The best possible forecast averaged 0.693 nats. So B pays an extra 0.511 nats a day for using the wrong forecast.

That extra cost is the Kullback–Leibler divergence, or KL divergence for short. It's named after the two people who introduced it.

Here $H(p)$ is the entropy of $p$: the surprise of the best forecast. The double bar $\|$ only separates the two lists. The first list, $p$, is reality. The second, $q$, is the forecast.

In plain words: KL divergence is your cross-entropy minus the best one possible. It's the extra surprise you pay for forecasting $q$ when reality is $p$.

Check it with the town. Each outcome adds its chance times the log of a ratio:

• sun: 0.5 × ln(0.5 ÷ 0.9) ≈ 0.5 × (−0.588) ≈ −0.294
• rain: 0.5 × ln(0.5 ÷ 0.1) ≈ 0.5 × 1.609 ≈ 0.805

Together that's about 0.511 nats, the same as 1.204 − 0.693.

import math

def kl(p, q):
    return sum(p_i * math.log(p_i / q_i) for p_i, q_i in zip(p, q) if p_i > 0)

town = [0.5, 0.5]
forecast_b = [0.9, 0.1]
print(round(kl(town, forecast_b), 3), round(kl(forecast_b, town), 3))

It prints 0.511 0.368. The if p_i > 0 part skips outcomes that never happen, since they add nothing. The second number swaps the two lists, and the answer changes. Two facts are worth knowing:

• KL divergence is 0 exactly when $q = p$, and positive otherwise.
• It isn't symmetric. Swapping $p$ and $q$ usually gives a different number. So it isn't a true distance. It's still the standard way to measure how different two sets of chances are.

You'll meet it in two places:

• RLHF and preference tuning, where people's ratings tune a model. KL is a penalty that keeps the tuned model's outputs close to the original model's.
• Distillation, where a small "student" model learns from a large "teacher". KL is the loss that makes the student's predictions match the teacher's.

import numpy as np

def kl(p, q):
    p, q = np.asarray(p, float), np.asarray(q, float)
    mask = p > 0
    return np.sum(p[mask] * np.log(p[mask] / q[mask]))

teacher = [0.70, 0.20, 0.08, 0.02]
student_close = [0.65, 0.22, 0.10, 0.03]
student_far = [0.25, 0.25, 0.25, 0.25]
print("KL(teacher || close student):", round(kl(teacher, student_close), 4))
print("KL(teacher || uniform):      ", round(kl(teacher, student_far), 4))
print("KL(teacher || teacher):      ", kl(teacher, teacher))

The close student scores 0.0069, almost nothing. A student that gives every word 0.25 scores 0.5344. The teacher compared with itself scores exactly 0.

Perplexity: the loss as a number of choices

A loss like 1.5 nats is hard to picture. Perplexity turns it into something you can picture: a number of equally likely choices.

Take a fair six-sided die. Each face has a chance of 1/6. A model that always gives the right answer 1/6 has a loss of −ln(1/6) ≈ 1.792 nats per guess. Undo the log with exp: raise $e$, a fixed number about 2.718, to that power. You get 6 back. So perplexity 6 means the model is, on average, as unsure as someone guessing a roll of a fair die.

For a whole text, average the loss over every token first:

Here $N$ is the number of tokens in the text. $t$ counts through them, from 1 to $N$. $q(\text{token}_t)$ is the chance the model gave to the token that really came at position $t$.

In plain words: find the loss for every token and take the average. Then undo the log with exp. The result is how many equally likely choices the model is as unsure as.

Lower is better. A perfect model, always sure and always right, has perplexity 1. Phase 5 uses perplexity to compare language models.

Watch out: a loss in bits and a loss in nats aren't comparable. From nats you use exp, as here. From bits you raise 2 to the power instead.

import math

probs = [0.5, 0.2, 0.1]   # chance given to each correct token
losses = [-math.log(p) for p in probs]
avg = sum(losses) / len(losses)
print([round(x, 3) for x in losses])
print(round(avg, 3), round(math.exp(avg), 2))