TL;DR
Offline distillation trains a student on teacher-generated data. On-policy distillation instead samples from the student and asks the teacher to supply token distributions or feedback on those states. The student therefore learns where it actually makes mistakes, echoing the motivation behind DAgger in imitation learning.
The Distribution-Shift Problem
A student deployed autoregressively conditions on its own previous tokens. Small errors move it away from the teacher-data distribution, after which the offline dataset provides little guidance. On-policy collection repeatedly refreshes training examples from the current student policy.
For student prefix (x_{<t}\sim\pi_S), a token-level objective can minimize
[ \operatorname{KL}(\pi_T(\cdot|x_{<t});||;\pi_S(\cdot|x_{<t})). ]
def distillation_loss(student_log_probs, teacher_probs):
return -(teacher_probs * student_log_probs).sum(-1).mean()
Distillation Is Not RL
The teacher provides a dense target rather than a scalar return, so no credit-assignment algorithm is required. That makes training simpler, but the student is bounded by the teacher signal and inherits its biases. Full RL remains useful when an external verifier can reward solutions the teacher would not generate or when exploration itself is the goal.
Practical Recipe
Start from supervised distillation, generate a bounded batch with the student, query the teacher only for retained prefixes, and mix on-policy examples with a stable replay set to reduce forgetting. Track teacher-query cost, student KL, task accuracy, and diversity. Keep the teacher temperature and prompt format fixed when comparing checkpoints.
References
- Ross, S., Gordon, G., and Bagnell, D. (2011). A Reduction of Imitation Learning and Structured Prediction to No-Regret Online Learning.
- Agarwal, R. et al. (2024). On-Policy Distillation of Language Models: Learning from Self-Generated Mistakes.
- Hinton, G., Vinyals, O., and Dean, J. (2015). Distilling the Knowledge in a Neural Network.