Blockwise Policy-Drift Gating for On-Policy Distillation
作者: Liwen Zheng, Haiyun Jiang
分类: cs.LG, cs.AI, cs.CL
发布日期: 2026-06-23
备注: 8 pages
💡 一句话要点
提出块状策略漂移门控以解决长时间推理任务中的脆弱性问题
🎯 匹配领域: 支柱二：RL算法与架构 (RL & Architecture)
关键词: 在线蒸馏 策略漂移 长时间推理 数学推理 机器学习
📋 核心要点
- 现有的在线蒸馏方法在长时间推理任务中表现脆弱，尤其是在采样的轨迹上进行训练时，容易受到局部匹配的影响。
- 本文提出的块状策略漂移门控方法，通过计算行为学生与当前学生的对数概率变化，利用这些变化来改进OPD的鲁棒性。
- 在六个Qwen3数学推理基准上，固定64-token块门控将平均求解率从0.4978提升至0.5160，显示出显著的性能提升。
📝 摘要（中文）
本文提出了一种块状策略漂移门控（Blockwise Policy-Drift Gating）方法，用于在回滚重用的情况下进行在线蒸馏（OPD）。该方法通过计算行为学生与当前学生在采样路径上的对数概率变化，聚合这些变化并使用均值归一化的门控来重新加权OPD位置损失。实验结果表明，固定64-token块门控显著提高了多个数学推理基准的求解率，验证了局部策略漂移作为OPD回滚的有效控制信号。
🔬 方法详解
问题定义：本文旨在解决在线蒸馏（OPD）在长时间推理任务中脆弱性的问题。现有方法在处理长时间轨迹时，容易受到局部教师支持匹配的影响，导致性能不稳定。
核心思路：提出块状策略漂移门控，通过计算行为学生与当前学生在采样路径上的对数概率变化，聚合这些变化并使用均值归一化的门控来重新加权OPD位置损失，从而提高鲁棒性。
技术框架：整体架构包括三个主要模块：首先，计算行为学生与当前学生的对数概率变化；其次，将这些变化在固定块上进行聚合；最后，利用得到的门控信号重新加权OPD损失。
关键创新：最重要的创新点在于引入了块状门控机制，作为一种轻量级的学生端控制器，能够有效地处理回滚重用中的策略漂移问题，与现有方法相比，提供了更为稳定的性能。
关键设计：在实验中，采用固定64-token块进行门控设计，损失函数中引入了均值归一化的门控信号，确保了在不同任务中的适应性和鲁棒性。
📊 实验亮点
实验结果显示，固定64-token块门控将多个数学推理基准的平均求解率从0.4978提升至0.5160，且在Teacher-TopK/LSM上，Block64在四个基准中表现最佳，验证了局部策略漂移作为控制信号的有效性。
🎯 应用场景
该研究的潜在应用领域包括智能体训练、自动化决策系统和复杂任务的求解等。通过提高在线蒸馏的鲁棒性，该方法能够在多种长时间推理任务中实现更高的求解率，具有重要的实际价值和未来影响。
📄 摘要（原文）
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.