TRACE：通过词元路由自策略对齐在关键处蒸馏

摘要

策略内自蒸馏（self-OPD）通过让策略在特权上下文下自我教学，使可验证奖励强化学习（RLVR）的奖励信号更密集。我们发现，当这种指导覆盖完整响应时，全词元KL将梯度花费在大多冗余的位置上并放大特权信息泄露，在长程数学训练中导致熵上升、推理变短以及分布外性能退化。我们提出TRACE（Token-Routed Alignment for Critical rEasoning），只在标注者标记的关键片段上进行蒸馏：对正确rollout的关键片段施加前向KL，对局部错误片段可选地施加反向KL，对其余全部词元使用GRPO，并在短暂预热后将KL通道退火移除。我们的分析通过两种效应解释TRACE：前向KL为学生分配不足的教师支持词元提供非消失的提升，而片段掩码与衰减使累积的特权梯度暴露保持有界。在四个留出数学基准和GPQA-Diamond上，TRACE平均比GRPO提高2.76个百分点，并保住了Qwen3-8B基础模型在GPQA-Diamond上的OOD分数，而GRPO和全词元self-OPD基线均出现退化。在在线自标注下增益依然存在（+1.90个百分点，约为强API增益的69%），缓解了TRACE仅引入外部标注者能力的担忧。跨规模来看，最佳路由动作取决于基础模型：在Qwen3-8B上是对关键片段的前向KL，而在Qwen3-1.7B上则转向对错误片段的反向KL。

TRACE：通过词元路由自策略对齐在关键处蒸馏：论文配图 — 图 2：TRACE 管道：(1) 学生样本推出 y^\hat{y}，验证者返回 RR； (2) 注释器 πA\pi_{A} 生成跨度掩码和粗略类型标签； (3) 教师接收类型标签作为私有诊断前缀，并在 y^<t\hat{y}_{<t} 上因果计算 logits； (4) 每个跨度类上的路由 KL 操作 — 默认为 𝒦y\mathcal{K}_{y} 上的 FKL，ℰy\mathcal{E}_{y} 和 𝒩y\mathcal{N}_{y} 上为 no-KL — 与 𝒩y\mathcal{N}_{y} 上的 GRPO 结合，并且 KL 权重在预热后衰减为零。