🤖 AI

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

本文提出 TRACE，一种基于令牌路由的自蒸馏方法，该方法在可验证奖励的强化学习过程中，选择性地对标注者标记的关键片段应用 KL 散度，从而减少梯度浪费并防止特权信息泄露，与标准 GRPO 及全响应自蒸馏基线相比，显著提升了长程数学推理能力和分布外泛化能力。

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

发布于 2026-05-12

📖 1 分钟阅读☕ 轻松阅读

CC BY 4.0

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

原始论文采用 CC BY 4.0 许可（http://creativecommons.org/licenses/by/4.0/）。 ✨ 这是对下方论文的AI生成解释。它不是由作者撰写或认可的。如需技术准确性，请参阅原始论文。阅读完整免责声明

想象一下，你正在教一名学生解决复杂的数学问题。你有一位“教师”（一个智能 AI）和一位“学生”（你试图训练的那个 AI）。

过去，研究人员尝试了一种称为在线策略自蒸馏（On-Policy Self-Distillation）的方法。其想法很简单：让学生生成一个答案，然后让教师审视整个答案，并指出：“你这里做对了，这里也做对了，还有这里……"一路贯穿始终。随后，学生会尝试模仿教师的每一个动作。

问题：“过度教师”效应
该论文认为，这种“全盘照抄”的方法实际上有害，尤其对于漫长且复杂的推理任务而言。

类比：想象一名学生正在参加一场长达 10 页的考试。教师逐页审阅，并用红色标出每一个单词，说道：“必须完全照此执行。”
结果：学生感到不堪重负。他们停止独立思考。他们的答案变得更短（因为害怕写太多），开始随机猜测（其“熵”值上升），最终考试失败。论文将这种现象称为“粒度不匹配”。学生并不需要被纠正每一个单词；他们只需要在少数关键步骤上获得帮助，这些步骤是他们可能卡住或犯错的地方。

解决方案：TRACE（关键推理的令牌路由对齐）
作者提出了一种名为TRACE的新方法。TRACE 不像教师那样纠正整篇文章，而是像一位聪明的导师，只指出那些具体且关键的瞬间。

以下是 TRACE 的工作原理，使用一个简单的类比：

1. “聚光灯”（令牌路由）

TRACE 不使用灯光照亮整页纸，而是使用聚光灯。

标注者：一个助手（甚至可以是学生自己！）审视学生的答案，并仅标记出“关键片段”。这些是逻辑要么极其精彩（关键片段），要么极其危险错误（错误片段）的少数句子。
规则：聚光灯仅覆盖答案中约**25%**的内容。其余部分则保持原样，不受干扰。

2. “双工具”方法（FKL 与 RKL）

一旦聚光灯聚焦于特定部分，TRACE 会根据所见内容使用两种不同的工具：

如果学生陷入困境（分配不足）：教师会说：“嘿，你忽略了这重要的一步！看着我，按这种方式做。”这称为前向 KL（FKL）。它促使学生关注正确的路径。
如果学生自信地犯错（过度自信）：学生说：“我确定这是对的！”但实际上它是错的。教师说：“不，停下！你太自信了。退后一步。”这称为反向 KL（RKL）。它抑制学生在错误方向上的过度自信。

3. “渐隐提示”（衰减）

TRACE 最关键的部分在于教师的帮助是暂时的。

类比：想象教师在练习考试中给学生提示。起初，提示清晰响亮。但随着学生进步，提示变得越来越微弱，直到完全消失。
结果：经过短暂的“热身”期（训练中的约 40 步）后，教师完全停止发声。学生被留下，依靠自己的逻辑（强化学习）来解决问题，但他们已经从短暂的指导期间习得了关键习惯。这防止了学生对教师产生依赖。

为何有效（结果）
该论文在数学问题（如解方程）和通用知识问题（GPQA）上测试了这种方法。

“全令牌”方法失败了：当其他方法试图纠正整个答案时，AI 的性能崩溃，并开始给出更短、缺乏创意的答案。
TRACE 成功了：通过仅纠正关键部分随后逐渐退场，TRACE 使 AI 的数学成绩相比标准方法有了显著提升（平均约 2.76%）。
“自学”的额外收益：即使 AI 充当自己的导师（没有人类或更强大的 AI 协助），TRACE 依然运作良好。这证明该方法并非通过引入超级智能教师的知识来“作弊”，而是真正教会了 AI 如何更好地学习。

总结：
TRACE 就像一位教练，不会因球员在比赛中的每一次失误而大喊大叫。相反，教练等待关键时刻，给出尖锐、具体的纠正，然后退后一步，让球员独自比赛。这使球员保持专注、自信，并具备应对漫长而艰巨挑战的能力，而不会导致倦怠。

技术摘要：TRACE——通过 Token 路由实现关键自我在线策略对齐的蒸馏

1. 问题陈述

本文探讨了**在线策略自我蒸馏（Self-OPD）**在应用于具有可验证奖励的长视野强化学习（RLVR），特别是数学推理任务时的局限性。尽管 Self-OPD 方法（如 SDPO、SRPO）通过使用策略快照作为教师来指导模型，从而稠化稀疏的奖励信号，但作者识别出长视野训练中存在一种关键的失效模式：

粒度不匹配： 现有方法在响应的所有 Token上应用知识蒸馏（KL）损失。然而，分析表明，在正确的轨迹中，约 97% 的概率质量在学生和教师之间已经共享，这意味着大多数 KL 梯度是冗余的。
“均匀蒸馏税”： 对所有 Token 应用 KL 迫使模型将梯度花费在不需要修正的位置上。这导致：
- 特权信息泄露： 学生模型吸收了特权上下文中存在的关联（例如已验证的轨迹），而这些关联在推理时并不存在，从而导致分布外（OOD）性能下降。
- 熵增与推理崩溃： 作者在长视野训练中观察到“三症状崩溃”：响应长度急剧缩短，每 Token 的 Actor 熵上升（通常是基线的 4 倍以上），且验证准确率显著下降。
核心问题： 正确的蒸馏单元并非整个响应，而是学生分布实际上需要修正的那一小部分决定性推理 Token。

2. 方法论：TRACE

作者提出了TRACE（关键推理的 Token 路由对齐），该方法仅将蒸馏信号路由到特定的 Token 跨度，而非均匀应用。

核心组件

Token 路由自我 OPD： TRACE 不使用全局 KL 损失，而是利用特权标注器（ $\pi_A$ ）来识别轨迹中的特定跨度：
- 关键跨度（ $K_y$ ）： 正确轨迹（奖励=1）中的决定性推理步骤。
- 错误跨度（ $E_y$ ）： 错误轨迹（奖励=0）中的根本原因推理错误。
- 非跨度（ $N_y$ ）： 所有剩余的 Token。
- 约束： 总跨度覆盖率被限制在 $\alpha = 0.25$ （即 25% 的 Token），以限制特权信息的暴露。
粗略诊断路由：
- 标注器输出一个粗略类型标签（例如 case_split、missed_case），但不输出跨度文本或完整轨迹内容。
- 教师（ $\pi_T$ ）接收问题及该粗略标签作为私有诊断前缀。它在学生生成的前缀上因果地计算 logits，但无法看到具体的跨度位置或内容。
- 学生（ $\pi_S$ ）使用路由损失进行训练，其中 KL 项仅应用于被掩码的跨度。
非对称散度路由：
- 关键跨度上的前向 KL（FKL）： 应用于正确轨迹。FKL（ $KL(\pi_T \parallel \pi_S)$ ）为学生分配不足的 Token 提供“提升”，无论学生当前的概率质量如何。
- 错误跨度上的反向 KL（RKL）（可选）： 应用于错误轨迹。RKL（ $KL(\pi_S \parallel \pi_T)$ ）抑制学生过度自信但教师不偏好的 Token。
- 非跨度上无 KL： 这些 Token 仅依赖标准的 GRPO（组相对策略优化）损失。
有限暴露调度：
- KL 权重（ $\lambda_k$ ）在短暂的预热（例如 40 步）后衰减至零。
- 这确保了在整个训练视野中，累积的特权梯度暴露保持有限，防止 OOD 降解关联的长期积累。

3. 理论分析

本文为路由策略提供了理论依据：

非对称逐点压力（推论 2）：
- FKL 对分配不足的 Token（学生概率 $p \ll$ 教师概率 $q$ ）有效，因为梯度压力与差距 $(q-p)$ 成比例，且当 $p \to 0$ 时不会消失。
- RKL 对自信但错误的 Token（ $p \gg q$ ）有效，因为梯度与学生的过度自信（ $p$ ）成比例。
累积暴露界（命题 3）： 通过结合跨度掩码（将带宽限制为 $\alpha$ ）和 KL 衰减（限制持续时间），TRACE 将累积的特权梯度暴露限制在 $O(\alpha \Lambda^2)$ ，避免了持久性全 Token KL 基线中观察到的无界增长。
信号与泄露： 该方法在关键跨度上最大化对齐信号，同时最小化泄露导致虚假关联的特权信息的风险。

4. 主要贡献

诊断： 识别了全 Token Self-OPD 中的“均匀蒸馏税”，将熵增、认知抑制和特权梯度泄露与粒度不匹配联系起来。
方法与理论： 提出了 TRACE，它根据 Token 类别（关键、错误、非跨度）路由 $\{FKL, RKL, \emptyset\}$ 并衰减特权通道。理论解释了为何在强基线 regime 中 FKL 对关键跨度占主导，以及为何在弱基线 regime 中最佳角落会转向错误跨度的 RKL。
实证证据：
- 性能： 在 Qwen3-8B 上，TRACE 在 5 个基准测试上的平均表现比 GRPO 提高了2.76 个百分点。
- 稳定性： 与 SDPO/SRPO 不同，TRACE 保持了基线模型在 GPQA-Diamond 上的 OOD 性能，而基线方法则出现退化。
- 自我标注： 该方法使用在线自我标注器（学生策略本身），恢复了约 69% 来自强大外部 API 标注器的增益，证明增益源于路由机制，而不仅仅是外部能力。
- 跨尺度适应： 在较弱的基线（Qwen3-1.7B）上，主导路由角落自然地从“关键跨度上的 FKL"逆转为“错误跨度上的 RKL”，与理论预测一致。

5. 结果

Qwen3-8B（强基线）： TRACE（关键跨度上的 FKL）在 5 个基准测试（MATH-500、AIME 24/25、AMC 23、GPQA-Diamond）上实现了**81.51%**的平均准确率，优于 GRPO（78.75%）和最强的先前 Self-OPD 基线 RLSD（78.99%）。关键在于，它保持了基线模型的 GPQA-Diamond 分数（58.33% vs 基线 58.27%），而 GRPO 和全 Token 基线则显著下降。
Qwen3-1.7B（弱基线）： TRACE（错误跨度上的 RKL）实现了**60.16%**的平均分，是唯一超过基线模型（58.77%）的训练方法。全 Token 基线（SDPO/SRPO）发生崩溃，比 GRPO 低 11–19 分。
消融研究：
- 掩码定位： 用随机 25% 掩码或全 Token KL 替换 TRACE 的特定跨度掩码会导致性能下降，证实了哪些Token 被选中比数量更重要。
- 标注器质量： 使用正在积极训练的学生作为其自身的标注器带来了 +1.90 分的提升，表明该方法不依赖于引入外部的“超人”推理能力。

6. 意义与主张

本文主张 TRACE 为长视野 Self-OPD 的不稳定性提供了一种损失几何解决方案，而非特定领域的修复。其意义在于：

解耦蒸馏与不稳定性： 通过将“特权”信号限制在关键跨度并使其衰减，TRACE 避免了在 SDPO 和 SRPO 中观察到的熵增和推理崩溃。
自适应路由： 它证明了最佳蒸馏方向（FKL 与 RKL）并非固定的超参数，而是取决于基线模型的错误特征（分配不足与过度自信）。
实际可行性： 它证明了无需外部教师或昂贵的每查询评分标准即可实现高质量蒸馏，而是依赖于可自我生成的轻量级粗略诊断标注。

作者强调，TRACE 是一个“角落路由框架”而非固定配方，其中具体的路由动作（FKL 与 RKL）应根据基线模型的残差错误模式进行选择。