On policy distillation focuses teacher tokens on key successful and failed examples
Spend Teacher Tokens Where They Matter: Success-Referenced On-Policy Distillation
Artificial Intelligence
Summary
Teaching large AI models can be very costly because it usually requires checking every example they try and correcting them. The authors propose a smarter method where the AI learns mostly from its mistakes where there is also a successful attempt to compare against. This way, the teacher AI only needs to spend its effort where it really helps, saving a lot of work without hurting the AI’s learning. Their method uses just a small fraction of teaching tokens but still keeps the AI’s problem-solving skills high.
What this means in practice
- •For ai model trainers: Allocate teacher feedback only to student attempts where both success and failure exist, reducing computational overhead while maintaining model quality.
- •For machine learning platform engineers: Implement token-efficient supervision strategies to reduce cloud compute costs by focusing teacher input on key divergent student outputs.
Authors
Xiang Chen, Futao Su, Kong Wang, Jiayi Chen, TanLin Li
Abstract
On-policy distillation (OPD) combines student-generated rollouts with dense token-level supervision from a teacher, but providing such supervision for every rollout requires substantial teacher computation. We introduce Success-Referenced On-Policy Distillation (SR-OPD), which reduces this cost by selecting which prompts and rollouts receive teacher supervision. When the student produces both successful and failed rollouts for the same prompt, a successful rollout can serve as a natural reference for selecting failed rollouts. SR-OPD therefore focuses on such prompts and prioritizes failed rollouts whose hidden-state trajectories show sustained divergence from a successful reference, while accounting for estimated teacher-input cost. Across three teacher-student pairs and six mathematical reasoning benchmarks, SR-OPD uses only 3.46-5.02% of the teacher-input tokens required by Vanilla OPD in the one-pass setting while maintaining comparable reasoning performance. Under a controlled setting matched to 5% of Vanilla OPD's teacher-input budget, further experiments support both key design choices: focusing supervision on prompts with both successful and failed rollouts, and using successful rollouts to guide failure selection. These results indicate that a student's own successful behavior can serve as a useful reference for allocating teacher supervision under a fixed teacher-input budget.