Module 28: Knowledge Distillation

Introduction

A frontier model is expensive to serve. Knowledge distillation is the oldest and most widely-used way to shrink one: train a small student network to imitate a large teacher — not by copying its weights (that was m27), but by copying its behavior, the full probability distribution it puts over outputs.

The trick is that a teacher’s soft probabilities say far more than a label does. Shown a photo of a car, a one-hot label says only car. The teacher’s distribution also says truck is far more plausible than cat — it has learned the similarity structure of the world. That extra signal, which Hinton called dark knowledge, is what a student learns from. A small model trained on soft targets can reach accuracy a small model trained on hard labels cannot.

Why it matters for LLMs:

  • Deployment. Every small model you actually run — a phone-sized assistant, a fast draft model (m16) — is likely a distilled one.
  • Reasoning on a budget. In 2025, DeepSeek showed that distilling a giant reasoning model’s traces into a small one beats running expensive RL (m13) on that small model directly.
  • It is the same loss everywhere. The soft-label idea reappears as the tuned lens’s training objective (m21), as sequence-level KD for translation, and as on-policy distillation for chat models.

What You’ll Learn

After this module, you can:

  • Soften a distribution with temperature and read the dark knowledge it exposes.
  • Write Hinton’s distillation loss from scratch and explain the $T^2$ factor.
  • Prove — numerically — that high-temperature distillation is logit matching.
  • Distinguish word-level from sequence-level KD for a language model, and see that sequence-level KD is just SFT on the teacher’s own outputs.
  • Explain on-policy distillation (GKD) and why the choice of divergence (forward vs reverse KL) trades mode-covering for mode-seeking.
  • Retell the R1 distillation result and why it reshaped how small reasoners are built.

Prerequisites

This module requires familiarity with:

Intuition: Dark Knowledge

Cross-entropy against a one-hot label only ever pushes up the probability of the correct class. It is indifferent to how the student distributes the leftover mass. But the teacher is not indifferent — it has opinions about the wrong answers, and those opinions encode real structure.

Temperature is the dial that exposes them. Divide the logits by $T$ before the softmax: at $T=1$ you get the ordinary (peaky) distribution; as $T$ grows the peak flattens and the plausible runner-ups rise into view. Drag the temperature below and watch the teacher’s faith in truck — not cat — emerge from behind the car spike.

TipTry This
  1. Slide $T$ from 1 to 12. At $T=1$ the car bar swallows everything. As $T$ rises, truck climbs while cat, plane, ship stay low — the teacher is telling you a truck looks like a car and a cat does not.
  2. Watch the ratio in the top-left. A one-hot label sets it to 0/0. The soft target keeps truck several times larger than cat — that is the signal a student gets for free.

The Math: Temperature and the Distillation Loss

Write the teacher’s logits $v$ and the student’s $z$, both over a vocabulary of $V$ classes. The temperature-softened distributions are

\[ p^{\text{teacher}}_i = \frac{\exp(v_i / T)}{\sum_j \exp(v_j / T)}, \qquad q^{\text{student}}_i = \frac{\exp(z_i / T)}{\sum_j \exp(z_j / T)}. \]

The distillation loss asks the student to reproduce the teacher’s soft distribution (via KL divergence), optionally blended with the ordinary hard-label cross-entropy $H$:

\[ \mathcal{L} = \underbrace{\alpha \, H(y,\, q^{T=1})}_{\text{hard label}} \;+\; \underbrace{(1-\alpha)\, T^2 \, \mathrm{KL}\!\left(p^{\text{teacher}} \,\|\, q^{\text{student}}\right)}_{\text{soft target}}. \]

Two details earn their keep. The gradient of the soft term through the softmax scales like $1/T^2$, so the explicit $T^2$ factor rescales it back to the same magnitude as the hard term — you can change $T$ without retuning $\alpha$. And $\mathrm{KL}$, not plain cross-entropy, keeps the “match the whole distribution” reading exact.

Step by Step

Code: Distillation From Scratch

Everything above is a few lines of PyTorch. The reference implementation lives in distillation.py; here we build the core pieces inline.

import torch
from distillation import softmax_T, soft_targets, distillation_loss, kl_divergence

# A teacher confident it's a "car" (index 2), with a real hint toward "truck" (3).
teacher_logits = torch.tensor([0.2, 0.6, 6.0, 3.2, 0.4, 1.1])

hard = softmax_T(teacher_logits, temperature=1.0)
soft = softmax_T(teacher_logits, temperature=4.0)
print("T=1  (peaky):", [f"{v:.3f}" for v in hard.tolist()])
print("T=4  (soft) :", [f"{v:.3f}" for v in soft.tolist()])
print("truck/cat ratio at T=4:", round((soft[3] / soft[0]).item(), 1), "x")
T=1  (peaky): ['0.003', '0.004', '0.926', '0.056', '0.003', '0.007']
T=4  (soft) : ['0.093', '0.102', '0.395', '0.196', '0.097', '0.116']
truck/cat ratio at T=4: 2.1 x

The loss itself is the hard/soft blend, with the $T^2$ factor on the soft term:

# One batch of 4 examples over a 6-class vocabulary.
torch.manual_seed(0)
student_logits = torch.randn(4, 6, requires_grad=True)
teacher_batch = torch.randn(4, 6)
labels = torch.randint(0, 6, (4,))

loss = distillation_loss(student_logits, teacher_batch, labels, temperature=4.0, alpha=0.5)
loss.backward()
print(f"distillation loss: {loss.item():.4f}")
print("student receives gradients:", student_logits.grad is not None)
distillation loss: 1.6313
student receives gradients: True

High Temperature is Logit Matching

Hinton’s key result: at high $T$ (with per-example zero-meaned logits) the gradient of the $T^2$-scaled soft loss w.r.t. a student logit approaches $\tfrac{1}{V}(z_i - v_i)$ — exactly the gradient of $\tfrac{1}{2}\sum_i (z_i - v_i)^2$. So distillation at high temperature is regression onto the teacher’s logits. We can check it numerically: the gap between the true autograd gradient and that approximation collapses as $T$ grows.

from distillation import verify_logit_matching

torch.manual_seed(0)
s = torch.randn(3, 8)
t = torch.randn(3, 8)
for T in (2.0, 10.0, 50.0, 200.0):
    err = verify_logit_matching(s, t, temperature=T)
    print(f"T={T:6.1f}   max|grad - logit-match|  =  {err:.2e}")
T=   2.0   max|grad - logit-match|  =  3.73e-02
T=  10.0   max|grad - logit-match|  =  8.26e-03
T=  50.0   max|grad - logit-match|  =  1.67e-03
T= 200.0   max|grad - logit-match|  =  4.18e-04
NoteKey Insight

Temperature interpolates between two regimes. At $T=1$ distillation cares mostly about the top class — close to a hard label. At high $T$ it becomes pure logit regression, forcing the student to reproduce the teacher’s entire score vector, dark knowledge and all. The sweet spot ($T \in 2..10$ in most papers) keeps the runner-up structure without drowning the correct answer.

Dark Knowledge, Measured

Does the extra signal actually help a small student? It does exactly where you’d expect: when labels are scarce. Below, a wide teacher is trained on plenty of data; a narrow student then sees only 32 labeled points. The hard student gets just those one-hot labels; the kd student also sees the teacher’s soft targets on the same 32 inputs — and generalizes measurably better on held-out data.

from distillation import train_student

hard_run = train_student("hard", seed=0)
kd_run = train_student("kd", seed=0)
print(f"teacher (trained on all data): {hard_run['teacher_acc']:.3f}")
print(f"student, hard labels only    : {hard_run['final_acc']:.3f}")
print(f"student, + soft targets (KD) : {kd_run['final_acc']:.3f}"
      f"   (+{kd_run['final_acc'] - hard_run['final_acc']:.3f})")
teacher (trained on all data): 0.806
student, hard labels only    : 0.725
student, + soft targets (KD) : 0.819   (+0.094)
TipTry This
  1. Read the gap. The KD curve sits above the hard curve for essentially the whole run — same 32 labels, more knowledge per label.
  2. In distillation.py, raise n_labeled toward the full set. The gap shrinks: dark knowledge matters most when labels are the bottleneck.

Distilling a Language Model: Word-level vs Sequence-level

A language model outputs a next-token distribution at every position, so there are two honest ways to distill it.

Word-level KD matches the teacher’s full distribution at each position — a dense KL over the whole vocabulary, every step. It is the direct generalization of what we did above:

from distillation import word_level_kd_loss, sequence_level_kd_loss

torch.manual_seed(0)
# A short sequence: per-position logits over a 12-word vocabulary.
teacher_seq = torch.randn(5, 12)
student_seq = torch.randn(5, 12, requires_grad=True)

wl = word_level_kd_loss(student_seq, teacher_seq, temperature=1.0)
print(f"word-level KD (mean per-token KL): {wl.item():.4f}")
word-level KD (mean per-token KL): 0.5614

Sequence-level KD (Kim & Rush, 2016) takes a different route: let the teacher generate an output sequence (e.g. by beam search), then train the student with ordinary next-token cross-entropy on those tokens. The student learns to reproduce the sequences the teacher actually produces — and crucially, sequence-level KD is just SFT on the teacher’s outputs:

import torch.nn.functional as F

# Pretend the teacher generated these tokens (here: its own argmax path).
teacher_tokens = teacher_seq.argmax(dim=-1)

seq_kd = sequence_level_kd_loss(student_seq, teacher_tokens)
plain_sft = F.cross_entropy(student_seq, teacher_tokens)
print(f"sequence-level KD : {seq_kd.item():.4f}")
print(f"plain SFT on teacher tokens: {plain_sft.item():.4f}")
print("identical:", torch.allclose(seq_kd, plain_sft))
sequence-level KD : 2.8593
plain SFT on teacher tokens: 2.8593
identical: True
NoteKey Insight

Word-level KD gives the richest signal (the whole distribution, every step) but requires the teacher’s logits at train time. Sequence-level KD needs only the teacher’s text — cheaper, and it matches the teacher’s sequence behavior rather than its per-token marginals. It is the form behind most “trained on outputs of a bigger model” recipes, R1 distillation included.

On-Policy Distillation (GKD)

Both losses above train on a fixed set of sequences. But at inference the student decodes its own tokens, drifting into states the training set never covered — the exposure bias you met in m08. On-policy distillation (Agarwal et al.’s GKD) fixes this by training on sequences the student itself samples, scored by the teacher. The student is graded exactly where it actually goes.

Sampling on-policy opens a second choice: which divergence to minimize. The forward and reverse KL pull in opposite directions.

from distillation import forward_kl, reverse_kl, js_divergence

p = torch.tensor([0.45, 0.10, 0.45])   # a bimodal teacher
q = torch.tensor([0.34, 0.32, 0.34])   # a spread-out student
print(f"forward KL(p‖q) [mode-covering]: {forward_kl(p, q).item():.4f}")
print(f"reverse KL(q‖p) [mode-seeking] : {reverse_kl(p, q).item():.4f}")
print(f"JS divergence   [symmetric]    : {js_divergence(p, q).item():.4f}")
forward KL(p‖q) [mode-covering]: 0.1360
reverse KL(q‖p) [mode-seeking] : 0.1816
JS divergence   [symmetric]    : 0.0380
  • Forward KL $\mathrm{KL}(p\|q)$ is mode-covering: the student is punished wherever the teacher has mass but the student does not, so it spreads to cover every teacher mode (and smears probability across the valleys between them).
  • Reverse KL $\mathrm{KL}(q\|p)$ is mode-seeking: the student is punished for mass where the teacher has none, so it commits to a single mode. This yields decisive students — often what you want for a chat model.

Fit a single-peaked student to a two-peaked teacher under each divergence and the difference is stark:

TipTry This
  1. Toggle the divergence. Reverse KL parks a narrow student squarely on one peak; forward KL spreads a wide student across both, piling mass into the empty valley between them.
  2. This is why the “right” objective depends on the goal: cover the teacher faithfully (forward / word-level KL) or imitate its decisiveness (reverse / on-policy).

Reasoning Distillation: The R1 Result

In January 2025, DeepSeek trained R1, a strong reasoning model, largely with reinforcement learning (the GRPO lineage of m13). Then they did something that reset expectations for small models: they took R1’s ~800K curated reasoning traces and ran plain supervised fine-tuning (sequence-level KD) on small, open base models — Qwen2.5 and Llama — with no RL stage at all.

The distilled students beat the RL baseline they were compared against:

Model How it was built AIME 2024 (pass@1)
QwQ-32B-Preview large-scale RL 50.0%
R1-Distill-Qwen-1.5B SFT on R1 traces 28.9%
R1-Distill-Qwen-7B SFT on R1 traces 55.5%
R1-Distill-Qwen-32B SFT on R1 traces 72.6%

The paper’s own conclusion is blunt: “distilling more powerful models into smaller ones yields excellent results, whereas smaller models relying on large-scale RL … require enormous computational power and may not even achieve the performance of distillation.” A 1.5B distilled model doing meaningful AIME math, and a 32B distilled model beating a 32B RL model, is dark knowledge at the frontier — the teacher’s reasoning traces are a far richer signal than a scalar reward.

WarningDistillation inherits the teacher’s ceiling — and its flaws

A student trained purely on a teacher can rarely exceed it, and it copies the teacher’s biases and mistakes wholesale. Distillation moves capability down in size, it does not create new capability. To push past the teacher you still need signal the teacher lacks — RL against a verifier (m13), fresh data, or on-policy exploration.

Common Pitfalls

Pitfall Why it bites Fix
Dropping the $T^2$ factor The soft gradient shrinks like $1/T^2$; without the factor the soft term vanishes at high $T$ and $\alpha$ needs retuning per temperature. Keep $T^2$ on the soft term (as in distillation_loss).
Mismatched temperatures Softening the teacher but not the student (or vice-versa) makes the KL compare different distributions. Use the same $T$ for both; divide by $T$ before the softmax.
$T=1$ “distillation” At $T=1$ the soft target is nearly one-hot — you’ve thrown away the dark knowledge you came for. Use $T \in 2..10$; verify the runner-ups are actually visible.
Reading student logits as probabilities of correctness The student learns the teacher’s distribution, calibration flaws included. Evaluate the student’s own task metric, not its agreement with the teacher.
Expecting the student to beat the teacher Pure distillation is bounded by the teacher. Add verifier-RL / fresh data if you need to surpass it.

Exercises

Exercise 1: The temperature sweet spot

Sweep the distillation temperature and watch how much dark knowledge survives. Report the runner-up (truck) probability at each $T$.

from distillation import softmax_T

logits = torch.tensor([0.2, 0.6, 6.0, 3.2, 0.4, 1.1])  # cat dog car truck plane ship
for T in (1.0, 2.0, 4.0, 8.0):
    p = softmax_T(logits, T)
    # Your turn: print T, the top class prob p[2], and the runner-up p[3].
    print(f"T={T}: car={p[2]:.3f}  truck={p[3]:.3f}")
T=1.0: car=0.926  truck=0.056
T=2.0: car=0.660  truck=0.163
T=4.0: car=0.395  truck=0.196
T=8.0: car=0.268  truck=0.189

Exercise 2: Reverse-KL is mode-seeking

Using forward_kl and reverse_kl, confirm that for a bimodal teacher a broad student scores better under forward KL while a narrow, on-mode student scores better under reverse KL.

from distillation import forward_kl, reverse_kl

teacher = torch.tensor([0.45, 0.05, 0.0, 0.05, 0.45])   # two modes at the ends
broad = torch.tensor([0.20, 0.20, 0.20, 0.20, 0.20])    # covers everything
narrow = torch.tensor([0.02, 0.06, 0.02, 0.10, 0.80])   # commits to one mode
# Your turn: compare forward_kl(teacher, ·) and reverse_kl(teacher, ·) for both.
print("forward: broad", round(forward_kl(teacher, broad).item(), 3),
      "narrow", round(forward_kl(teacher, narrow).item(), 3))
print("reverse: broad", round(reverse_kl(teacher, broad).item(), 3),
      "narrow", round(reverse_kl(teacher, narrow).item(), 3))
forward: broad 0.591 narrow 1.098
reverse: broad 5.434 narrow 0.953

Exercise 3: Distill your own tiny student

Call train_student with a harder task (raise noise, lower n_labeled) and confirm the KD student’s advantage over the hard-label student grows as labels get scarcer.

from distillation import train_student

# Your turn: try n_labeled in {16, 24, 48} and compare final_acc for "hard" vs "kd".
for nl in (16, 32):
    h = train_student("hard", n_labeled=nl, seed=0)["final_acc"]
    k = train_student("kd", n_labeled=nl, seed=0)["final_acc"]
    print(f"n_labeled={nl}: hard={h:.3f}  kd={k:.3f}  gap={k - h:+.3f}")
n_labeled=16: hard=0.569  kd=0.669  gap=+0.100
n_labeled=32: hard=0.725  kd=0.819  gap=+0.094

Summary

Key takeaways:

  1. Distillation copies behavior, not weights. A student learns from the teacher’s full probability distribution — the “dark knowledge” a one-hot label discards.
  2. Temperature is the dial. softmax(z/T) softens the distribution; higher $T$ exposes the runner-up classes that carry the similarity structure.
  3. The loss is a hard/soft blend with a $T^2$ factor. $\mathcal{L} = \alpha H + (1-\alpha) T^2 \mathrm{KL}(p\|q)$; the $T^2$ keeps the two gradients comparable across temperatures.
  4. High-$T$ distillation is logit matching — provable and numerically verifiable: the student regresses onto the teacher’s raw logits.
  5. Language models distill two ways. Word-level KD matches the per-token distribution; sequence-level KD is plain SFT on the teacher’s generated tokens.
  6. On-policy distillation (GKD) trains on the student’s own samples to kill exposure bias, and the choice of divergence trades mode-covering (forward KL) for mode-seeking (reverse KL).
  7. R1 reset the small-model playbook. SFT on a strong reasoner’s traces beat large-scale RL on the same small model — distillation moves frontier capability down in size.

What’s Next

You have now built the language model from tensors all the way to the current frontier — architecture, training, alignment, reasoning, efficiency, and the compression and knowledge-transfer techniques (quantization, merging, and now distillation) that put a frontier model on real hardware. Circle back to Module 13: Reasoning to see the RL side of the distillation-vs-RL trade, or Module 16: Speculative Decoding where a small (often distilled) draft model makes a large one faster.

Going Deeper

Core Papers:

Practical Resources: