On-Policy Distillation and Exposure Bias
Why a student trained only on teacher trajectories cannot recover from its own errors, how generating from the student fixes the state distribution, and what reverse KL buys and costs.
A student trained on teacher outputs is only ever asked what to do in situations the teacher created. At inference it creates its own situations, makes a small error at token 40, and finds itself in a state no training example resembled. The error compounds. This is exposure bias, it is the oldest known problem in sequence generation training, and on-policy distillation is the direct answer to it.
The state distribution mismatch
Frame distillation as imitation learning. Behavioural cloning trains on state-action pairs from an expert, and its classical failure is that errors move the learner off the expert's state distribution, where its training gives no guidance, so errors compound quadratically in the horizon rather than linearly. This is the standard analysis behind DAgger in imitation learning, and it transfers directly to sequence models.
On-policy distillation applies the same fix. Generate sequences from the student, then ask the teacher for its distribution at each of those positions and train the student to match it there. The supervision now lands exactly on the states the student actually reaches, including the ones it reaches by mistake, so it learns what a good model would do after a bad token.
Generalised knowledge distillation (Agarwal et al., 2024, arXiv:2306.13649) formalises this with two knobs: the fraction of on-policy versus fixed data, and the choice of divergence.
Forward versus reverse KL
The divergence choice is not cosmetic.
Forward KL, \(\mathrm{KL}(p_{\text{T}} \| p_{\text{S}})\), is mode-covering. It is infinite where the teacher has mass and the student has none, so the student is forced to cover everything, including modes it lacks capacity to represent, and it ends up spreading probability across regions with nothing in them.
Reverse KL, \(\mathrm{KL}(p_{\text{S}} \| p_{\text{T}})\), is mode-seeking. It is infinite where the student has mass and the teacher has none, so it punishes the student for producing anything the teacher would not, and permits ignoring teacher modes entirely. For a capacity-limited student this is usually the better failure: a model that does a subset of things well beats one that does everything badly.
Reverse KL is also naturally on-policy, since its expectation is taken over the student's own distribution, which is why the two ideas arrive together. Jensen-Shannon and generalised interpolations between the two are available and are mostly a way to tune how aggressively the student is allowed to drop modes.
Cost and the RL connection
The bill is generation. Every training step needs the student to produce sequences and the teacher to score them, so a step costs an inference pass over both models plus the backward pass, rather than a backward pass over a cached dataset. That is several times the cost of offline distillation per step, offset partly by needing fewer steps.
Structurally this is reinforcement learning with the teacher as a dense reward model, and it inherits the machinery: rollout buffers, on-policy batching, and a KL term that keeps the student from drifting. The practical difference from RLHF is that the reward here is a full distribution over the vocabulary rather than a scalar, which is a far richer signal per token and is why on-policy distillation converges much faster than RL from scalar rewards.
When it breaks
The teacher must be able to score arbitrary sequences. On-policy distillation needs teacher logits on student-generated text, which requires white-box access. An API that returns only sampled text supports sequence-level distillation and not this. That access requirement, more than compute, is what decides which method a team can use.
Early training generates garbage. A randomly initialised or weak student produces sequences so far off-distribution that the teacher's opinion about them is close to meaningless, and training on them wastes compute. Standard practice is to start with offline supervised fine-tuning and switch to on-policy once the student is competent, which makes the schedule between the two a real hyperparameter.
Mode dropping is a feature until it is not. Reverse KL lets the student abandon modes, which is desirable when they are beyond its capacity and undesirable when the dropped mode is a language, a domain, or a demographic. The loss will not tell you which happened; only evaluation stratified by the axis you care about will.
It is not a substitute for RL on the tasks RL is for. On-policy distillation makes the student imitate the teacher on the student's own states. It cannot exceed the teacher, and where the objective is verifiable correctness rather than imitation, reinforcement learning against a verifier reaches places the teacher never did.
14 flashcards for this concept
Click a card to reveal the answer.