Knowledge Distillation Mechanics
The full distillation loss — temperature-scaled soft targets, KL divergence, and why the loss needs a T² correction term
Knowledge Distillation Mechanics
TL;DR
Distillation trains a student to match a teacher's full probability distribution over outputs, not just its single correct label — and a "softened" version of that distribution, produced by dividing logits by a temperature before the softmax, carries far more information about which wrong answers the teacher considers plausible. KL divergence measures the gap between the two distributions, and a T² correction keeps that loss term's gradient scale comparable to an ordinary hard-label loss it's typically combined with.
| Property | Value |
|---|---|
| Level | Intermediate |
| Reading time | ~18 minutes |
| Prerequisites | Quantization From First Principles |
| You will understand | The exact math of the distillation loss, why temperature matters, and where the T² term comes from |
Hard Labels vs Soft Targets
A standard classification label is a one-hot vector — all the probability mass on the correct class, zero everywhere else:
hard label for "cat": [1, 0, 0, 0, 0] (classes: cat, dog, bird, car, tree)A trained teacher model's raw output, after a normal softmax, might look like:
teacher output: [0.92, 0.06, 0.01, 0.0003, 0.0007]That distribution already carries more information than the hard label. It says the teacher considers "dog" far more plausible than "bird," which is far more plausible than "car" — a hierarchy of near-misses the hard label simply can't express. Training a student to match this whole distribution, not just the top answer, teaches it something about the relationships between classes that hard labels never could.
Temperature: Making the Soft Signal Louder
The problem: with a confident teacher, the raw softmax output above is already close to one-hot — 0.92 vs 0.06/0.01/near-zero doesn't leave much usable signal in the smaller values. Temperature scaling fixes this by dividing the logits by T > 1 before the softmax, which flattens the distribution and makes the relative differences between the smaller probabilities much more visible.
def softmax_with_temperature(logits, T):
return torch.softmax(logits / T, dim=-1)A worked numeric example
Take teacher logits [4.0, 1.5, -0.5, -4.0, -3.5] for the same five classes:
T = 1 (ordinary softmax):
softmax([4.0, 1.5, -0.5, -4.0, -3.5]) ≈ [0.895, 0.075, 0.010, 0.0003, 0.0005]
→ almost one-hot, barely distinguishable non-max values
T = 4:
softmax([1.0, 0.375, -0.125, -1.0, -0.875]) ≈ [0.421, 0.229, 0.140, 0.058, 0.066]
→ the same ranking, but now with real, comparable magnitude —
"dog" (0.229) is clearly more plausible than "car" (0.058) to the studentRaising T from 1 to 4 didn't change which class the teacher favors — it revealed the relative confidence among the non-max classes that was nearly invisible at T=1. That revealed structure is exactly the extra signal distillation is trying to transfer.
The Loss: KL Divergence Between Distributions
KL divergence measures how much one probability distribution diverges from a reference one — the natural choice when the goal is literally "make the student's output distribution match the teacher's":
KL(P_teacher ‖ P_student) = Σ P_teacher(i) · log( P_teacher(i) / P_student(i) )def distillation_loss(student_logits, teacher_logits, T):
student_log_probs = F.log_softmax(student_logits / T, dim=-1)
with torch.no_grad():
teacher_probs = F.softmax(teacher_logits / T, dim=-1)
return F.kl_div(student_log_probs, teacher_probs, reduction="batchmean")teacher_logits are wrapped in torch.no_grad() because the teacher is frozen — only the student's parameters should receive gradient updates from this loss.
Why the Loss Needs a T² Correction
Distillation loss is typically combined with an ordinary hard-label loss against the true labels:
total_loss = α · hard_label_loss + (1 - α) · T² · distillation_lossThe T² factor exists because raising the temperature shrinks the distillation loss term's gradient magnitude by roughly 1/T². Dividing logits by T before the softmax scales down the gradients that flow back through it by the same factor — so without correcting for it, a high-temperature distillation term would contribute a much weaker training signal than the hard-label term it's combined with, purely as an artifact of the temperature choice rather than anything meaningful about how important that term should be. Multiplying by T² restores the distillation term's gradient to roughly the scale it would have had at T=1, so the temperature can be tuned for its intended purpose — controlling how much soft-label information is revealed — without also silently changing how strongly that term influences training.
The combined loss lets the student learn from two complementary signals at once: the ground-truth label (hard_label_loss, weighted by α) and the teacher's richer, near-miss-aware distribution (distillation_loss, weighted by 1 - α).
Where This Sits Relative to On-Policy Distillation
Everything on this page describes offline, static distillation: the teacher's outputs are computed once (or ahead of time) on a fixed dataset, and the student trains against those fixed targets. This is the foundational mechanism, and it's straightforward — but it has a known weakness: the student only ever sees the teacher's judgment on inputs the dataset happened to contain, never on situations the student's own generation might wander into during actual use.
On-policy distillation, covered practically via TRL's GKDTrainer in the Small Language Models track, changes what gets fed into this same loss mechanism — instead of only training on fixed teacher-generated or ground-truth examples, the student generates its own rollouts, and the teacher scores or corrects those instead. The loss machinery (temperature-scaled soft targets, KL divergence, the T² correction) is the same foundation either way; what changes is whose generations the loss gets computed over.
Concept Checks
Check yourself
Why does a confident teacher's ordinary (T=1) softmax output carry less usable distillation signal than a temperature-softened version?
Because a confident teacher's raw output is already close to one-hot — the correct class dominates and the other probabilities are tiny and hard to distinguish from each other. Raising the temperature flattens the distribution, making the relative differences among the smaller, non-max probabilities much larger and more informative, revealing which wrong answers the teacher considers more or less plausible.
Why is KL divergence a more natural loss for distillation than, say, ordinary cross-entropy against a one-hot label?
Because KL divergence directly measures the gap between two full probability distributions, which is exactly what distillation is trying to minimize — the distance between the student's output distribution and the teacher's. Cross-entropy against a one-hot label only measures how well the student predicts the single correct class, discarding all the relative-plausibility information the teacher's softened distribution carries about the other classes.
What specifically goes wrong if you scale up the temperature for distillation without multiplying the distillation loss term by T²?
The distillation loss term's gradient magnitude shrinks by roughly 1/T² as a side effect of dividing logits by T before the softmax, so at a high temperature it would contribute a much weaker training signal relative to the hard-label loss term purely because of the temperature choice, not because that term was actually meant to matter less. Multiplying by T² restores its gradient scale so temperature can be tuned for what it's meant to control — how much soft-label detail is revealed — without unintentionally also changing that term's influence on training.
Key Concepts Recap
| Concept | One-line summary |
|---|---|
| Hard label | A one-hot vector — all probability on the correct class |
| Soft target | A teacher's full output distribution, carrying near-miss information hard labels can't |
| Temperature scaling | Dividing logits by T > 1 before softmax to reveal relative confidence among non-max classes |
| KL divergence | Measures how much the student's distribution diverges from the teacher's — the natural distillation loss |
| T² correction | Restores the distillation loss's gradient scale after temperature shrinks it by roughly 1/T² |
| Combined loss | A weighted sum of hard-label loss and (T²-corrected) distillation loss |
| Offline vs on-policy | This page's static targets vs GKD-style training on the student's own corrected rollouts — same loss mechanics, different source of examples |
Next
With compression covered from two angles, the next page addresses a different kind of adaptation — teaching a model to prefer better outputs, not just imitate examples: Alignment & DPO Math.
Quantization From First Principles
The exact math behind mapping float weights to integers, worked numerically, plus symmetric vs asymmetric, per-tensor vs per-channel, and the straight-through estimator behind QAT
Alignment & DPO Math
The Bradley-Terry model of preference, why classic RLHF needs a separate reward model, and the full derivation of DPO's reward-model-free loss