Policy distillation improves by weighting teacher actions by student benefit
When to Follow the Teacher: Advantage-Based On-Policy Distillation
Machine LearningArtificial Intelligence
Summary
Teaching an AI agent usually involves it copying a teacher’s behavior exactly, but that doesn’t always help the student AI learn better. The authors created a method that lets the student AI pay more attention to teacher actions that actually improve its chances of success and less to actions that don’t help. Their approach estimates how useful each teacher action is without needing extra trial runs, making learning more efficient. Tests show this method works better than existing ones on tasks like math, coding, and medical question answering.
What this means in practice
- •For machine learning engineers: Train AI agents more effectively by focusing teacher guidance on actions that improve student performance without extra data collection.
- •For ai software developers: Improve automated tutoring systems by selectively reinforcing demonstrations that truly aid AI learners during on-policy training.
Authors
Mahdi Sabbaghi, George Pappas, Hamed Hassani, Adel Javanmard
Abstract
Online Policy Distillation (OPD) often outperforms offline methods such as Supervised Fine-Tuning (SFT) by training on trajectories sampled directly from the student. However, standard OPD trains the student to follow the teacher at every state of the trajectory, assuming that teacher guidance is always beneficial. We instead distinguish between what the teacher prefers and what actually helps the current student succeed. We introduce Advantage-Based On-Policy Distillation (AOPD), which weights teacher actions by their counterfactual advantage under the student. Teacher actions that improve the student's expected outcome are reinforced, while those that do not are downweighted. Directly computing these advantages requires trying alternative teacher actions at each step and running new rollouts from each. Instead, we show that the AOPD objective can be estimated using only standard student rollouts and their terminal rewards. Our estimator uses teacher importance weights on the student-sampled actions to provide unbiased advantage-weighted supervision, without any additional rollouts or reward evaluations. This allows AOPD to focus teacher supervision on actions that help the current student while preserving the rollout efficiency of standard OPD. Experiments on math, coding, and Medqa show gains over distillation and RL baselines in both external-teacher and self-distillation settings.