Distillation Math — Teacher to Student

What distillation actually transfers

A standard classifier is trained against a one-hot target: for the correct class the label is 1, every other class is 0. That target throws away everything the data could say about how classes relate. A distillation target is instead the teacher’s full probability vector over all classes — the soft target.

Why is the soft vector richer? Because a well-trained teacher does not merely know the answer; it knows the runners-up and their relative plausibility. In a language model over a 50,000-token vocabulary, the distribution after ‘The capital of France is’ puts most mass on Paris but leaves a graded tail across other cities, and that tail is a compressed statement about the model’s learned geometry. Training the student to reproduce the whole vector hands it a far denser supervisory signal per example than a single correct token ever could — which is exactly why a student can learn from fewer examples, or reach accuracy its size alone would not predict.

The softmax and its temperature

Both models turn logits into probabilities with a softmax, but distillation adds a temperature T that divides the logits first:

p_i(T) = exp(z_i / T) / Σ_j exp(z_j / T)

At T = 1 this is the ordinary softmax. As T grows the logits are squashed toward each other, so the distribution gets softer — the peak comes down and the tail comes up. In the limit T → ∞ it approaches the uniform distribution; as T → 0 it sharpens to a one-hot argmax.

The point of raising T is that the interesting dark knowledge lives in the small probabilities. At T = 1 a 0.001 versus a 0.0001 is numerically invisible next to a 0.9; raising the temperature magnifies those ratios so the student can actually feel them in the loss. The same T is applied to both teacher and student while computing the soft loss.

The KL-divergence objective

To make the student’s softened distribution q(T) match the teacher’s p(T), we minimise the Kullback–Leibler divergence from teacher to student:

L_soft = KL(p || q) = Σ_i p_i log(p_i / q_i)
       = Σ_i p_i log p_i  -  Σ_i p_i log q_i

The first term is the teacher’s negative entropy — a constant with respect to the student’s parameters — so gradient descent only ever sees the second term, the cross-entropy -Σ_i p_i log q_i. Minimising KL and minimising this soft cross-entropy are therefore the same optimisation. KL is always ≥ 0 and equals 0 only when q = p exactly, so the objective is honestly measuring ‘how far apart are these two distributions,’ with the teacher’s high-probability classes weighting the sum most heavily.

The combined loss: soft plus hard

Pure imitation has a weakness: the teacher is wrong sometimes, and a student that only mimics inherits those mistakes. So the standard recipe interpolates between the soft target and the true hard label with a mixing weight α:

L = α · T^2 · KL(p(T) || q(T))  +  (1 - α) · CE(y_true, q(1))

The first term pulls the student toward the teacher’s soft distribution at temperature T; the second is ordinary cross-entropy against the ground-truth one-hot label, computed at T = 1. The blend keeps the student anchored to the truth while still absorbing the teacher’s richer structure. Typical settings use α around 0.5–0.9 and T in the 2–10 range, tuned per task. The mysterious T^2 factor on the soft term is not cosmetic — the next section explains why it must be there.

Why the gradients need a T-squared correction

Consider the gradient of the soft cross-entropy with respect to a student logit v_j. Because the student’s logits were divided by T before the softmax, the chain rule leaves a factor of 1/T:

∂L_soft / ∂v_j = (1/T) · (q_j(T) - p_j(T))

Hinton et al. show that for large T, if the logits are zero-meaned, this is approximately (1/T^2)(v_j - z_j) — the gradient magnitude shrinks as 1/T^2. Raising the temperature to expose dark knowledge therefore quietly weakens the soft gradient by the same factor. If you then combine it with the hard loss, the soft term is drowned out. Multiplying L_soft by T^2 exactly cancels the shrinkage, so the soft and hard gradients stay comparable in scale and a single α keeps meaning the same thing as you retune T.

A gradient reading: soft targets as logit matching

The approximation ∂L_soft/∂v_j ≈ (1/T^2)(v_j - z_j) carries a lovely interpretation. After multiplying by T^2, the student’s update is driven by (v_j - z_j) — the difference between the student’s logit and the teacher’s logit. In the high-temperature limit, distillation is essentially regressing the student’s logits onto the teacher’s logits, class by class.

This tells you what the two temperature regimes emphasise. At high T every logit, including the tiny negative ones, contributes to the match, so the student is asked to reproduce the teacher’s full logit landscape — all the dark knowledge. At low T the loss concentrates almost entirely on getting the top class right and largely ignores the tail. Choosing T is thus choosing how much of the teacher’s fine structure you insist the student copy versus how much you let it focus on the headline answer.

A worked numeric example

Take a three-class problem with teacher logits z = [3, 1, 0]. At T = 1 the softmax is:

exp([3,1,0]) = [20.09, 2.72, 1.00],  sum = 23.81
p(1) = [0.844, 0.114, 0.042]

Class 3 sits at 0.042 — small and easy to ignore. Now soften with T = 2, dividing the logits first (z/T = [1.5, 0.5, 0]):

exp([1.5,0.5,0]) = [4.48, 1.65, 1.00],  sum = 7.13
p(2) = [0.629, 0.231, 0.140]

The peak drops from 0.844 to 0.629 and class 3 climbs from 0.042 to 0.140 — more than tripling. The ratio between class 2 and class 3 also becomes gentler (2.7× down to 1.65×). That amplified tail is precisely the signal the student now trains against, information the raw one-hot label [1,0,0] never contained.

What to match: logits, features, and attention

Matching output probabilities is the classic recipe, but a transformer offers richer targets. Feature-based distillation adds a term that pulls the student’s hidden states toward the teacher’s, usually with a mean squared error ||h_student - W·h_teacher||^2, where a learned projection W reconciles the differing widths. Attention transfer matches the attention maps softmax(QK^T / sqrt(d_k)) layer by layer, teaching the student where to look rather than only what to answer.

These intermediate signals matter most when the student is much shallower than the teacher, because output-only matching gives one supervision point at the very end of a deep stack. TinyBERT and DistilBERT combine several of these terms — output KL, hidden-state MSE, attention matching, and sometimes an embedding-layer loss — and the practical art is weighting them so no single term dominates. Layer mapping (which teacher layer supervises which student layer) becomes a design choice in its own right.

Shapes, cost, and the compute budget

The soft loss operates over the vocabulary or class axis. For a language model with sequence length N and vocabulary V, teacher and student each emit logits of shape [N, V], and the KL is summed over V for each of the N positions — cheap arithmetic next to the forward passes that produced the logits.

The real cost is that both models run on every training batch. The teacher does a forward pass (no backward pass needed if it is frozen), the student does forward and backward. So a distillation step costs roughly one extra teacher forward pass over plain student training. Two standard economies help: precompute and cache the teacher’s soft targets once when the dataset is fixed, turning the teacher cost into a one-time expense; or, when memory is tight, keep the teacher in lower precision since its outputs only need to be approximately right to guide the student.

Why distillation is the engine of CPU-SLM

Small language models that run on a CPU exist largely because of distillation. A 7B or 70B teacher captures broad competence; the deployable artifact is a 1B-or-smaller student that must fit in a few gigabytes and answer in real time without a GPU. Training that student from scratch on raw text would need the same enormous corpus and compute the teacher consumed, and would still tend to underperform.

Distillation short-circuits that. The teacher has already done the expensive job of discovering structure in language; its soft targets hand that structure to the student as a dense, per-token curriculum, so the student converges faster and lands higher than same-size scratch training. The result stacks cleanly with the other CPU tricks — quantisation to int8 or int4, pruning, and efficient KV-cache layouts — because a distilled student is simply a better small model to begin with, and a better starting point survives compression with more of its quality intact.

Common pitfalls

Several mistakes recur. Forgetting the T^2 factor leaves the soft loss too weak once the temperature is raised, and the student quietly ignores the teacher — a silent failure, since training still ‘works,’ just without the benefit. Mismatched temperatures at inference are another trap: T is a training-time device, and the deployed student must run at T = 1.

Trusting a weak teacher caps the student at the teacher’s ceiling plus noise; distillation transfers errors as faithfully as insights, so a miscalibrated or overconfident teacher poisons the tail that made soft targets valuable in the first place. Finally, watch the tokenizer and vocabulary alignment: matching output distributions only makes sense if student and teacher share a vocabulary, otherwise the two [N, V] tensors do not even index the same classes. When teacher and student differ there, you fall back to feature- or attention-level matching, or re-tokenize the teacher’s outputs before comparing.