Knowledge Distillation: Making Models Smaller

When 99.2% accuracy is too slow to matter

NPCI's UPI network settles tens of billions of transactions a month, which averages out to several thousand transactions every second during peak load — festival weekends, salary days, the last over of an IPL final. Picture a bank's fraud desk that has trained a large transformer-based classifier on years of transaction history: device fingerprint, geolocation drift, merchant category, time-of-day pattern, graph features pulled from the payer-payee network. The model reaches 99.2% accuracy on held-out data, comfortably ahead of every earlier model the bank has shipped. There is one problem: scoring a single transaction takes 400 milliseconds, and UPI's settlement window budgets well under 100 milliseconds end-to-end, most of which is already spent on network hops and the two banks' own ledger checks. A model this accurate that cannot fit inside the window is not a fraud filter — it is a research artifact sitting in a notebook. The engineers now face the standard dilemma of production machine learning: the model that generalizes best is rarely the model you can afford to serve.

The reflexive fix — train a smaller network directly on the same labeled transactions — usually falls short. A six-layer network trained from scratch on one-hot fraud / not-fraud labels typically lands several accuracy points below what the large model achieves, even holding architecture family and training compute roughly comparable. Knowledge distillation is the technique that closes most of that gap. Instead of training the small "student" model on the raw ground-truth labels alone, you additionally train it to imitate the full output distribution of the large "teacher" model. A trained teacher's near-misses carry more supervisory signal per example than the ground-truth label by itself, and that signal is exactly what a small model — with too few parameters to rediscover the same structure unaided from raw data — needs to borrow.

The teacher-student framework: what distillation actually trains

Formalize the setup. The teacher is a large network, already trained to convergence and then frozen — its weights never update again during distillation. The student is a smaller network, with fewer layers, narrower hidden dimensions, or both, initialized fresh and trained from scratch. Both networks see the same input x and each produces a vector of logits — raw, pre-softmax scores — over the same K classes: z_T for the teacher, z_S for the student.

At inference time, any classifier converts logits to probabilities with the ordinary softmax, q_i = exp(z_i) / Σ_j exp(z_j). The trouble at this default temperature is that a well-trained network tends to produce very peaked distributions: one class near 1.0, everything else pinned near 0. That is exactly what you want for a final prediction, but it is nearly useless as a training signal, because it discards the information about which of the wrong classes were almost right. Hinton, Vinyals, and Dean's 2015 paper, "Distilling the Knowledge in a Neural Network" (arXiv:1503.02531), introduced a temperature-scaled softmax to recover exactly that information:

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

At T = 1 this is the ordinary softmax. Raising T "softens" the distribution: it compresses the gaps between logits before exponentiating, so probability mass spreads out over the non-winning classes instead of collapsing onto the argmax. Hinton called the resulting structure "dark knowledge." For the fraud model, the class-similarity information — that a transaction the teacher scores "Legit" with high confidence was nonetheless, in the teacher's internal representation, far closer to "Suspicious" than to outright "Fraud" — never appears in a one-hot label at all, and barely survives an ordinary T = 1 softmax either.

The full training loss for the student combines two terms: a soft-target loss that pulls the student's softened distribution toward the teacher's softened distribution, and a hard-target loss that pulls the student's ordinary (T = 1) prediction toward the ground-truth label, weighted by α:

L_total = α · T² · L_soft(q_T, q_S) + (1 − α) · L_hard(ŷ, y)

L_soft is typically KL divergence or cross-entropy between the two softened distributions, both computed at the same temperature T; L_hard is ordinary cross-entropy against the true label at T = 1, exactly as if the student were being trained without a teacher at all. The T² factor is not decorative — it corrects a gradient-scaling artifact, derived below.

A worked example: what temperature actually does to a distribution

Take a concrete instance. The teacher scores one transaction with logits z = [3.0, 1.0, 0.0] over three classes, [Legit, Suspicious, Fraud]. At T = 1:

exp(3.0) = 20.09,  exp(1.0) = 2.72,  exp(0.0) = 1.00,  sum = 23.80
q(T=1) = [0.844, 0.114, 0.042]

The teacher is 84% confident in "Legit," reasonable for a genuine transaction, but look at the ratio between the two non-winning classes: 0.114 / 0.042 ≈ 2.72. The teacher's internal geometry says "Suspicious" is nearly three times as likely as "Fraud," but a downstream reader of the T = 1 output would barely register the 4.2% Fraud mass at all. Raise T to 4 — divide every logit by 4, then exponentiate:

exp(0.75) = 2.117,  exp(0.25) = 1.284,  exp(0) = 1.00,  sum = 4.401
q(T=4) = [0.481, 0.292, 0.227]

The same ratio, 0.292 / 0.227 ≈ 1.284, is now visible as a modest gap between two much larger numbers — a shape a small network's loss function can actually fit. Push T further, to 20:

exp(0.15) = 1.162,  exp(0.05) = 1.051,  exp(0) = 1.00,  sum = 3.213
q(T=20) = [0.362, 0.327, 0.311]

All three classes now sit within five percentage points of uniform (1/3 = 0.333). The ratio has collapsed to 0.327 / 0.311 ≈ 1.051, barely distinguishable from 1. The general pattern: for two classes with logit gap Δz, the ratio of their temperature-scaled probabilities is exp(Δz / T), which tends to 1 as T → ∞ for any fixed Δz. Push the temperature high enough and every class becomes equally likely, regardless of what the teacher actually learned. The diagram below plots exactly these nine numbers as proportional bars.

Why the loss gets multiplied by T²

Let u_i = z_i / T, so q_i = softmax(u)_i. For cross-entropy loss C = −Σ_i p_i log q_i against a fixed soft target p — the teacher's distribution, which carries no gradient since the teacher is frozen — the standard softmax-cross-entropy gradient with respect to u_i is the familiar result ∂C/∂u_i = q_i − p_i. Applying the chain rule through u_i = z_i / T:

∂C/∂z_i = (∂C/∂u_i) · (∂u_i/∂z_i) = (q_i − p_i) / T

The gradient the student receives from the soft-target loss shrinks in proportion to 1/T. Left uncorrected, raising temperature to expose dark knowledge would simultaneously and silently shrink the soft-target term's contribution to training relative to the hard-target term, which carries no such 1/T factor — it is always computed at T = 1. Multiplying L_soft by T² restores the gradient magnitude the soft-target term contributes; Hinton et al. show the correction is close to exact in the high-temperature regime, where the softened logits are small relative to T and q, p can be treated as a small perturbation around the uniform distribution. This is why the recommended loss is α·T²·L_soft rather than α·L_soft — an implementation detail that keeps the α weighting meaningful across different choices of T, and one that is easy to silently get wrong when copying a distillation loss between codebases.

Tracing one full training step

The following PyTorch snippet computes exactly the quantities above for a full training step: teacher logits [3.0, 1.0, 0.0] against student logits [2.5, 1.0, 0.2], at T = 4, with α = 0.7 and the true label "Legit" (index 0). It was run to confirm every printed number below; nothing here is estimated.

import torch
import torch.nn as nn
import torch.nn.functional as F

z_teacher = torch.tensor([3.0, 1.0, 0.0])
z_student = torch.tensor([2.5, 1.0, 0.2], requires_grad=True)
true_label = torch.tensor([0])          # index 0 = "Legit"
T = 4.0
alpha = 0.7

soft_teacher = F.softmax(z_teacher / T, dim=0)
soft_student_log = F.log_softmax(z_student / T, dim=0)

kl = nn.KLDivLoss(reduction="sum")
L_soft = kl(soft_student_log, soft_teacher) * (T ** 2)
L_hard = F.cross_entropy(z_student.unsqueeze(0), true_label)
L_total = alpha * L_soft + (1 - alpha) * L_hard

print("soft_teacher:", soft_teacher.tolist())
print("L_soft:", L_soft.item())
print("L_hard:", L_hard.item())
print("L_total:", L_total.item())

L_total.backward()
print("grad wrt student logits:", z_student.grad.tolist())

soft_teacher: [0.4810, 0.2918, 0.2272]
L_soft: 0.0457
L_hard: 0.2802
L_total: 0.1161
grad wrt student logits: [-0.1757, 0.0890, 0.0868]

Two things are worth tracing by hand from this output. First, note nn.KLDivLoss computes true KL divergence, D_KL(p‖q) = Σ_i p_i log(p_i / q_i), not the plain cross-entropy used earlier in this chapter's derivation — the two differ by the constant H(p), the teacher's own entropy, which does not depend on the student and therefore contributes nothing to the gradient. That is why L_soft above (0.0457) looks much smaller than the "L_soft" figure a plain cross-entropy would report; both are legitimate choices, and most production codebases, including the original Hinton implementation, use the KL form because it goes to exactly zero when student matches teacher perfectly. Second, look at the gradient's first component, −0.1757 — the largest-magnitude entry, and negative, meaning gradient descent will push z_S[0] (the "Legit" logit) up. That is the correct direction: the true label is Legit, and both loss terms agree the student should become more confident there. The two smaller, positive components on Suspicious and Fraud will be nudged down, but not by equal amounts — L_soft alone is pulling the student's Suspicious/Fraud split toward the teacher's 0.292/0.227 ratio, not toward zero on both, which is exactly the "dark knowledge" transfer the temperature was introduced to enable in the first place.

DistilBERT: distillation at production scale

The clearest large-scale demonstration of this technique in NLP is DistilBERT, introduced by Victor Sanh, Lysandre Debut, Julien Chaumond, and Thomas Wolf at Hugging Face in 2019 ("DistilBERT, a distilled version of BERT: smaller, faster, cheaper and lighter," arXiv:1910.01108). BERT-base has 12 transformer encoder layers and roughly 110 million parameters. DistilBERT halves the depth to 6 layers while keeping the same hidden size and vocabulary, initializing each retained student layer directly from the corresponding teacher layer rather than from random weights — taking one layer out of every two — which the authors found sped up convergence substantially compared to training the student from scratch. The result has roughly 66 million parameters, about 40% fewer than BERT-base, runs about 60% faster at inference, and retains about 97% of BERT's language-understanding performance as measured on the GLUE benchmark suite.

What makes DistilBERT a useful case study for this chapter is its loss function, which extends the two-term recipe above into three terms. Alongside the temperature-scaled distillation loss over the model's masked-language-modeling output distribution (response-based distillation, exactly as derived above) and the standard masked-language-modeling loss against the actual masked tokens (this chapter's L_hard), DistilBERT adds a cosine embedding loss that aligns the direction of the student's hidden-state vectors with the teacher's at the same layer positions. This third term is a form of feature-based distillation layered on top of response-based distillation: it constrains not just what the student's final output distribution looks like, but how its internal representations are oriented relative to the teacher's, which the authors found meaningfully improved downstream task performance beyond output-matching alone. For a system like the fraud classifier from this chapter's opening, the same three-term recipe transfers directly: match the softened fraud/legit/suspicious distribution, match the true label, and — if the teacher and student share compatible internal layer widths — match selected hidden-state directions too.

Beyond matching outputs: feature and relation distillation

Response-based distillation, matching only the final softened output, is the simplest variant and the one this chapter has derived in full, but it is not the only one. Adriana Romero, Nicolas Ballas, Samira Ebrahimi Kahou, Antoine Chassang, Carlo Gatta, and Yoshua Bengio's FitNets (2014, published at ICLR 2015) introduced feature-based distillation: rather than waiting until the final layer, a "hint" loss trains an intermediate student layer to reproduce an intermediate teacher layer's activations directly, using a learned linear regressor to bridge any mismatch in layer width. This lets a student that is much thinner than the teacher, not just shallower, still absorb structure from deep inside the teacher's representation rather than only from its final decision. A third family, relation-based distillation, goes further still: instead of matching individual activations at all, it matches the relationships between pairs or triples of examples — the pairwise distances or angles the teacher induces across a batch — so the student learns the teacher's geometry of similarity between examples, not merely its output at any single input in isolation. All three families compose: DistilBERT's cosine embedding term is itself a lightweight feature-based addition layered on top of a response-based backbone, and nothing prevents adding a relation-based term to the same total loss.

Common misconception: "raise the temperature as high as you can"

Given that raising T from 1 to 4 revealed useful structure that T = 1 hid, it is tempting to conclude that distillation should always be run at the highest temperature the training loop tolerates, on the theory that more softening always means more information transferred. The worked example above shows exactly why this is wrong: pushing T from 4 to 20 did not reveal more structure, it erased structure that T = 4 had already exposed. At T = 20 the three classes sat within five points of the uniform 1/3 baseline, and the informative Suspicious-versus-Fraud ratio had collapsed from 1.284 to 1.051 — closer to indistinguishable than at T = 4. The mathematics behind this is not subtle: the ratio between any two softened probabilities is exp(Δz / T), and that expression converges to exp(0) = 1 for every pair of classes simultaneously as T grows without bound, regardless of how different their original logits were. Temperature does not add information; it redistributes a fixed quantity of information — the shape of the logit vector — across a probability simplex, and past a certain point that redistribution simply flattens everything toward the maximum-entropy uniform distribution, which contains no information about the teacher's learned distinctions at all. In practice, published distillation results typically use T somewhere in the range of 2 to 10; the correct temperature is an empirical choice made by validating the student's downstream accuracy, not a value pushed toward infinity on the assumption that softer is always more informative.

Active recall

Attempt each question before reading its answer.

1. Why does a small student network trained directly on one-hot ground-truth labels typically underperform the same architecture trained by distillation from a larger teacher, given identical training data?
2. A teacher scores an image with logits [4, 1, −1] over three classes. Compute the softmax probabilities at T = 1 and at T = 2, showing the intermediate exponentials.
3. In the worked training step above (teacher logits [3, 1, 0], student logits [2.5, 1, 0.2], α = 0.7, T = 4), suppose T is raised to 8 with every other quantity held fixed. Trace what happens to: (a) q_T and q_S, (b) the unscaled KL divergence, (c) L_soft = T²·KL, (d) L_hard, (e) L_total, (f) the gradient magnitude reaching the student. Which of these barely move, and why?
4. Name DistilBERT's three loss terms and state, in one line each, what each one supervises.
5. A classmate argues: "To transfer the maximum amount of information from teacher to student, set T as high as the training loop allows." Using the T = 1 / 4 / 20 table, explain precisely why this claim is false, and state what actually determines a good choice of T.
6. In the UPI fraud scenario, both the distilled student and a same-sized model trained from scratch on the same labeled transactions cost identically little to run at inference. Why is the distilled version still likely to be more accurate?

Answers.

1. One-hot labels only tell the student which single class is correct; they say nothing about how the correct class relates to the incorrect ones. A trained teacher's softened output additionally encodes that relational structure — how similar the teacher judges each wrong class to be to the right answer — which acts as a richer, per-example supervisory signal than a single bit of "which class" information. A small student, with too few parameters to independently rediscover that similarity structure from raw data the way the larger teacher did, learns it faster and generalizes better when it is handed directly through the soft-target loss, rather than being forced to reconstruct it unaided from hard labels alone.

2. At T = 1: exp(4) = 54.60, exp(1) = 2.72, exp(−1) = 0.37, sum = 57.68, giving q(T=1) = [0.947, 0.047, 0.006]. At T = 2, divide logits by 2 first: exp(2) = 7.39, exp(0.5) = 1.65, exp(−0.5) = 0.61, sum = 9.64, giving q(T=2) = [0.766, 0.171, 0.063]. Note how much sharper this distribution is than the [3,1,0] example in the chapter body: a wider logit spread (5 here, versus 3 there) means even T = 2 leaves the top class heavily dominant, since the ratio-collapsing effect of temperature has more spread to work against.

3. Running the exact computation at T = 8: q_T becomes [0.406, 0.316, 0.279] and q_S becomes [0.388, 0.321, 0.291] — both distributions move further toward uniform than at T = 4, as expected. (a) Both q_T and q_S flatten. (b) The unscaled KL divergence shrinks from 0.00286 to 0.00071, roughly a fourfold drop — consistent with KL scaling approximately as 1/T² near the uniform distribution, since doubling T roughly quarters the divergence. (c) L_soft = T²·KL barely moves, from 0.0457 to 0.0453 (under a 1% change), because the T² factor very nearly cancels the roughly-1/T² shrinkage in (b) — the small residual (under 1%, rather than exactly 0%) is the expected gap between the exact high-temperature approximation and this specific, moderate value of T. This is the T² correction from the "why T²" section working essentially as designed. (d) L_hard is completely unchanged at 0.2802, because it is computed at T = 1 regardless of the distillation temperature chosen. (e) L_total moves only from 0.1161 to 0.1158, under a 0.3% change, since its only T-dependent piece (L_soft) itself barely moved. (f) The gradient's dominant component moves only from −0.1757 to −0.1729, about 1.6%. The lesson: although the softmax distributions themselves changed substantially between T = 4 and T = 8, essentially every downstream training quantity — loss and gradient alike — stayed nearly fixed, precisely because the T² scaling was designed to make training outcomes insensitive to the exact temperature chosen.

4. (i) A temperature-scaled distillation loss over the masked-language-modeling output distribution, matching the student's soft predictions to the teacher's. (ii) The standard masked-language-modeling loss, training the student to predict masked tokens correctly against the real (hard) labels. (iii) A cosine embedding loss aligning the direction of the student's hidden-state vectors with the teacher's corresponding layer, transferring internal representational structure rather than only final-output behavior.

5. False, because temperature does not create new information — it only redistributes a fixed amount of information (the shape of the teacher's logit vector) across the probability simplex. The ratio between any two classes' softened probabilities is exp(Δz/T), which tends to 1 for every class pair simultaneously as T grows, so past some point raising T further erases the very distinctions it was meant to expose rather than clarifying them. In the T = 1/4/20 table, T = 4 exposes the Suspicious-versus-Fraud relationship clearly (ratio 1.28), while T = 20 nearly erases it (ratio 1.05) — worse than T = 4, not better. The correct choice of T is an empirical one, validated against the student's actual downstream accuracy, typically landing somewhere between 2 and 10 in published work, not pushed to the largest value the training loop permits.

6. Both models have identical inference cost by construction, so the comparison is purely about what each learned during training, not what either costs to run. The teacher was very likely trained with far more capacity, data, or compute than the small model could use directly — larger context windows over transaction history, ensembling, or heavier regularization that only pays off at large parameter counts — and it converged to a set of internal similarity judgments between classes that a same-sized model trained from scratch, with only one-hot labels to learn from, has no direct access to and would need vastly more data or training time to rediscover independently. Distillation transfers that already-learned structure directly through the soft-target loss, effectively letting the small model borrow generalization capacity it could not have developed on its own within the same parameter budget.

Think About It

Think about this: How would you explain knowledge distillation: making models smaller to a friend who has never seen a computer? What real-world analogy would you use? Imagine you had to build a system using these concepts — what would be your first step? Try this: before moving on, write down three things you learned and one question you still have.

Key Takeaways — Summary and Recap

Let us recap what we covered: the core ideas behind knowledge distillation: making models smaller, how they connect to real-world applications, and why they matter for your journey in computer science. Remember these key points as you move forward. For competitive exam preparation (CBSE, JEE, BITSAT), focus on understanding the WHY behind each concept, not just the WHAT.