On-policy distillation

Standard knowledge distillation has a shape almost everyone learns first. Take a large teacher, have it generate a pile of high-quality outputs, then fine-tune a small student on that pile with ordinary cross-entropy. It works. It is also the reason so many distilled models feel fine for a sentence and then quietly come apart.

The problem is not the teacher’s quality. It is which states the student is supervised on.

The distribution the student never trains on

A language model generates token by token, and each token it emits becomes part of the input for the next one. So the inputs a model sees at inference are inputs it generated itself.

Now look at what off-policy distillation actually teaches. Every training prefix comes from the teacher. The student only ever learns “given a prefix a strong model would have written, what comes next?”

At inference, the student writes its own prefix. The first slightly awkward token puts it on a prefix the teacher would never have produced — which is precisely the region where it has no supervision at all. The next prediction is a little worse, which pushes it further off-distribution, and the error compounds.

This is exposure bias, and it is the same compounding-error problem imitation learning has wrestled with for years. The fix there was DAgger: stop training on the expert’s states, start training on the learner’s states and ask the expert what it would have done.

The swap

On-policy distillation keeps the teacher but changes who writes the training data:

1. Sample a completion from the student for a prompt.
2. Run the teacher over that completion — one forward pass, no generation — to get its next-token distribution at every position.
3. Train the student to match the teacher’s distribution on its own tokens.

The student is now corrected exactly where it is weak, because the tokens being corrected are the ones it actually produced. Nothing about the teacher changed. What changed is that supervision follows the student around instead of waiting for it in a place it never goes.

You will see this called generalized knowledge distillation in the literature, and it is the same idea as DAgger transplanted into sequence modelling.

Forward KL, reverse KL, and what you are asking for

Cross-entropy against teacher text is minimizing forward KL, KL(teacher ‖ student). Forward KL is mode-covering: it is heavily penalised when the student assigns near-zero probability to something the teacher finds plausible. So the student spreads mass to cover every branch the teacher might take — including branches a small model cannot follow through on.

On-policy setups usually optimise reverse KL, KL(student ‖ teacher), which is mode-seeking. It punishes the student for putting mass anywhere the teacher considers unlikely, and is indifferent to teacher modes the student ignores.

That asymmetry matters more than it sounds:

• Forward KL gives you a student that hedges. It keeps a little probability on several continuations, which for a capacity-limited model means it commits to none of them cleanly.
• Reverse KL gives you a student that picks one thing the teacher endorses and does it properly.

For a student much smaller than its teacher, the mode-covering objective is asking for something the student does not have the capacity to deliver. Mode-seeking asks for less and gets more.

Some recipes interpolate — a Jensen–Shannon-style mix, or forward KL early and reverse KL late. The interpolation is a knob, not a principle; the principle is that you are choosing between “cover everything the teacher might say” and “say one thing the teacher would approve of”.

What it costs

Per optimiser step, on-policy distillation is meaningfully more expensive than training on a fixed corpus:

The expensive part is student generation, not the teacher. The teacher never generates; it scores a completed sequence in a single forward pass, which is the cheap direction.

Against that, on-policy distillation is far more sample-efficient per token of supervision, because every token it spends is spent on a state the student will actually encounter. The usual result is fewer total tokens for a better student — you are just paying for them in compute rather than in a dataset you built once.

A recipe that holds up

Nobody starts with on-policy distillation from a base model, because early on the student’s samples are too poor to be worth correcting:

1. SFT first. Fine-tune on teacher-generated (or human) data until the student is coherent. This is ordinary off-policy distillation and it is the right tool for getting off the ground.
2. Switch to on-policy. Once the student can produce plausible completions, start sampling from it and scoring with the teacher. This is where the compounding-error fix happens.
3. Keep some off-policy data mixed in. A small fraction of teacher text stops the student drifting into a narrow, self-reinforcing style — the same entropy-collapse risk that shows up in RL fine-tuning.
4. Sample at the temperature you will deploy at. Training on greedy rollouts and serving at temperature=0.8 reintroduces the distribution gap you just paid to close.

Where it fits against RL

On-policy distillation sits between supervised fine-tuning and reinforcement learning, and it is worth being precise about the difference.

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

That is a far denser signal, and it is why distillation converges in a fraction of the steps. The catch is the ceiling: the student is chasing the teacher, so it cannot exceed it. RL has no teacher and no such ceiling — which is exactly why the two are usually used in sequence rather than as alternatives.

Check yourself

In one breath

Train the student on its own outputs, graded by the teacher, and you supervise the states it will actually be in — which is the whole problem with distilling from transcripts. Use a mode-seeking objective so a small model commits instead of hedging, warm up with ordinary SFT first, and expect to trade dataset cost for compute cost.