Title: Prune-OPD: Efficient and Reliable On-Policy Distillation for Long-Horizon Reasoning

Abstract:

On-policy distillation (OPD) utilizes dense teacher rewards to bolster the reasoning capabilities of language models. However, applying OPD to long-horizon tasks reveals a significant vulnerability: as a student’s generated prefix inevitably deviates from the teacher’s reasoning path, the teacher’s dense reward signal becomes locally unexploitable. Persisting in token generation and evaluation along these "drifted" trajectories not only compromises reward quality but also results in substantial computational inefficiency.

To mitigate these issues, we present Prune-OPD, a framework designed to dynamically synchronize training resources with the quality of supervision. By continuously assessing the local alignment between student and teacher predictions—such as through top-$k$ overlap metrics—Prune-OPD identifies prefix-drift events as they occur. When significant drift is detected, the system monotonically reduces the weight of subsequent unreliable rewards and initiates dynamic rollout truncation. This mechanism stops futile generation processes, allowing for the reallocation of compute power exclusively to reliable teacher supervision.

Evaluated across various teacher-student pairings, Prune-OPD consistently aligns computational effort with supervision reliability. In scenarios where prefix drift renders dense teacher rewards ineffective, the framework cuts training time by 37.6% to 68.0% while maintaining, and frequently enhancing, performance on rigorous benchmarks such as AMC, AIME, and HMMT. Conversely, when student-teacher compatibility is strong, Prune-OPD automatically extends the training window to retain long-context supervision. These findings indicate that Prune-OPD enhances OPD efficiency not through indiscriminate rollout shortening, but by strategically directing computation toward locally exploitable teacher rewards.

Source: arXiv Generated at: 2026-06-02 00:00:00 UTC