Teacher adaptation improves student learning from on-policy data
Meet Students Where They Are: Student-Prefix Teacher Adaptation for On-Policy Distillation
Computation and Language
Summary
It can be hard for a teaching computer program to help another program learn when the learner makes mistakes or tries unexpected things. The paper shows a way to make the teacher better at understanding these learner mistakes by adapting it before it starts teaching. This adaptation helps the teacher give useful guidance that matches what the student actually does, leading to better student performance on math reasoning tasks.
What this means in practice
- •For machine learning engineers: Improve training of language models by adapting the teacher model to better support student model learning from its own generated data.
- •For ai system trainers: Enhance reasoning model deployment by refining teacher feedback to handle learner mistakes in on-policy training scenarios.
Authors
Xue Zhang, Songming Zhang, Yunlong Liang, Fandong Meng, Yufeng Chen, Jinan Xu
Abstract
On-policy distillation (OPD) mitigates the student's exposure bias by optimizing on student-generated trajectories. However, student-generated prefixes may contain reasoning errors or follow patterns unfamiliar to the teacher, undermining the reliability of its token-level supervision. This degradation becomes more pronounced as prefixes grow longer. To address this teacher-side distribution mismatch, we propose Student-Prefix Teacher Adaptation (SPTA), which adapts the teacher to the states on which it is queried during distillation. Specifically, before OPD starts, we collect student rollout prefixes with random lengths and continually train the teacher to continue from the prefixes using reinforcement learning with verifiable rewards (RLVR). By treating each prefix as context and rewarding correct final answers, the teacher learns to recover from student-induced states rather than imitate the student's reasoning. Then, the adapted teacher is used for standard OPD, with no changes to the distillation objective. Experiments on mathematical reasoning benchmarks show that our SPTA improves average student performance over standard OPD and OPD with task-level RLVR-adapted teachers. These results highlight the importance of adapting teachers not only to solve problems independently, but also to provide reliable guidance on the states their students actually visit.