Prune-OPD: 장기 추론을 위한 효율적이고 신뢰성 있는 온-폴리시 증류

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

온-폴리시 증류(OPD)는 풍부한 교사 모델의 보상을 활용하여 추론 모델을 향상시킵니다. 그러나 OPD를 장기적인 작업에 적용할 때 중요한 문제가 발생합니다. 학습 모델이 생성하는 초기 부분이 필연적으로 교사 모델의 사고 과정과 달라짐에 따라, 교사 모델의 풍부한 보상은 지역적인 활용 가치를 잃게 됩니다. 이러한 '편차'가 발생한 경로에서 계속해서 토큰을 생성하고 평가하는 것은 보상의 품질을 저하시킬 뿐만 아니라 막대한 계산 자원을 낭비합니다. 이러한 문제를 해결하기 위해, 본 논문에서는 훈련 예산을 지도 품질에 따라 동적으로 조정하는 프레임워크인 **Prune-OPD**를 제안합니다. Prune-OPD는 학습 모델과 교사 모델의 예측 간의 지역적인 일관성을 지속적으로 모니터링합니다 (예: 상위 $k$개 항목의 중첩을 통해). Prune-OPD는 이러한 편차를 실시간으로 감지하고, 심각한 편차가 감지되면 이후의 신뢰할 수 없는 보상을 점진적으로 감소시키고 동적 롤아웃 중단을 트리거합니다. 이를 통해 학습 과정은 더 이상 의미 없는 생성을 중단하고 계산 자원을 신뢰할 수 있는 교사 모델의 지도를 위해 엄격하게 재할당합니다. 다양한 교사-학습 모델 조합에서 Prune-OPD는 일관되게 계산 자원을 지도 신뢰성에 맞게 조정합니다. 초기 부분의 편차로 인해 교사 모델의 풍부한 보상이 신뢰할 수 없게 될 경우, Prune-OPD는 훈련 시간을 37.6%에서 68.0%까지 줄이면서 동시에 AMC, AIME, HMMT와 같은 어려운 벤치마크에서 성능을 유지하거나 향상시킵니다. 학습 모델과 교사 모델의 일관성이 높을 경우, Prune-OPD는 자동으로 훈련 범위를 확장하여 장기적인 컨텍스트 지도를 유지합니다. 이러한 결과는 Prune-OPD가 롤아웃 길이를 무작정 줄이는 것이 아니라, 지역적으로 활용 가능한 교사 모델의 보상에 계산 자원을 재할당함으로써 OPD를 개선한다는 것을 시사합니다.

On-policy distillation (OPD) leverages dense teacher rewards to enhance reasoning models. However, scaling OPD to long-horizon tasks exposes a critical flaw: as the student's generated prefix inevitably diverges from the teacher's thought process, the teacher's dense reward loses local exploitability. Continuing to generate and evaluate tokens on these ``drifted'' trajectories not only degrades reward quality but also incurs massive computational waste. To address this, we introduce \textbf{Prune-OPD}, a framework that dynamically aligns training budgets with supervision quality. By continuously monitoring the local compatibility between student and teacher predictions (e.g., via top-$k$ overlap), Prune-OPD detects prefix-drift events in real time. Upon detecting severe drift, it monotonically down-weights subsequent unreliable rewards and triggers dynamic rollout truncation. This allows the training process to halt futile generation and reallocate compute strictly to reliable teacher supervision. Across diverse teacher-student combinations, Prune-OPD consistently aligns computation with supervision reliability. When prefix drift makes dense teacher rewards unreliable, it reduces training time by 37.6\%--68.0\% while preserving, and often improving, performance on challenging benchmarks (AMC, AIME, HMMT). When student-teacher compatibility remains high, it automatically preserves long-context supervision by expanding the training window. These results suggest that Prune-OPD improves OPD not by blindly shortening rollouts, but by reallocating computation toward locally exploitable teacher rewards.