Reward-Gated On-Policy Distillation

2026-07-04

AI Summary

This work addresses the vulnerability of conventional policy distillation to unreliable supervision signals from teacher models—such as high-confidence rewards for incorrect solutions or penalties for student behaviors that are reasonable but follow different reasoning paths. To mitigate this, the authors propose a verification-aware policy distillation method that dynamically integrates sparse verifier rewards with dense teacher logits through a reward-gating mechanism, activating token-level supervision only when the verifier deems the signal reliable. The approach employs on-policy trajectory sampling, token-level knowledge distillation, and reverse KL divergence optimization. Evaluated on reasoning and code generation benchmarks, it significantly outperforms vanilla reverse KL distillation and TSD-KD baselines, achieving gains of 2.9 and 4.9 points respectively at 1K generation length, and improving by 8.2 points over the unaligned student model in long-generation scenarios.

Abstract

On-policy distillation is a powerful way to transfer reasoning ability from a strong teacher to a smaller student: the student samples trajectories from its own policy, and the teacher provides dense token-level supervision on the states the student actually visits. However, this supervision is not always reliable: a teacher can assign high likelihood to plausible but incorrect solutions, or low likelihood to correct student solutions that follow different reasoning paths. Unconditionally distilling the teacher can therefore reinforce bad modes or erase useful student behavior. To address these limitations, we introduce RG-OPD: Reward-Gated On-Policy Distillation that uses verifier feedback to decide when teacher logits should be trusted. RG-OPD bridges sparse verifier rewards and dense teacher logits, preserving token-level supervision while filtering misleading teacher signals. Across reasoning and coding benchmarks, RG-OPD produces stronger distilled students, outperforming both vanilla reverse-KL distillation and the recent TSD-KD baseline. At 1K generation length, RG-OPD improves over reverse-KL by 2.9 points and over TSD-KD by 4.9 points; in the long-generation setting, it improves over the untuned student by 8.2 points. Our code is available at https://github.com/UoC-tail/RG-OPD.

Problem

• on-policy distillation
• teacher supervision
• reward gating
• misleading signals
• dense token-level supervision

Innovation

• Reward-Gated Distillation
• On-Policy Distillation
• Verifier Feedback
• Token-Level Supervision
• Knowledge Distillation