arXiv:2605.10194cs.AIcs.LG2026-05被引 10

只在关键片段教学,让模型学得更准更稳。

TRACE: Distilling Where It Matters via Token-Routed Self On-Policy Alignment

论文配图:TRACE: Distilling Where It Matters via Token-Routed Self On-Policy Alignment
图 1 · 摘自论文原文
  • 只对标注的关键段落进行强化学习自我教学。
  • 数学任务平均提升2.76个百分点,长序列泛化能力更强。
  • 适合追求高可靠性的推理模型训练,尤其关注鲁棒性。

基于策略的自蒸馏(self-OPD)通过提供特权上下文,使策略在可验证奖励的强化学习中实现密集训练。我们发现,当指导覆盖完整响应时,全令牌KL散度会将梯度集中在冗余位置,加剧特权信息泄露,导致熵增、推理缩短及长序列数学训练中的分布外性能下降。为此提出TRACE:仅在标注的关键片段上进行蒸馏——对正确轨迹的关键段落施加前向KL,对局部错误段落可选反向KL,其余部分采用GRPO,并在短暂预热后逐渐衰减KL通道。分析表明,前向KL为学生低估的教师支持项提供持续提升,而段落掩码与衰减机制控制了累积特权梯度暴露。在四个独立数学基准及GPQA-Diamond上,TRACE平均比GRPO提升2.76个百分点,且在GPQA-Diamond上保持了Qwen3-8B基线的分布外得分,而GRPO与全令牌自蒸馏基线均出现退化。在线自标注场景下仍取得1.90个百分点提升(约强API收益的69%),说明其不依赖外部标注。跨规模实验显示最优路由策略因模型而异:在Qwen3-8B上为关键段落前向KL,而在Qwen3-1.7B上则转为错误段落反向KL。

原文摘要 · Abstract (English)

On-policy self-distillation (self-OPD) densifies reinforcement learning with verifiable rewards (RLVR) by letting a policy teach itself under privileged context. We find that when this guidance spans the full response, all-token KL spends gradients on mostly redundant positions and amplifies privileged-information leakage, causing entropy rise, shortened reasoning, and out-of-distribution degradation in long-horizon math training. We propose Token-Routed Alignment for Critical rEasoning (TRACE), which distills only on annotator-marked critical spans: forward KL on key spans of correct rollouts, optional reverse KL on localized error spans, and GRPO on all remaining tokens, with the KL channel annealed away after a short warm-up. Our analysis explains TRACE through two effects: forward KL provides non-vanishing lift to teacher-supported tokens that the student under-allocates, while span masking and decay keep cumulative privileged-gradient exposure finite. On four held-out math benchmarks plus GPQA-Diamond, TRACE improves over GRPO by 2.76 percentage points on average and preserves the Qwen3-8B base OOD score on GPQA-Diamond, where GRPO and all-token self-OPD baselines degrade. Gains persist under online self-annotation (+1.90 percentage points, about 69% of the strong-API gain), reducing the concern that TRACE merely imports external annotator capability. Across scales, the best routed action is base-dependent: on Qwen3-8B it is forward KL on key spans, while on Qwen3-1.7B it shifts to reverse KL on error spans.

强化学习自蒸馏推理优化模型训练

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。