More efficient reinforcement learning by selecting key tokens for distillation
Not Every Token Is Worth Distilling: Selective Supervision for Direct-OPD
Machine LearningArtificial Intelligence
Summary
Teaching small AI models to become better can be done by copying improvements onto bigger models, but not every teaching point is helpful. The authors show that giving feedback on all parts without discrimination can sometimes mislead, so they suggest focusing only on the parts where the teacher AI shows clear improvement. This selective focus helps the bigger model learn better without extra work. Testing on different AI sizes and tasks, this new way mostly improved accuracy or kept it as good as before.
What this means in practice
- •For machine learning engineers: Improve accuracy of large language models by selectively transferring improvements from smaller models without added computational cost.
- •For nlp system developers: Enhance downstream language understanding tasks by focusing model training on states with significant policy changes.
Authors
Yibo Zhao, Zixuan Yang, Yunshi Lan, Xiang Li
Abstract
Direct On-Policy Distillation (Direct-OPD) transfers reinforcement-learning-induced policy improvements from a small model to a larger student by using the token-level log-ratio between post-RL and pre-RL checkpoints as dense supervision on the student's own rollouts. This transfer rewards the policy shift at every state, yet the log-ratio measures only relative change: it can stay fixed even as the probability mass that both checkpoints assign to the student's candidate tokens vanishes. Through an exact construction, we show that the Direct-OPD reward and its update can remain unchanged while the Jensen-Shannon divergence (JSD) and both KL directions between the checkpoints vanish with this mass, and we note that a small JSD bounds how much the teacher's behavior changed. Motivated by this analysis, we propose Selective Supervision for Direct-OPD (S$^2$D-OPD), which ranks student-sampled states by their teacher-reference JSD and masks Direct-OPD supervision at low-divergence states, retaining only the top 10% of states per response. Across two teacher pairs and four student models ranging from 1.7B to 8B parameters, S$^2$D-OPD improves held-out accuracy over dense Direct-OPD on AIME and HMMT benchmarks in seven of eight settings and matches it in the eighth, without extra forward passes. Our code is available at https://anonymous.4open.science/r/S2D-OPD-8868.