REOPD: Reliability-Adaptive Reward Extrapolation for On-Policy Distillation

2026-08-12

AI Summary

Traditional policy distillation methods relying on global reward extrapolation are prone to reward hacking and training instability, often necessitating costly hyperparameter tuning. This work proposes a reliability-adaptive reward extrapolation framework that, for the first time, enables token-level adaptive extrapolation without requiring a verifier, reward model, or additional sampling. The approach dynamically adjusts extrapolation coefficients via token-level compatibility weights and batch-level adaptive budgets, formulated as λ_{b,t} = 1 + γ_b q_t, thereby enhancing optimization stability while preserving alignment with the teacher policy. Empirical results demonstrate that the method outperforms G-OPD in both single-teacher mathematical tasks and multi-teacher settings, and achieves comparable performance on single-teacher code generation tasks, confirming its effectiveness across domains and under diverse teacher configurations.

Abstract

On-policy distillation (OPD) trains a student on its own trajectories under dense token-level supervision from a teacher. Reward-extrapolation methods such as ExOPD amplify the teacher-reference log-likelihood ratio to move beyond direct imitation, but apply a single global coefficient $λ$ to every token. This can drive the student to fit extreme peaks in the implicit reward, causing reward hacking and unstable training, and the optimal $λ$ varies across domains, requiring costly sweeps. We propose REOPD, a reliability-adaptive reward extrapolation framework for OPD. REOPD combines a token-level compatibility weight with a batch-level adaptive budget, yielding a token-wise coefficient $λ_{b,t}=1+γ_b q_t$ that preserves teacher alignment while selectively extrapolating along reliable teacher-reference directions. It requires no verifier, reward model, value model, or extra rollout beyond standard OPD. REOPD outperforms G-OPD on single-teacher mathematics and on both domains in the multi-teacher setting, while matching G-OPD on single-teacher code, demonstrating effective fine-grained reliability adaptation across domains and teacher configurations.

Problem

• on-policy distillation
• reward extrapolation
• reward hacking
• adaptive coefficient
• token-level supervision

Innovation

• on-policy distillation
• reward extrapolation
• reliability adaptation
• token-level supervision
• adaptive coefficient