Knowledge distillation trains a smaller, cheaper "student" model to mimic a larger, more capable "teacher" model — transferring much of the teacher's knowledge into a form that's far cheaper to run in production.
The Key Insight: Soft Labels Carry More Information
Training the student only on hard, one-hot labels (like a standard classifier) discards a lot of useful information the teacher actually has. Instead, distillation trains the student to match the teacher's full, soft probability distribution over classes — including the relative probabilities assigned to incorrect classes, sometimes called "dark knowledge." A teacher confidently predicting "dog" but assigning meaningfully more probability to "wolf" than to "car" is communicating genuinely useful information about class similarity that a hard label alone completely discards.
The Distillation Loss
This is exactly the KL Divergence Loss usage previewed in KL Divergence Loss. \(q^T\) denotes softmax with a temperature \(T>1\) (see Temperature (Sampling)) applied to both teacher and student logits — a higher temperature softens both distributions, making the "dark knowledge" in the smaller, non-predicted-class probabilities more prominent and easier for the student to learn from. \(\alpha\) balances this distillation term against a standard hard-label loss.
Code
import torch
import torch.nn.functional as F
def distillation_loss(student_logits, teacher_logits, true_labels, T=3.0, alpha=0.5):
hard_loss = F.cross_entropy(student_logits, true_labels)
soft_teacher = F.softmax(teacher_logits / T, dim=1)
soft_student = F.log_softmax(student_logits / T, dim=1)
soft_loss = F.kl_div(soft_student, soft_teacher, reduction='batchmean') * (T ** 2)
return alpha * hard_loss + (1 - alpha) * soft_loss
# teacher is FROZEN (no gradients); only the student is trained
teacher_model.eval()
with torch.no_grad():
teacher_logits = teacher_model(x)
student_logits = student_model(x)
loss = distillation_loss(student_logits, teacher_logits, true_labels)Why the Temperature-Squared Factor
Dividing logits by a large temperature \(T\) shrinks the gradients' magnitude roughly proportionally to \(1/T^2\) — multiplying the soft loss by \(T^2\) compensates for this, keeping the soft-loss term's contribution to the overall gradient at a comparable scale to the hard-loss term, regardless of the specific temperature chosen.
Common Mistakes
- Distilling with temperature \(T=1\) (no softening at all) — this loses much of the "dark knowledge" benefit, since the teacher's distribution stays sharply peaked and nearly indistinguishable from a hard label for confident predictions.
- Forgetting to freeze the teacher model — the teacher should remain fixed throughout distillation; only the student's weights should ever be updated.
Interview Relevance
Q: "Why does knowledge distillation typically outperform training a small model directly on hard labels alone?" The teacher's full soft probability distribution — including relative probabilities across incorrect classes — carries meaningful information about class similarity ("dark knowledge") that a single hard, one-hot label completely discards. Training the student to match this richer signal, especially with temperature-softened distributions that make this dark knowledge more prominent, gives the student a more informative training target than hard labels alone provide.
Practice Question
Why is a high temperature specifically used during distillation, rather than the standard temperature of 1 used for normal classification?
