Blockwise Policy-Drift Gating for On-Policy Distillation

2026-06-22

AI Summary

This work addresses the instability of policy distillation in long-horizon reasoning tasks when relying on student-sampled trajectories. To mitigate this issue, the authors propose a lightweight, block-level policy drift gating mechanism that dynamically reweights position-wise losses without altering the teacher targets or rollout policies. The method computes the log-probability divergence between the behavioral and current policies over fixed-length token blocks, then applies mean normalization to generate a gating signal. Notably, it leverages local policy drift—measured as the discrepancy between old and new policies—as a practical control signal, substantially enhancing distillation robustness. Evaluated on the Qwen3 mathematical reasoning benchmark, the approach with 64-token blocks improves average pass@8 from 0.4978 to 0.5160 and achieves state-of-the-art average performance across four benchmarks under the Teacher-TopK/LSM setting.

Abstract

On-policy distillation (OPD) trains a student policy using teacher signals computed on trajectories sampled by the student itself. Recent work shows that sampled-token OPD can be fragile on long-horizon reasoning tasks and that local teacher-support matching is a simple and effective repair. This paper introduces blockwise policy-drift gating, a lightweight student-only old-current drift controller for OPD under rollout reuse. The method computes log-probability shifts between the behavior student and the current student on the sampled token path, aggregates these shifts over fixed blocks or spans, and uses the resulting detached, mean-normalized gates to reweight OPD position losses. It does not change teacher targets, teacher top-K supports, or the rollout policy. In a six-variant Qwen3 math reasoning benchmark with a uniform 200-step training budget for all trained variants, we use pass@8 as the primary problem-level solve-rate metric. Fixed 64-token block gating improves sampled-token OPD mean pass@8 from 0.4978 to 0.5160 across AIME24, AIME25, MATH500, and AMC23. On Teacher-TopK/LSM, Block64 gives the best four-benchmark mean pass@8 among trained students. The results identify local old-current policy drift as a practical control signal for reused OPD rollouts and motivate block-level gating as a simple default for improving solve-rate robustness.

Problem

• on-policy distillation
• policy drift
• long-horizon reasoning
• rollout reuse
• student policy

Innovation

• blockwise gating
• policy drift
• on-policy distillation
• rollout reuse
• log-probability shift