Model Distillation

Model Distillation

Distillation trains a small student $q_{\theta }$ to reproduce the behaviour of a large teacher $p$. Here we revise it for language models: what the loss is, which divergence to use, where the training text should come from, what you can do with only API access, and why the whole thing works. The note follows a tutoring session (2026-11-13, 2026-11-14); the derivations are the ones done there, and the last section lists the slips made on the way.

The thing to remember: almost every design choice reduces to three questions.

• Which divergence? Forward or reverse KL.
• Whose text? Teacher-written, student-written, or a mix.
• What does the teacher expose? Text, probabilities, or internals.

The basic loss

Gradient of cross-entropy with respect to the logits

Student logits $z$, $q=\mathrm{softmax}(z)$, target distribution $p$, loss $L=-{\sum }_i{p}_i\log q_i$.

Since $\log q_j=z_j-\log {\sum }_k{e}^{z_k}$, the Jacobian of the log-softmax is

$$\frac{\partial \log q_j}{\partial z_i}={\delta }_{ij}-q_i$$

The probability that appears belongs to the logit you differentiate with respect to ($i$), not to the token in the numerator ($j$). Then

$$\begin{array}{rl}\frac{\partial L}{\partial z_i}& =-\sum _j{p}_j({\delta }_{ij}-q_i)\\ & =-p_i+q_i\sum _j{p}_j\\ & =q_i-p_i\end{array}$$

Hard label versus soft target

Take $q=(0.7,0.2,0.1)$, hard label on token 1, teacher $p=(0.6,0.39,0.01)$.

Target	Gradient $q-p$	Token 2	Token 3
Hard $(1,0,0)$	$(-0.3,+0.2,+0.1)$	pushed down	pushed down
Soft $(0.6,0.39,0.01)$	$(+0.1,-0.19,+0.09)$	pushed up	pushed down

With a hard label every wrong logit is pushed down in proportion to the student's own $q_i$: the signal contains nothing the student did not already believe. The soft target changes the sign of the update on token 2. That is the precise meaning of "dark knowledge".

Temperature

A confident teacher, e.g. $p=(0.98,0.015,0.005)$, gives a soft target that is nearly one-hot, so the information about wrong tokens is numerically invisible. Divide the logits by $T>1$ before the softmax to flatten the distribution:

$$p^T=\mathrm{softmax}(v/T),q^T=\mathrm{softmax}(z/T)$$

• Apply $T$ to both sides. If only the teacher is softened, the student at $T=1$ learns a flattened distribution and is underconfident at inference.
• Matching at $T$ implies matching at $T=1$. $\mathrm{softmax}(z/T{)}_i\propto \mathrm{softmax}(z{)}_i^{1/T}$, so softening is a renormalised $T$-th root. If $q_i^{1/T}/Z_s=p_i^{1/T}/Z_t$ for all $i$, raising to the power $T$ gives $q\propto p$, and two normalised distributions that are proportional are equal. $T$ only changes which errors the loss emphasises during training.

The $T^2$ factor

With $u=z/T$, the chain rule gives $\partial u_i/\partial z_i=1/T$, hence

$$\frac{\partial L}{\partial z_i}=\frac{1}{T}\left(q_i^T-p_i^T\right)$$

The difference shrinks too. For large $T$, use $e^x\approx 1+x$ on numerator and denominator ($N$ tokens, mean logit $\bar{z}$):

$$\begin{array}{rl}\mathrm{softmax}(z/T{)}_i& \approx \frac{1+z_i/T}{N(1+\bar{z}/T)}\\ & \approx \frac{1}{N}+\frac{z_i-\bar{z}}{NT}\end{array}$$

So $q_i^T-p_i^T\approx \frac{(z_i-\bar{z})-(v_i-\bar{v})}{NT}$ and the gradient scales as $1/T^2$.

The mixed objective

$$L=\alpha \mathrm{CE}(\text{hard label},q)+(1-\alpha )T^2\mathrm{CE}(p^T,q^T)$$

Both terms are cross-entropies; they differ in target and temperature. Forgetting $T^2$ at $T=4$ makes the distillation term about 16 times weaker than intended, and the student behaves almost like one trained on hard labels only. Introduced in Hinton's Distilling the Knowledge in a Neural Network.

Choice of divergence

Forward versus reverse KL

	Forward $KL(p\parallel q)$	Reverse $KL(q\parallel p)$
Formula	${\sum }_i{p}_i\log (p_i/q_i)$	${\sum }_i{q}_i\log (q_i/p_i)$
Weighted by	$p$ (teacher)	$q$ (student)
Forbids	$q\to 0$ where $p>0$ (term blows up)	$q>0$ where $p\approx 0$ (term blows up)
Does not see	mass placed where $p\approx 0$ (term weighted by $p$)	modes dropped by the student (term weighted by $q\to 0$)
Small student	one wide peak over all modes, mass in the gaps	commits to one mode, stays out of the gaps
Sampled text	varied, drifts into incoherence	coherent, low variety
Name	mass-covering, mean-seeking	mode-seeking, zero-forcing

• Soft-target cross-entropy is forward KL: $-{\sum }_i{p}_i\log q_i=KL(p\parallel q)+H(p)$ and $H(p)$ does not depend on the student.
• Concrete forward failure: teacher has $0.99$ on three plausible tokens, the student manages $0.90$. Over 200 sampled tokens the student emits about 20 implausible ones against about 2 for the teacher.
• Concrete reverse failure: teacher has $(0.5,0.3,0.19)$ on three plausible tokens; a student with all mass on the first pays only $\log 2\approx 0.69$. Dropping plausible options is cheap, so about half of the teacher's variety at that position disappears.

Which KL can you compute?

Over whole sequences:

$$\begin{array}{rl}KL(p\parallel q)& ={\mathbb{E}}_{y\sim p}\left[\log p(y)-\log q(y)\right]\\ KL(q\parallel p)& ={\mathbb{E}}_{y\sim q}\left[\log q(y)-\log p(y)\right]\end{array}$$

Setting	What you can do with $p$	Easy divergence
Variational Inference	evaluate up to a constant, cannot sample	reverse
Distillation	can sample (the teacher is a model you run)	forward

• Forward needs only teacher text. $\log p(y)$ has no student parameters, so it drops out; what is left is $-\log q(y)$ on teacher samples, i.e. ordinary hard-label training. A dataset generated once stays valid because the teacher is frozen.
• Reverse needs fresh samples from the current student at every step (a stored set is stale after one update) and the teacher's log-probabilities on them. It is an on-policy method.

Token level versus sequence level

Same mean, different variance

On a teacher-written prefix you can train with (a) the token the teacher sampled, one-hot, or (b) the teacher's full next-token distribution.

• (a) gradient $q-e_Y$ with $Y\sim p$; (b) gradient $q-p$.
• $\mathbb{E}[e_Y]=p$, so the two gradients are equal in expectation. (a) is a single-sample estimate of (b).

Coordinate $i$ of $e_Y$ is a Bernoulli variable with parameter $p_i$ (note $1^2=1$, so $\mathbb{E}[1^2]=p_i$):

$$\begin{array}{rl}\mathrm{Var}(q_i-1[Y=i])& =p_i(1-p_i)\\ \text{total variance}& =\sum _i{p}_i(1-p_i)=1-\sum _i{p}_i^2\end{array}$$

• Zero exactly when the teacher is certain ($p$ one-hot): hard labels lose nothing there.
• Largest for uniform $p$: $1-1/N$.
• Across a dataset both methods share the variance from which prefixes were drawn; (a) adds $\mathbb{E}[1-{\sum }_i{p}_i^2]$ on top. That is the entire gap.

Why hard labels still work

Each long prefix occurs once, so "50k samples per prefix" would make sequence-level distillation impossible. It is not: the weights are shared, the gradient is averaged over all positions, and what is learned at one prefix transfers to similar ones. Sequence-level distillation from teacher text (Kim and Rush's Sequence-Level Knowledge Distillation) is the most common form in practice.

Exposure bias and on-policy distillation

Compounding errors

Training prefixes come from the teacher or from data; inference prefixes come from the student. If each token leaves the training distribution with probability $\epsilon$ and never returns, a generation of $n$ tokens survives with probability $(1-\epsilon {)}^n$. For $\epsilon =0.01$, $n=500$: $0.99^{500}\approx 0.7\%$. A student that is 99% reliable per token almost never finishes on-distribution.

• The argument appears for compositional tasks in Dziri's Faith and Fate; the rigorous imitation-learning version is Ross and Bagnell's.
• The assumption "once off, always off" is often false: a model can recover if it has been trained on what to do after a mistake.

On-policy distillation

The student writes the training sequences; the teacher supplies its next-token distribution at every position of them.

• At position $t+1$ after a bad token at $t$, the teacher's distribution is conditioned on the flawed prefix, so it describes the best continuation from that bad state. The student gets recovery targets for exactly the mistakes it makes.
• Training is still one parallel pass: average the per-position cross-entropy over the student-written sequence.
• This is DAgger applied to language models. See Agarwal's On-Policy Distillation of Language Models (GKD).
• Injected or templated errors are a weaker substitute: you have to guess what the student's mistakes look like.

Mixing teacher and student prefixes

Pure on-policy training from a weak student fails: its text is garbage the teacher has never seen, so the teacher's distributions on it are unreliable.

Method	Prefix source	Note
Supervised KD	fixed teacher or data text	train/inference mismatch
On-policy (GKD)	student	unreliable teacher targets early; sensitive to student initialisation
Teacher-mixed sampling (MiniLLM)	blend of both next-token distributions at each step	reverse-KL objective; mixing stops the student gaming the teacher with degenerate text
Interleaved sampling (SKD)	student proposes, teacher replaces tokens outside its top-$K$	behaves like supervised KD early, like on-policy later, with no hand-tuned schedule

Black-box distillation

What each level of access allows

Teacher exposes	Available	Ruled out
Text only	sequence-level (forward KL, hard labels)	soft targets, reverse KL, on-policy scoring, internals
Text + reasoning	sequence-level on rationale + answer	same as above
Log-probabilities	token-level soft targets, on-policy, reverse KL	internals
Weights	everything, including hidden states and attention	nothing

• Hiding reasoning text blocks sequence-level distillation of the reasoning. Withholding log-probabilities blocks token-level and on-policy methods. They are different protections.
• Visible reasoning does not enable on-policy distillation: the teacher wrote that text.

Rationales

Training on reasoning followed by answer supervises many more tokens per example and makes the answer easier to predict from what precedes it. Hsieh's Distilling Step-by-Step shows small models need far fewer examples when rationales are an extra target. The cost: outputs are much longer, so by $(1-\epsilon {)}^n$ a small student has more chances to derail before the answer.

Protecting against distillation

• Providers hide raw chains of thought partly for this reason (competitive advantage is on OpenAI's own list of reasons).
• Savani's Antidistillation Sampling keeps traces visible but perturbs sampling to poison them: in their MMLU example a 72% teacher yields a 52% student from ordinary traces and about 40% from poisoned ones.
• Hiding is not airtight; 2026 reports describe raw traces recovered through provider APIs.

Building a dataset without prompts

• Write a few seed prompts, have the teacher generate more, then have it answer them.
• Generated prompts are narrower than real ones: every call samples from the same high-probability region.
• Counters: higher temperature, a different random subset of seeds per call, conditioning on a sampled topic or style, removing near-duplicates. Detection: compare token statistics with a sample of real prompts.

One sample or many per prompt?

• Default: more distinct prompts. Ten answers to one prompt overlap; ten prompts cover ten contexts that shared weights can pool.
• Extra samples are informative only where the teacher's answers vary ($1-\sum p_i^2$ large): open-ended tasks, several valid solutions, genuine uncertainty.
• Agreement among samples is an empirical confidence estimate without log-probabilities; use it to filter or down-weight.

Intermediate representations

Hidden-state matching

Student state $h_t^S\in {\mathbb{R}}^{d_S}$, teacher state $h_t^T\in {\mathbb{R}}^{d_T}$, learned projection $W\in {\mathbb{R}}^{d_T\times d_S}$:

$$L_{\text{hid}}=\frac{1}{n}\sum _t{\Vert W{h}_t^S-h_t^T\Vert }^2$$

• Project the student up, not the teacher down. $\Vert h^S-W{h}^T{\Vert }^2$ with $W$ and the student both free has the trivial solution $W=0$, $h^S=0$. Projecting up has a fixed, input-dependent target, so one constant vector cannot fit it.
• Used as an auxiliary loss next to output distillation, not alone.

What the loss constrains

• Basis-free. If student B has states $R{h}_t$ for an invertible $R$, it reaches the same loss with $W_B=W_A{R}^{-1}$. No individual neuron means anything to the loss.
• Linear recoverability. The loss is small iff the teacher's state can be reconstructed from the student's by a linear map. Information present only in a nonlinear encoding counts as lost.
• PCA floor. $W{h}^S$ varies in at most $d_S$ dimensions. With (centred) teacher variances ${\lambda }_1\ge \cdots \ge {\lambda }_{d_T}$ along principal directions, the best case captures the top $d_S$ directions and the loss cannot go below ${\sum }_{i>d_S}{\lambda }_i$.

Layer mapping

Residual view: $h_{l+1}=h_l+f_l(h_l)$. With a 24-layer teacher and a 6-layer student, the uniform rule pairs student layer $m$ with teacher layer $4m$, so one student block must approximate a composition of four teacher blocks.

$$L=\sum _m{\Vert W_m{h}_m^S-h_{\varphi (m)}^T\Vert }^2$$

Mapping	$\varphi$	Assumption
Uniform (TinyBERT)	$\varphi (m)=4m$	the small network computes things in the same order and proportion
Last layer only (MiniLM)	one term	none on intermediate layers; they organise freely
Soft	target ${\sum }_j{\alpha }_{mj}{h}_j^T$, $\alpha$ learned	the student chooses which teacher layers to track

Attention matching

One head: $Q=X{W}_Q$, $K=X{W}_K$ with $X\in {\mathbb{R}}^{n\times d}$, $W_Q,W_K\in {\mathbb{R}}^{d\times d_k}$, and $A=\mathrm{softmax}(Q{K}^{\top }/\sqrt{d_k})$.

• $W_Q,W_K$ are weights and depend on the width. The map $A$ is $n\times n$ whatever $d$ and $d_k$ are; row $t$ is a distribution over positions.
• So no projection is needed: compare maps directly, by row-wise cross-entropy $-{\sum }_j{a}_{t,j}^T\log a_{t,j}^S$ or by squared error.
• Remaining mismatch: head count. Heads have no natural order, so pairing head $k$ with head $k$ is arbitrary. MiniLMv2 concatenates queries and keys across heads and re-splits them into the same number of groups on both sides.

A complete white-box loss

$$\begin{array}{rl}L=& -T^2\sum _i{p}_i^T\log q_i^T\\ & +\beta \sum _m\sum _t{\Vert W_m{h}_{m,t}^S-h_{4m,t}^T\Vert }^2\\ & +\gamma \sum _m\sum _{\text{heads}}\sum _{i,j}{\left(a_{ij}^{T,4m}-a_{ij}^{S,m}\right)}^2\end{array}$$

Practical issues

The capacity gap

A much stronger teacher can give a worse student. Its distributions depend on distinctions the student cannot represent; under forward KL the student averages over them and smears mass onto tokens that are wrong in each specific context, and those positions seed compounding errors. A weaker teacher's distributions can be fitted without smearing.

Remedies:

• Reverse KL on student-written text: the student commits to the modes it can represent.
• Simplify the target: filter or truncate the teacher's distribution.
• Teacher assistant (Mirzadeh's Improved Knowledge Distillation via Teacher Assistant): distil 70B into 7B, then 7B into 1B. Each gap is smaller, but the copy errors chain.

Toy model of the trade-off (invented numbers): a step across ratio $r$ keeps $f(r)$ of performance; $k$ equal steps keep $f(70^{1/k}{)}^k$.

$k$	ratio per step	$f$	kept overall
1	70	0.80	0.80
2	8.4	0.93	0.86
4	2.9	0.97	0.89
8	1.7	0.98	0.85

Interior optimum: a few intermediate models. Each one also costs a full distillation run.

A two-teacher variant

Idea from the session: train the 1B student on 70B-written text with soft targets from a 7B model that was not itself distilled from the 70B one. This is the mixed objective with two teachers: hard labels from the 70B model (unbiased, noisy, the only route for its knowledge) and soft targets from the 7B model (low variance, biased towards what the 7B model already believes). Where the teachers disagree the two terms conflict, and those are the positions where the 70B model knows more. Open question whether it beats the chain; not checked against the multi-teacher literature.

Tokenizer mismatch

Anything that passes through plain text survives; anything defined per token or per position breaks.

Method	Different tokenizers?	Why
Token-level soft targets	breaks	distributions over different vocabularies
Sequence-level on teacher text	works	re-tokenise the teacher's string with the student's tokenizer
On-policy, per-token scoring	breaks	positions and vocabularies do not line up
On-policy, whole-string $\log p(y)$	works	the teacher scores the string with its own tokenizer; one sparse number per sequence
Hidden states	breaks	position $t$ covers different text; needs alignment (e.g. pooling over words)
Attention maps	breaks	$n\times n$ with different $n$

Dense signal without matching vocabularies: Minixhofer's Universal Cross-Tokenizer Distillation via Approximate Likelihood Matching compares the likelihood of chunks of text that both tokenizations delimit identically. See Tokenization.

Why distillation works

Targets as estimators of the true distribution

Let $p^{\ast }$ be the true next-token distribution. Both the human token and the teacher's distribution $p_T$ estimate it.

$$\mathbb{E}\Vert \text{target}-p^{\ast }{\Vert }^2=\Vert \text{bias}{\Vert }^2+\text{variance}$$

Target	Bias	Variance	Error for $p^{\ast }=(0.5,0.3,0.2)$
Human token $e_Y$	0	$1-\sum p_i^{\ast 2}$	0.62
Label smoothing $0.9e_Y+0.1u$	$0.1(u-p^{\ast })$, squared norm 0.0005	$0.81\times 0.62$	0.50
Teacher $p_T=(0.6,0.25,0.15)$	$\mathbb{E}[p_T]-p^{\ast }$	small	0.015

• The teacher saw the same single token at this prefix. Its low variance comes from pooling many similar contexts through shared weights; averaging $m$ samples divides variance by $m$.
• The price is bias, from treating similar contexts alike and from the teacher's own errors. It is visible as the teacher's excess loss on held-out text: $H(p^{\ast })+KL(p^{\ast }\parallel p_T)$.
• Label smoothing also reduces variance, but with a context-independent vector: it adds no information. The gain of distillation is context-specific information, not smoothing. See Bias Variance Trade-off, Entropy.

Three reasons, then: information about unsampled tokens, lower-variance targets, and both bought with a bias.

Connections

• Variational Inference: same two divergences, opposite one is tractable, because there you cannot sample from $p$.
• Imitation learning: exposure bias is behaviour cloning's compounding error; on-policy distillation is DAgger. See Inverse Reinforcement Learning.
• Quantal Response Equilibrium: the temperature softmax is the logit response function; $T$ plays the role of the inverse rationality parameter.
• Ensemble Methods: the original motivation was compressing an ensemble into one network.
• Low-Rank Adaptation: both are rank arguments. LoRA limits the rank of a weight update; the projection $W$ limits the rank of the reconstruction of the teacher's activations.
• LLM-to-LLM alignment in social simulation: fitting one model's action distribution to another's is distillation of a policy. The divergence decides whether mixed strategies survive (forward) or collapse to a pure one (reverse).

Personal pitfalls

Slips made during the session, with the correction. Review these first.

Slip	Correction
Wrote ${\sum }_{j\ne i}{p}_j{q}_j$ in the gradient; later gave hard-label gradients $0.7,0.7$	$\partial \log q_j/\partial z_i=-q_i$: the index of the logit, so the sum is $q_i(1-p_i)$
"Higher temperature gives bigger gradients"	Gradient scales as $1/T^2$; hence the $T^2$ factor
Left out $T^2$ twice; wrote $(p-q)T^2$ as the loss	$(q-p)$ is the gradient. Loss: $-T^2\sum p^T\log q^T$
Attached "mean" to reverse KL (twice)	Forward is mean-seeking; reverse is mode-seeking
"Soft targets minimise reverse KL"	The sum is weighted by $p$: forward KL plus $H(p)$
"Forward KL is hard to compute"	True in VI only. In distillation you can sample the teacher, so forward is the easy one
"A stored dataset suits reverse KL"	Reverse needs samples from the current student; stored ones are stale after one step
ABAB teacher, all-A reverse-KL student	That student has infinite reverse KL; reverse only loses where $p$ is multi-modal at one position
"Visible reasoning enables on-policy distillation"	Teacher wrote the text; on-policy needs teacher log-probs on student text
Fix for exposure bias: inject template errors	The student should write the prefixes; its own mistakes are the relevant ones
"Sequence-level is the average of token-level"	It is the single-sample version, equal in expectation
"Averaging many samples makes the estimate unbiased"	Averaging reduces variance by $1/m$; bias is separate
Compared only bias${}^2$ of label smoothing with the teacher's error	Total error is bias${}^2$ + variance
"Extra samples help where the teacher always answers the same"	Variance is 0 there; they help where answers vary
Trivial solution assigned to "project student up"	It is "project teacher down": $W=0$, $h^S=0$
Attention map is $d\times d$, depends on width	Weights depend on width; the map is $n\times n$
Tokenizer mismatch: sequence-level breaks, hidden and attention survive	The reverse: text survives, per-position quantities break
"$\log p(y)$ of a string is not defined across tokenizers"	The teacher tokenises the string itself; it is defined, just sparse
Uniform total variance $1-1/N^2$; layer step 3 for 24 vs 6	$1-1/N$; step 4

References

Checked during the session:

• Agarwal et al., On-Policy Distillation of Language Models: Learning from Self-Generated Mistakes, ICLR 2024.
• Gu, Dong, Wei, Huang, MiniLLM: Knowledge Distillation of Large Language Models, ICLR 2024.
• Xu et al., Speculative Knowledge Distillation: Bridging the Teacher-Student Gap Through Interleaved Sampling, ICLR 2025.
• Hsieh et al., Distilling Step-by-Step!, Findings of ACL 2023.
• Minixhofer, Vulić, Ponti, Universal Cross-Tokenizer Distillation via Approximate Likelihood Matching, NeurIPS 2025.
• Savani et al. (Kolter group), Antidistillation Sampling, arXiv 2504.13146. Venue not checked.

Cited from memory, venue not re-checked:

• Hinton, Vinyals, Dean, Distilling the Knowledge in a Neural Network, 2015.
• Kim and Rush, Sequence-Level Knowledge Distillation, EMNLP 2016.
• Ross, Gordon, Bagnell, A Reduction of Imitation Learning and Structured Prediction to No-Regret Online Learning (DAgger), AISTATS 2011.
• Dziri et al., Faith and Fate: Limits of Transformers on Compositionality, NeurIPS 2023.
• Jiao et al., TinyBERT, Findings of EMNLP 2020.
• Wang et al., MiniLM, NeurIPS 2020; MiniLMv2, Findings of ACL 2021.
• Mirzadeh et al., Improved Knowledge Distillation via Teacher Assistant, AAAI 2020.

Flagged, not peer-reviewed (used only for the claims about hidden reasoning):

• Simon Willison, notes on OpenAI o1, quoting OpenAI's stated reasons (2024).
• eWeek on o3-mini chain-of-thought summaries (2025).
• Xint blog on leaked reasoning traces; Cloud Security Alliance research note (2026).

Review checklist

• Derive $\partial L/\partial z_i=q_i-p_i$ without the index slip.
• Hard vs soft gradient on the $(0.7,0.2,0.1)$ example: which logit changes sign.
• Why $T$ goes on both sides, and why matching at $T$ implies matching at $T=1$.
• High-$T$ expansion of the softmax and the $1/T^2$ gradient scaling.
• Forward vs reverse KL: what each forbids, what each ignores, what the text looks like.
• Which KL is tractable in VI and which in distillation, and why.
• Hard-label gradient variance $1-\sum p_i^2$; when it is zero and when maximal.
• $(1-\epsilon {)}^n$ and what on-policy distillation changes.
• Why pure on-policy from a weak student fails; how SKD and teacher-mixed sampling fix it.
• Access table: text, reasoning, log-probs, weights.
• Hidden-state loss: projection direction, basis invariance, PCA floor.
• Attention map shape and why no projection is needed.
• Capacity gap mechanism and the teacher-assistant trade-off.
• Tokenizer mismatch: what survives and what breaks.
• Bias-variance comparison of human token, label smoothing, teacher target.