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

📄 arXiv: 2605.10194v1 📥 PDF

作者: Jiaxuan Wang, Xuan Ouyang, Zhiyu Chen, Yulan Hu, Zheng Pan, Xin Li, Lan-Zhe Guo

分类: cs.AI, cs.LG

发布日期: 2026-05-11

备注: work in progress

💡 一句话要点

提出TRACE框架：通过基于Token路由的自策略对齐，实现高效的强化学习蒸馏

🎯 匹配领域: 支柱二：RL算法与架构 (RL & Architecture)

关键词: 强化学习 自蒸馏 推理优化 Token路由 语言模型对齐 长程推理

📋 核心要点

  1. 现有全Token自蒸馏方法在长程推理中存在梯度冗余，导致特权信息泄露，引发模型推理能力退化与分布外性能下降。
  2. TRACE提出基于Token路由的对齐策略，仅针对关键推理片段进行差异化蒸馏，并结合KL通道衰减机制控制梯度暴露。
  3. 实验表明TRACE在数学推理与复杂问答任务中显著优于GRPO，且在不同规模模型上展现出灵活的路由策略适应性。

📝 摘要（中文）

在线策略自蒸馏（self-OPD）通过特权上下文引导模型自我学习，从而增强强化学习（RLVR）的效果。研究发现，当引导覆盖整个响应时，全Token的KL散度会将梯度浪费在冗余位置，导致特权信息泄露，引发熵增、推理缩短及长程数学任务中的分布外（OOD）退化。为此，本文提出TRACE（Token-Routed Alignment for Critical rEasoning），仅在标注者标记的关键片段上进行蒸馏：对正确路径的关键片段施加前向KL，对局部错误片段施加反向KL，其余Token则使用GRPO优化，并随训练进程衰减KL通道。分析表明，TRACE通过限制累积特权梯度暴露，有效提升了模型对关键Token的分配能力。在四项数学基准及GPQA-Diamond上，TRACE平均超越GRPO 2.76个百分点，且在保持OOD性能的同时，通过在线自标注验证了其泛化潜力。

🔬 方法详解

问题定义：在基于特权上下文的强化学习（RLVR）中，全Token的自蒸馏会导致梯度在无关位置浪费，造成特权信息泄露，进而引发推理链缩短、熵增以及模型在长程任务中的分布外性能退化。

核心思路：引入“关键片段”概念，将蒸馏过程从全序列精简为针对性路由。通过区分关键推理步骤与错误片段，分别施加前向与反向KL约束，从而实现对模型推理逻辑的精准引导，同时通过KL衰减机制避免过度拟合特权信息。

技术框架：TRACE框架包含三个核心组件：对正确路径的关键片段执行前向KL蒸馏；对局部错误片段执行反向KL以纠正偏差；对剩余非关键Token采用GRPO进行常规强化学习优化。训练过程中，KL约束项随训练步数进行退火处理。

关键创新：提出Token路由机制，将蒸馏目标从“全序列对齐”转变为“关键逻辑对齐”。这种方法有效解决了特权信息泄露问题，并证明了通过限制累积梯度暴露，可以显著提升模型在长程推理任务中的鲁棒性。

关键设计：采用动态路由策略，根据模型规模调整优化重点：在大模型（如Qwen3-8B）中侧重于关键片段的前向KL，而在小模型（如Qwen3-1.7B）中则更倾向于通过反向KL纠正错误片段。此外，KL通道的温和衰减确保了模型在后期能从受限引导转向自主推理。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

TRACE在四项数学基准测试及GPQA-Diamond上平均提升2.76个百分点，显著优于GRPO及全Token自蒸馏基线。在GPQA-Diamond任务中，TRACE成功保持了Qwen3-8B的基准OOD分数，而其他方法均出现明显退化。在线自标注实验中，TRACE保留了约69%的强API引导增益，证明了其在无需外部标注下的强大实用价值。

🎯 应用场景

该方法适用于大规模语言模型的推理能力增强，特别是在数学竞赛、复杂逻辑推理及科学问答等长程任务中。其通过自标注机制降低了对外部高质量标注数据的依赖，为资源受限环境下的模型对齐与强化学习训练提供了高效的优化范式。

📄 摘要（原文）

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.