The Math That Makes It Work: KL Divergence and Temperature, Demystified

Every distillation pipeline, from Hinton's 2015 paper to whatever taught the newest "mini" model on the price list, comes down to one line of math:

L = α · CE(student, hard labels)  +  (1 − α) · T² · KL(teacher_soft ‖ student_soft)

One loss, two teachers. The first term is the answer key — the ground-truth labels from your dataset. The second is the mentor — the big model's softened probability distributions. α decides how the student splits its study hours between them.

By the end of this post you'll understand every symbol in that line, including the T² that most tutorials silently drop — and then wonder why their distillation underperforms. And you'll have the whole thing in about fifteen lines of PyTorch that you can paste into any training loop.

Quick recap: Part 2 showed what transfers (dark knowledge — the structure in the teacher's doubt) and introduced temperature. Part 3 showed where to extract it (outputs, hidden layers, or geometry). Today is response-based distillation's engine room.

Term 1: Cross-entropy — the answer key

You already know this one; it's the standard classification loss. For each example, look at the probability the student assigned to the correct class, take the negative log, average over the batch:

Confidently right → tiny loss. Confidently wrong → enormous loss.

Why keep it at all, if the teacher is so wise? Two reasons. The teacher is sometimes wrong, and without the answer key the student faithfully inherits every mistake. And hard labels provide a crisp, unambiguous signal that anchors training when the teacher's soft distribution is mushy. The mentor is brilliant; the textbook keeps everyone honest.

Term 2: KL divergence — the surprise meter

KL divergence measures how different two probability distributions are, but the useful intuition is surprise:

KL(P ‖ Q) = how surprised you'd be, on average, if you expected distribution Q but reality behaved like P.

If the teacher says (90% cat, 7% dog, 3% fox) and the student says (88% cat, 9% dog, 3% fox) — low surprise, tiny KL, well done. If the student says (60% cat, 1% dog, 39% truck) — the teacher is shocked. Not mainly because "cat" dropped, but because the student put mass on truck, which the teacher considers absurd, and starved dog, which the teacher considers plausible. KL punishes the student for getting the shape of the doubt wrong, not just the top answer. That's precisely what makes it the right tool for transferring dark knowledge — a plain "did you match the argmax" loss would throw the dark knowledge away.

One detail worth noticing: KL is asymmetric. KL(teacher ‖ student) is not KL(student ‖ teacher). We use the direction where the teacher is the reference — "how surprised is the teacher by the student" — which pushes the student to cover everything the teacher considers plausible.

Both distributions inside the KL term are computed at temperature T (logits divided by T before softmax — Part 2's volume dial), so the wrong-answer structure is actually audible while the student learns.

The T² factor: the part everyone skips

Here's the good story about gradients I promised.

When you divide logits by T inside a softmax, you don't just soften the probabilities — you also shrink the gradients flowing back through the loss. Work through the calculus (or trust me and check the appendix of Hinton's paper) and the gradient magnitude of the soft-label term scales as 1/T². At T = 4, the mentor's teaching signal arrives at roughly 1/16th strength.

Now look at the combined loss again. The hard-label term's gradients don't shrink — there's no temperature on them. So without correction, raising T quietly turns down the mentor and turns up the answer key. You think you're tuning "how much dark knowledge to transfer"; you're actually tuning two entangled things at once, and your carefully chosen α means something different at every temperature.

The fix is embarrassingly simple: multiply the soft term by T². The two gradient streams stay comparable, α means what you think it means, and T becomes a clean, independent dial.

The failure mode when you skip it is nasty precisely because it's quiet: nothing crashes, loss goes down, the model trains — just noticeably worse than it should, with the soft targets contributing almost nothing. If you take one practical thing from this whole series, take this.

Analogy for the pair (α, T): you're studying with a textbook and a tutor. α is how you split your hours between them. T is how loudly the tutor thinks out loud — at T=1 they only state final answers; at higher T they explain which wrong options were tempting. And T² is making sure that when the tutor speaks more softly per sentence, you sit proportionally closer, so an hour with them still teaches an hour's worth.

The code

The entire engine, ready to paste:

import torch
import torch.nn.functional as F

def distillation_loss(student_logits, teacher_logits, labels,
                      T: float = 4.0, alpha: float = 0.3):
    """The classic Hinton distillation loss.
    alpha weights the hard-label term; (1 - alpha) the soft term.
    """
    # Term 1 — the answer key
    hard = F.cross_entropy(student_logits, labels)

    # Term 2 — the mentor (both sides at temperature T)
    soft_targets = F.softmax(teacher_logits / T, dim=-1)
    log_student  = F.log_softmax(student_logits / T, dim=-1)
    soft = F.kl_div(log_student, soft_targets, reduction="batchmean")

    # T**2 keeps the two gradient streams comparable — do not delete
    return alpha * hard + (1.0 - alpha) * (T ** 2) * soft

Three implementation notes that save real debugging time. F.kl_div expects log-probabilities for the student and plain probabilities for the teacher — mixing that up is the classic silent bug, and your loss will still go down while learning garbage. The teacher is frozen: compute its logits under torch.no_grad(), no gradients, no optimizer. And at inference time, no temperature anywhere — T exists only inside this loss.

Choosing α and T without a PhD

Honest answer: these are empirical knobs, and anyone quoting universal best values is selling something. But the ranges are well-worn.

Temperature T — typical range 2 to 6. Think of the extremes. At T = 1 you're training on the teacher's raw confidence — usually so peaked it's barely different from hard labels, so why bother. As T → ∞ every distribution flattens toward uniform — the tutor mumbling "eh, anything's possible," pure noise, nothing to learn. The sweet spot is where wrong-answer structure is clearly visible but the right answer still obviously wins. Start at T = 4 (Hinton's own workhorse value). One useful heuristic: bigger gaps between teacher and student confidence patterns, or very confident teachers, tend to want higher T.

Alpha — typical range 0.1 to 0.5 on the hard-label term. Most of the signal should come from the mentor; that's the entire point of showing up. Push α down toward 0.1 when your teacher is excellent and your hard labels are noisy (very common with web-scraped or crowd-labeled data — the teacher's soft opinion is often more reliable than the label). Push it up toward 0.5 when the teacher is mediocre on your domain or you have pristine labels. Start at 0.3 and move it second — T usually matters more.

And when you're unsure: a small grid over T ∈ {2, 4, 6} × α ∈ {0.1, 0.3, 0.5} is nine cheap runs on a small student. Distillation's economics are forgiving here — students are exactly the models you can afford to train nine times.

What's next

That's the theory complete: what transfers, where from, and the precise loss that does it. Time to stop talking and actually shrink a model — teacher, student, training loop, results table, and the satisfying moment where the distilled student beats an identical model that studied alone. That's Part 5.