Distillation quietly flipped: stop training on the teacher's answers

There is a version of knowledge distillation that everybody learns first, and it is the one almost every distilled model on the market was built with. Take a big model. Have it write a few hundred thousand high-quality responses. Fine-tune a small model on that corpus with ordinary cross-entropy. Ship the small model.

It works well enough to be everywhere. It also has a specific, predictable failure: the resulting model is good for a sentence or two and then drifts — the reasoning goes slightly sideways, a citation appears that should not, the structure it started with quietly dissolves.

That failure is not about the teacher’s quality, and it is not about the student being too small. It is about which situations the student was ever supervised in.

The states nobody trained on

A language model’s input at step t includes everything it generated up to step t. That is the whole architecture. Which means the inputs it faces in production are inputs it produced itself.

Now look at what the standard recipe teaches. Every single training example has a prefix written by the teacher. The student learns one skill very well: given text a strong model would have written, predict the next token.

Then you deploy it and it has to write its own prefix.

The first token that is slightly off — not wrong, just not what the teacher would have chosen — puts the model on a prefix that appears nowhere in its training distribution. It is now predicting from a state it has no supervision for. Its next token is a little worse. That compounds. By the third paragraph it is somewhere the teacher would never have been, doing its best with no guidance at all.

Imitation learning hit this wall years ago and named it. The student policy’s errors compound because it is trained on the expert’s state distribution and evaluated on its own. The fix there was DAgger: roll out the learner, then ask the expert what it would have done in the states the learner actually reached.

On-policy distillation is that idea, applied to tokens.

The swap is small; the consequence is not

The change is three lines of pseudocode:

off-policy:  batch = teacher_corpus.sample()
             loss  = cross_entropy(student(batch), batch.targets)

on-policy:   rollout = student.generate(prompt)          # student writes it
             target  = teacher.forward(rollout)          # teacher grades it
             loss    = kl(student(rollout), target)       # match on those tokens

The teacher never generates anything. It scores a finished sequence in a single forward pass, which is the cheap direction — no autoregressive loop, fully parallel across positions.

What you get for that is supervision aimed precisely at the model’s weaknesses, because the tokens being corrected are the tokens it just produced. Every gradient is spent on a state the student will genuinely be in.

Forward KL makes small models hedge

There is a second change that comes along with the first, and it is easy to miss.

Cross-entropy on teacher text is minimising forward KL, KL(teacher ‖ student). That objective is mode-covering: it punishes the student severely for assigning near-zero probability to anything the teacher considered plausible. So the student learns to keep a little mass on every branch.

For a model with room to represent all those branches, fine. For a 3B student imitating a frontier teacher, it is asking for something the student cannot deliver. The result is a model that spreads probability thinly across several continuations and commits to none of them cleanly — which, at generation time, reads exactly like drift.

On-policy setups typically optimise reverse KL, KL(student ‖ teacher) — mode-seeking. It punishes the student for putting mass anywhere the teacher thinks is unlikely, and does not care about teacher modes the student ignores entirely.

The practical translation:

• Forward KL: “be able to say anything the teacher might say.”
• Reverse KL: “say one thing the teacher would approve of, properly.”

For a capacity-limited student, the second request is both easier and more useful. This is the part of the story that is genuinely about the objective and not just the data pipeline, and it is why the two changes are usually made together.

What it costs, plainly

On-policy distillation is more expensive per optimiser step, and it is worth being specific about where the money goes.

	Off-policy	On-policy
Training data	Built once, reused every epoch	Regenerated continuously
Per step	Student forward + backward	Student generation, then student forward + backward, plus one teacher forward
Teacher at train time	Not needed	Resident in memory next to the student
Infrastructure	A dataset and a trainer	A sampling loop coupled to the trainer

The dominant cost is student generation, not the teacher. People assume keeping a frontier teacher in memory is the expensive part; in practice it scores in one pass while the student is the thing running an autoregressive loop.

Against that, you need far fewer tokens of supervision, because none of them are wasted on states the student will never visit. The honest summary is that you have converted a data cost into a compute cost, and for a model you intend to actually deploy, that is usually the better trade.

The recipe people actually use

Nobody runs on-policy distillation from a base model, because early student samples are too poor to be worth grading.

1. SFT first. Ordinary off-policy distillation on teacher text, until the student is coherent. This is the right tool for getting off the ground and there is no reason to skip it.
2. Then switch. Once samples are plausible, start rolling out the student and scoring with the teacher. This is where the compounding error actually gets fixed.
3. Keep some teacher text in the mix. A small off-policy fraction stops the student collapsing into a narrow self-reinforcing style — the same entropy problem that shows up in RL fine-tuning.
4. Sample at your serving temperature. Training on greedy rollouts and then serving at temperature=0.8 puts back the distribution gap you just spent compute closing. This one catches people repeatedly.

Where it stops

On-policy distillation sits between supervised fine-tuning and reinforcement learning, and the difference from RL is worth stating because they get conflated.

RL from verifiable rewards hands back a single scalar per rollout — the answer was right, or it was not — and leaves the model to work out which of six hundred tokens deserved the credit. On-policy distillation hands back a full distribution at every position: not “that was wrong” but “at token 41 you chose therefore; a stronger model would have put most of its mass on however.”

That is a vastly denser signal, which is why distillation converges in a fraction of the steps. The trade is the ceiling. The student is chasing a teacher, so it cannot exceed one. RL has no teacher and therefore no ceiling — and is correspondingly slower and noisier.

Which is why serious post-training stacks run them in sequence rather than picking a side: distil to get most of the teacher’s capability cheaply, then use RL where you want to go past it.

The thing worth remembering

If you take one idea from this: a model is only reliable in the states it was supervised in, and it generates its own states. That sentence explains the drift, explains why more teacher data does not fix it, and explains why sampling from the student — the thing that sounds like a small engineering detail — is the part that matters.