The Math of Distillation — Why Soft Answers Teach More

Why is the distillation loss KL(teacher||student), what is the temperature T actually doing, and where does that mysterious T² factor in every implementation come from? A ground-up derivation of why a soft distribution carries more teaching signal than a correct answer.

Distilling the Knowledge in a Neural Network

Copy the answer, or copy the margin notes?

Imagine borrowing a classmate's worked problem set. You can copy it two ways. You can copy just the circled answers: "Q3 is B." Or you can copy the hesitation in the margins: "Q3 is B, though C is genuinely tempting; A and D are nowhere close."

Which set of notes teaches you more is obvious. The first says only that the answer is B. The second also tells you that B and C live near each other in the space of plausible answers, while A and D are in a different neighbourhood entirely. Same single question, very different amount of information transferred.

In machine learning these two objects are the hard label and the soft label. The ground truth shipped with your dataset is the former; the probability distribution produced by a trained large model (the teacher) is the latter. Distillation is the training procedure that copies the teacher's soft labels into a smaller model (the student), and this article is about the mathematics of that copy and nothing else. For the wider picture — sequence-level distillation, synthetic-data distillation, real model families — see Knowledge Distillation from Scratch.

What a hard label throws away

Take a classification problem with $C$ classes. A hard label picks one of $C$ options, so a single example carries at most $\log_2 C$ bits. Ten classes gives you a bit over 3.3 bits, and no arrangement of the data will squeeze out more.

A teacher's output, by contrast, is a vector of $C$ probabilities. They must sum to one, so it has $C-1$ degrees of freedom: one example now pins down $C-1$ real numbers in the student rather than a single choice. A vector like "dog 0.7, wolf 0.25, cat 0.03, car 0.02" says the answer is dog, and also that wolf was a near miss while car was never in the running — a ranking and a set of distances. Hinton and colleagues named this surplus dark knowledge.

The crucial point is that the surplus lives in the part where the teacher is wrong. If the teacher emitted a perfect 1.0 on the correct class every time, its output would be identical to the hard label and there would be nothing to distil. The teacher's uncertainty — the specific shape of its confusion — is the teaching material.

Softmax with a temperature

The trouble is that a well-trained teacher gets very confident. Once the output reads "dog 0.999, wolf 0.0008, cat 0.0001…", the ranking information still exists but is numerically flattened to nothing useful. So we deliberately soften the distribution with a temperature $T$:

Here $z_i$ is the logit of class $i$ — the raw score before the softmax — and $T$ is a positive number. Equation (1) says nothing more than: divide the logits by $T$, then apply the usual softmax. At $T=1$ it is the usual softmax. Raising $T$ shrinks the gaps between logits, so the distribution flattens; pushing $T$ toward zero stretches the gaps until only the largest class survives, which is exactly a hard label again.

So the temperature is a continuous dial between "look only at the winner" and "look at the relationships among all classes." Try it:

The loss: why KL?

We want to tell the student "produce the same distribution as the teacher." The standard measure of the gap between two probability distributions is the KL divergence. Writing the teacher's softened distribution as $p$ and the student's as $q$:

This measures how wasteful the student's distribution is when the teacher's is taken as the reference — see KL Divergence from Scratch for the full story.

Now split the logarithm:

$$D_{\mathrm{KL}}(p\|q) = \underbrace{-\sum_i p_i \log q_i}_{\text{cross-entropy } H(p,q)} \;-\; \underbrace{\left(-\sum_i p_i \log p_i\right)}_{\text{teacher's entropy } H(p)}$$

The second term depends only on the teacher, so it is a constant with respect to the student's parameters. Minimising $D_{\mathrm{KL}}(p\|q)$ and minimising the cross-entropy $H(p,q)$ therefore produce identical gradients. Whether a codebase writes the soft loss as a KL or as a cross-entropy against soft targets makes no difference to what gets learned. This is the same decomposition that underlies Entropy and Cross-Entropy.

The direction matters too. $D_{\mathrm{KL}}(p\|q)$ punishes the student severely for putting a small $q_i$ where the teacher put a large $p_i$. In other words: the student is not allowed to zero out an option the teacher considered plausible. Flip the arguments to $D_{\mathrm{KL}}(q\|p)$ and the student can satisfy the loss by collapsing onto a single mode — precisely the hard-label behaviour we were trying to escape. Distillation uses $D_{\mathrm{KL}}(p\|q)$ because it wants that asymmetry.

Where the $T^2$ comes from

Every implementation multiplies the soft loss by $T^2$. This is not a fudge factor; it falls out of the gradient.

Differentiating equation (2) with respect to the student logit $z_i$:

The difference between the student's and teacher's probabilities, divided by $T$. The $1/T$ appears from the chain rule, because we divided the logits by $T$ on the way in.

But the $1/T$ story does not end there: when $T$ is large, $q_i - p_i$ shrinks as well. For large $T$ the ratio $z_i/T$ is near zero, so $\exp(z_i/T) \approx 1 + z_i/T$; assuming the logits are centred at zero this gives $q_i \approx \frac{1}{C} + \frac{z_i}{CT}$, and likewise $p_i \approx \frac{1}{C} + \frac{v_i}{CT}$ for the teacher's logits $v_i$. Taking the difference:

$$\frac{\partial \mathcal{L}_{\text{soft}}}{\partial z_i} \approx \frac{1}{T} \cdot \frac{z_i - v_i}{CT} = \frac{z_i - v_i}{C\,T^2}$$

The gradient shrinks like $1/T^2$, so we multiply by $T^2$ to cancel it. With that factor in place, moving $T$ from 4 to 8 no longer quietly changes how much the soft term contributes, and the mixing weight $\alpha$ below stays the only knob controlling the balance. Forget the $T^2$ and raising the temperature silently switches the soft term off.

The approximation carries a bonus insight. The numerator is $z_i - v_i$, which means high-temperature distillation approaches a squared-error fit between the student's and teacher's raw logits. As you raise $T$, the objective slides smoothly from matching probabilities to matching logits.

The full loss, and the code

In practice the real ground-truth labels are kept alongside the soft targets:

$\alpha$ is the mixing weight in $[0,1]$, $y$ is the one-hot ground truth, and $q(1)$ is the student's distribution back at temperature 1. The hard term is deliberately left unsoftened: its job is to hit the actual answer.

In PyTorch that is all there is to it:

import torch.nn.functional as F

def kd_loss(student_logits, teacher_logits, labels, T=4.0, alpha=0.9):
    soft = F.kl_div(
        F.log_softmax(student_logits / T, dim=-1),   # input must be log-probs
        F.log_softmax(teacher_logits / T, dim=-1),   # pass the target as log too
        reduction="batchmean",                       # 'mean' also divides by C
        log_target=True,
    ) * (T * T)                                      # cancel the 1/T^2 in the gradient
    hard = F.cross_entropy(student_logits, labels)   # this one stays at T = 1
    return alpha * soft + (1 - alpha) * hard

F.kl_div takes (input, target) and computes $D_{\mathrm{KL}}(\text{target} \| \text{input})$, so the student goes in input and the teacher in target. Swap them and training still appears to converge, but the student starts dropping modes — killing off exactly the candidates the teacher kept alive.

How this shows up on the job

Who touches it, and when. This is the standard playbook for an ML engineer cutting inference cost: replacing a large internal model with a small production one. Compressing a classifier, deriving a small LLM from a large one, or self-distillation (training a same-size student on its own teacher's outputs as a regulariser) all differ in intent, but the loss keeps the shape derived above.

The knobs. Effectively two: the temperature $T$ and the mixing weight $\alpha$. Choose $T$ from how peaked the teacher actually is — and the fastest way to find out is not a learning curve but dumping a handful of teacher outputs and reading the runners-up by eye. If the top class sits at 0.999, there is room to raise $T$; if the mass already splits 0.6 / 0.3 / 0.1, a modest $T$ suffices. $\alpha$ encodes how much you trust the teacher: weight the soft term heavily when the teacher is clearly stronger than the student, and lean on the hard term in domains where the teacher's accuracy is shaky.

Failure modes that bite:

• Dropping the $T^2$. The soft term weakens on its own as $T$ rises, and you conclude "temperature doesn't help." Worse, the distortion propagates through your whole hyperparameter sweep.
• The wrong reduction. PyTorch's kl_div with 'mean' divides by the number of classes as well as the batch, shrinking the loss by a factor of $C$. 'batchmean' is the one that matches the definition.
• Leaving the temperature on at inference. $T$ is a training-time device. Serve a student still running at $T=4$ and every confidence score comes out flattened, breaking any downstream threshold or filter.
• Not putting the teacher in eval() under no_grad. Active dropout or BatchNorm makes the teacher's outputs jitter batch to batch, turning your targets into noise — and gradients flowing into the teacher waste compute and memory for nothing.
• Mismatched vocabularies or class definitions. If teacher and student use different tokenizers, logit index $i$ means different things on each side and a logit-space KL is simply undefined. That case forces a switch to sequence-level distillation, training on the teacher's generated text instead.

The interview version. Asked why soft labels help, the clean answer is: each example constrains $C-1$ numbers instead of 1, and the surplus encodes the similarity structure among classes. The natural follow-up — "what happens as $T \to 0$?" — answers itself: you recover hard-label training and the benefit disappears.

Summary

• The distillation loss is $D_{\mathrm{KL}}(p\|q)$; the teacher's entropy is constant, so the gradients are identical to a cross-entropy against soft targets
• Temperature $T$ is a continuous dial between "only the winner matters" and "the relationships among classes matter," collapsing to the hard label as $T \to 0$
• The soft gradient shrinks like $1/T^2$, so the $T^2$ factor makes the meaning of $\alpha$ independent of the temperature
• The direction is $D_{\mathrm{KL}}(\text{teacher} \| \text{student})$; reversing it lets the student collapse onto a mode