Distillation: logit KL, temperature, and on-policy vs off-policy

A 4.5B model can’t learn everything a 27B one knows, but it can learn a lot more from the 27B model than from raw text. The teacher’s outputs carry information that a single “correct” token doesn’t. Distillation is how you move it. Here’s the math, and the choices that decide whether it works.

Soft targets carry the information

Training on text gives one correct next token per position, a one-hot target. A teacher model gives a full distribution: “def” 70%, “class” 20%, “async” 8%. Those runner-up probabilities encode what Hinton et al. (2015) called dark knowledge: which wrong answers are nearly right.

To expose more of it, both models’ logits are softened by a temperature T before comparing them:

p_teacher = softmax(z_teacher / T)      p_student = softmax(z_student / T)
L_distill = T² · KL( p_teacher ‖ p_student )

The T² factor keeps gradient magnitudes comparable across temperatures, because softening scales the gradients down by 1/T². In practice the distillation term is mixed with ordinary cross-entropy on the real tokens.

Which way the KL points matters

KL divergence isn’t symmetric, and the direction changes what the student does with the probability mass it can’t afford to match:

• Forward KL, KL(teacher ‖ student): the student is penalised wherever the teacher has mass and the student doesn’t. It spreads out to cover everything the teacher might say, the mode-covering behaviour. For a small student, that can mean putting probability on answers it can’t actually produce well.
• Reverse KL, KL(student ‖ teacher): the student is penalised for putting mass where the teacher doesn’t. It concentrates on the teacher’s main modes and drops the rest, the mode-seeking behaviour. MiniLLM (Gu et al., 2024) argues this suits generative students better.

For a capacity-limited student, reverse KL often gives cleaner generations. Forward KL keeps more diversity. Pick deliberately.

Logits or sequences

• Logit distillation matches the full distribution at every position. It carries the most information, but the teacher and student must share a tokenizer, or the distributions don’t line up.
• Sequence-level distillation (Kim and Rush, 2016) trains the student on text the teacher generated. It works across different tokenizers and is just supervised fine-tuning on teacher outputs, but it keeps only the one sampled path and throws away the distribution.

Storing logits is the hidden cost. A full distribution over a vocabulary of about 152K tokens is 304 KB per token at FP16. For a billion training tokens, that’s 304 TB. The usual fixes are to keep only the top-k probabilities (top-64 is about 384 bytes per token, or 384 GB per billion tokens) or to run the teacher online during training.

Off-policy vs on-policy

Both methods above are off-policy: the student learns on sequences it didn’t write. At inference it generates its own text, makes its own mistakes, and ends up in states it never trained on. It doesn’t know how to recover from errors it was never shown.

On-policy distillation fixes the mismatch. The student generates, and the teacher scores every token of what the student wrote. Generalised Knowledge Distillation (GKD, Agarwal et al., 2024) is the standard formulation. The student learns how the teacher would continue from the student’s own mistakes, which is exactly what it needs at inference. The cost is generation during training, and a teacher in memory or reachable at every step.

Distilling reasoning: what DeepSeek-R1 showed

The best-known distillation result of the last few years is DeepSeek’s. Alongside R1 they released six smaller models, Qwen and Llama bases from 1.5B to 70B parameters. Each was fine-tuned on about 800K samples generated by R1, with plain sequence-level distillation and no RL on the student.

Two findings stood out:

• Reasoning transfers through text. The distilled models picked up long chain-of-thought behaviour from R1’s traces alone. DeepSeek reported R1-Distill-Qwen-32B ahead of OpenAI’s o1-mini on several reasoning benchmarks.
• Distilling beat RL for the small model. Running large-scale RL directly on the 32B base did markedly worse than distilling R1 into it. A strong teacher’s traces were worth more than the student discovering reasoning on its own at that compute.

The practical lesson: if a capable teacher exists, distil first and spend RL on what’s left (the reasoning note and the DeepSeek note cover the rest of that story).

How this fits my compass model

The plan for compass is to take a fine-tuned Qwen with strong terminal and agent behaviour, distil a larger Qwen into 4.5B variants, and then run RL per language and stack. The RL note covers that part. Distillation is the middle step: it moves capability down to a size that ships inside an IDE, before RL sharpens it on tasks where tests decide the reward. My distillation tooling is in the distiller repo.

Rules I use

1. Same tokenizer, logit distillation; different tokenizers, sequence-level. Don’t fight the vocabulary mismatch.
2. Store top-k, not full distributions, unless you can run the teacher online.
3. Finish on-policy. Off-policy gets the student most of the way cheaply; on-policy teaches it to recover from its own errors.
4. Evaluate on generations, not loss. A lower KL on teacher text doesn’t guarantee better outputs when the student writes alone.

Sources: Hinton, Vinyals and Dean, Distilling the Knowledge in a Neural Network (2015). Kim and Rush, Sequence-Level Knowledge Distillation (2016). Gu et al., MiniLLM (ICLR 2024). Agarwal et al., On-Policy Distillation of Language Models (GKD, ICLR 2024). DeepSeek-AI, DeepSeek-R1 (2025). Storage figures are computed for a 152,064-token vocabulary.