On-Policy Distillation

On-policy distillation is a language model post-training technique in which a student model generates its own sequences of tokens, and a stronger teacher model grades each token of that output rather than supplying the sequences itself, as standard imitation-based distillation would. Kevin Lu and colleagues at Thinking Machines Lab described the method in the October 2025 post "On-Policy Distillation." At each step, the student's output distribution is compared against the teacher's distribution for that same context using the per-token reverse KL divergence. The student is then updated to move toward the teacher's judgment specifically at the states the student itself tends to visit, rather than at states drawn from a fixed, teacher-authored dataset.

The technique sits between two more established approaches. Reinforcement learning trains on the student's own trajectories but supplies only a single sparse reward at the end of an episode, no matter how many tokens the trajectory contains. Off-policy distillation, or straightforward imitation of teacher-generated examples, gives a dense per-token training signal, but it trains the student on states drawn from the teacher's behavior rather than its own. That can leave the student fluent on the teacher's distribution while brittle on the states it actually encounters once deployed. On-policy distillation combines the dense, per-token signal of distillation with on-policy sampling, giving the student roughly as many bits of training signal as it has generated tokens per episode, rather than one bit per episode.

Thinking Machines reported using the method to train small models toward expert-level performance on math reasoning, and to keep a continuously updated assistant current, at substantially lower compute cost than reinforcement learning. Follow-up research, including "Rethinking On-Policy Distillation of Large Language Models" (arXiv:2604.13016) and its sequel, found that the method's success depends on the student and teacher sharing compatible reasoning patterns, and on the teacher offering genuinely new capability beyond what the student has already seen. The same research reported that training on as few as one or a small handful of prompts can recover most of the gains of training on a full dataset.