技术摘要：OISD——语言模型的在线策略内部自蒸馏

1. 问题陈述

近期针对大语言模型（LLM）的强化学习（RL）后训练方法，特别是那些利用可验证奖励的组相对策略优化（GRPO）方法，显著推动了数学推理能力的进步。然而，这些方法主要以最终输出为中心。它们利用稀疏的、结果层面的奖励（例如最终答案的正确性）来优化最终执行策略，而很大程度上忽视了编码在中间层中的预测信号、推理过程和内部表示。

虽然知识蒸馏（KD）提供了一种提供密集 token 级监督的机制，但现有范式存在特定的局限性：

• 离线策略蒸馏（Off-Policy Distillation）： 由于学生模型是在固定教师生成的轨迹上进行训练，而非其自身生成的轨迹，因此会产生分布不匹配。
• 在线策略蒸馏（On-Policy Distillation）： 虽然减少了不匹配，但通常依赖外部教师模型。
• 在线策略自蒸馏（On-Policy Self-Distillation）： 近期方法移除了外部教师，但依赖外部注入的特权信息（例如特权提示、真实解或辅助反馈）来监督模型。

目前存在一个空白，即在不依赖外部特权数据或多重执行策略的情况下，利用模型自身的内部计算作为监督源。此外，先前的研究表明，推理表示在各层之间系统地演变（从早期层的广泛候选预测到深层层的收敛决策），这表明中间层可以由更深层引导，以提高推理的一致性。

2. 方法论：OISD 框架

作者提出了OISD（在线策略内部自蒸馏），这是一个将模型的最终层视为“分离的内部教师”，并将选定的中间层视为“辅助内部学生”的框架。该过程发生在 rollout 和 GRPO 优化过程中的单次前向传播内。

核心架构

• 单一执行策略： 最终层（$L$）保持为生成 rollout 和优化 GRPO 目标的唯一策略。
• 内部教师 - 学生： 在同一前向传播过程中，选定中间层（$l < L$）的隐藏状态使用解嵌入矩阵投影到 token 空间（遵循logit lens框架）。这创建了一个临时的“学生”分布，该分布源自与“教师”（最终层）相同的上下文。
• 无外部特权： 监督信号完全源自模型在同一在线策略 rollout 上的自身内部状态，不需要外部真实解或特权提示。

对齐目标

OISD 采用两种互补的内部对齐机制，将预测信号从最终层转移到中间层：

1. Logit 对齐（“如何思考”）：

   • 目标： 转移高层推理行为和预测信念。
   • 机制： 计算中间层（$p^l_t$）和最终层（$q^L_t$）的温度缩放 token 分布之间的 Jensen–Shannon（JS）散度。
   • 加权： 损失由带符号的 GRPO 优势（$A_t$）加权。高优势 token 将中间层拉向最终层的分布，而低优势 token 将其推开，以避免强化不良的延续。
   • 公式： $\mathcal{L}^{(l)}_{\text{think}} = \mathbb{E}_t [ A_t \cdot \text{JS}(p^{l,\tau}_t, \text{sg}(q^{L,\tau}_t)) ]$

2. 注意力对齐（“看向哪里”）：

   • 目标： 强制一致的注意力模式和证据路由行为。
   • 机制： 计算中间层和最终层在采样的一组因果键（结合局部窗口和全局步长位置以管理内存）上的注意力权重之间的 JS 散度。
   • 加权： 与 logit 对齐类似，损失由带符号的优势加权。正优势鼓励中间层关注与最终层相同的证据；负优势则减少对低奖励模式的一致性。
   • 公式： $\mathcal{L}^{(l)}_{\text{attn}} = \mathbb{E}_t [ A_t \cdot \frac{1}{H} \sum_{h=1}^H \text{JS}(\tilde{a}^l_{t,h}, \text{sg}(\tilde{a}^L_{t,h})) ]$

训练目标

总损失结合了标准 GRPO 损失与两个内部对齐损失：
$$\mathcal{L}_{\text{OISD}} = \mathcal{L}_{\text{GRPO}} + \lambda_{\text{think}} \mathcal{L}_{\text{think}} + \lambda_{\text{attn}} \mathcal{L}_{\text{attn}}$$
其中 $\lambda_{\text{think}}$ 和 $\lambda_{\text{attn}}$ 控制蒸馏的强度。

3. 主要贡献

1. 新范式： 引入在线策略内部自蒸馏，这是一种推理 RL 范式，利用模型自身的内部计算进行监督，同时保持单一执行策略。
2. OISD 框架： 一种具体实现，使用最终层作为分离教师，中间层作为学生，在不依赖外部特权信息的情况下转移在线策略预测信号。
3. 双重对齐机制： 开发了两个互补目标：
   • Logit 对齐： 监督推理过程（“如何思考”）。
   • 注意力对齐： 监督证据路由（“看向哪里”）。
4. 实证验证： 证明了在四个数学推理基准测试中，相较于强大的推理 RL 基线（包括 GRPO 和 BuPO），OISD 取得了显著且一致的提升。

4. 实验结果

作者在 Qwen3-4B、Qwen3-8B 和 OctoThinker 模型上评估了 OISD，测试涵盖四个基准：AMC23、MATH500、AIME24 和 AIME25。

• 性能提升：
  • 在 Qwen3-4B 上，OISD 取得了 64.75 的平均分，优于 GRPO（55.08）9.67 分，优于 BuPO（58.51）6.24 分。
  • 在 AIME 基准测试上的提升尤为显著（例如，在 AIME24 上比 GRPO 高出 14.31 分）。
  • 在 Qwen3-8B 上，OISD 取得了 66.72 的最佳总体平均分，优于 GRPO（64.23）和 BuPO（66.36）。
• Pass@K 分析： 与 GRPO 相比，OISD 在所有 $K$ 值下均实现了更高的 Pass@K 率，表明采样解集的整体质量得到了提升，而不仅仅是平均准确率的提高。
• 消融研究：
  • Logit 和注意力对齐均产生了积极贡献。Logit 对齐在 AIME24 上显示出最大增益，而注意力对齐在平均表现上略好。
  • 结合两者取得了最佳结果，证实了它们的互补性。
• 与 OPSD 的比较： OISD（64.75）显著优于 OPSD（56.63），后者是一种依赖特权真实解的最先进在线策略自蒸馏方法。这突显了 OISD 纯内部监督的有效性。
• 定性分析：
  • 训练动态： OISD 显示出稳定的奖励提升和初始的熵增加（表明更广泛的探索），随后趋于稳定。
  • 注意力可视化： 在训练过程中，OISD 的中间层逐渐与最终层对关键证据 token（例如端点、上确界）的注意力保持一致。
  • Logit-Lens 可视化： 与经常产生碎片化中间预测的 GRPO 相比，OISD 在推理关键 token 上表现出中间层与最终预测更早的对齐。

5. 意义与主张

本文声称 OISD 开启了对在线策略内部自蒸馏的研究。其主要意义在于：

• 消除外部依赖： 它在训练过程中无需外部教师、特权提示或真实解，即可提供密集的内部监督。
• 策略一致性： 它保留了单一执行策略用于 rollout 生成和优化，避免了离线策略方法的分布不匹配问题。
• 增强推理： 通过将中间表示与最终层的推理和注意力模式对齐，OISD 鼓励模型学习更强、全局一致性更好的中间表示。

作者指出了局限性，包括专注于具有可验证答案的数学推理，以及缺乏对最佳层选择的深入研究。他们建议未来的工作应探索向更复杂的推理任务和多层内部监督的泛化。