OISD: 语言模型的在策略内部自蒸馏

arXiv cs.LG 论文

摘要

OISD是一种新的自蒸馏框架,通过在GRPO后训练期间将预测信号从最终层转移到中间层,使用logit和注意力对齐来提升语言模型推理能力。

arXiv:2605.29089v1 公告类型:新 摘要:最近基于强化学习(RL)的后训练方法主要使用稀疏的结果级奖励来优化最终输出策略,而很大程度上忽略了中间表示中编码的预测信号。在本文中,我们引入了一种称为在策略内部自蒸馏的新范式,并提出了OISD框架,该框架通过将在策略预测信号从最终层转移到中间表示来提升推理能力。在rollout和组相对策略优化(GRPO)优化过程中,最终层既作为策略,又作为选定的中间层的分离内部教师,通过两种互补机制引导这些中间层与其对齐:logit对齐(传递高级推理行为,即如何思考)和注意力对齐(强制执行一致的注意力模式,即关注何处),两者均不需要外部特权信息。我们的OISD与GRPO一起,使用带符号的优势加权的Jensen-Shannon对齐来蒸馏信息丰富的中间表示,同时在统一的行动策略下保持策略一致性。实验结果表明了OISD的有效性,在四个数学推理任务上,相较于强大的推理RL基线,取得了显著且一致的改进。代码将在https://github.com/THE-MALT-LAB/OISD发布。
查看原文
查看缓存全文

缓存时间: 2026/05/29 09:16

# OISD: 语言模型的同策略内部自蒸馏

来源:https://arxiv.org/html/2605.29089
刘新宇1,Darryl Cherian Jacob1,周阳1,王金东2,何攀1†
1 奥本大学
2 威廉与玛丽学院
† 通讯作者

###### 摘要

近期的强化学习(RL)后训练方法主要利用稀疏的结果级奖励来优化最终输出策略,而很大程度上忽略了中间表示中编码的预测信号。在本文中,我们引入了一种名为“同策略内部自蒸馏”的新范式,并提出了 **OISD** 框架,该框架通过将同策略预测信号从最终层传递到中间表示,从而提升推理能力。在数据收集和组相对策略优化(GRPO)期间,最终层既作为策略,又作为所选中间层的分离内部教师。所选中间层通过两种互补机制被引导与之对齐:logit对齐(传递高级推理行为,即“如何思考”),以及注意力对齐(强制最终层到所选中间层的注意力模式一致,即“看哪里”),两者均无需外部特权信息。我们的 **OISD** 与 GRPO 结合,采用带符号优势加权的 Jensen–Shannon 对齐来蒸馏信息丰富的中间表示,同时在统一的行动策略下保持策略一致性。实验结果表明,**OISD** 的有效性,在四个数学推理任务上相比强大的推理 RL 基线取得了显著且一致的改进。代码将在 https://github.com/THE-MALT-LAB/OISD 发布。

---

**OISD: 语言模型的同策略内部自蒸馏**

刘新宇1,Darryl Cherian Jacob1,周阳1,王金东2,何攀1†
1 奥本大学
2 威廉与玛丽学院
† 通讯作者

## 1 引言

参阅图注

**图1:四种蒸馏范式的比较:(a) 离线策略蒸馏,(b) 在线策略蒸馏,(c) 同策略自蒸馏,以及 (d) 同策略内部自蒸馏(我们的方法)。比较突出了教师-学生角色、数据生成以及蒸馏应用位置方面的差异。**

大型语言模型(LLM)在数学推理、编程和复杂指令遵循方面取得了显著进展[Yao 等人,2022 (https://arxiv.org/html/2605.29089#bib.bib2);Wei 等人,2022 (https://arxiv.org/html/2605.29089#bib.bib31)]。这一进展的一个关键驱动力是面向推理的后训练,即在预训练和监督微调之后进一步优化模型,以生成更准确、连贯和可靠的推理轨迹[Shao 等人,2024 (https://arxiv.org/html/2605.29089#bib.bib16);Guo 等人,2025 (https://arxiv.org/html/2605.29089#bib.bib17)]。

近期使用可验证奖励的强化学习(RLVR),尤其是 GRPO 风格的训练,显著改进了面向推理的后训练,而无需密集的人工监督[Shao 等人,2024 (https://arxiv.org/html/2605.29089#bib.bib16);Guo 等人,2025 (https://arxiv.org/html/2605.29089#bib.bib17);Zheng 等人,2025a (https://arxiv.org/html/2605.29089#bib.bib41);Yu 等人,2026 (https://arxiv.org/html/2605.29089#bib.bib40)]。然而,这些方法仍然高度注重最终输出,仅用稀疏的结果奖励优化最终的行动策略,而很大程度上忽略了中间预测、推理过程和内部表示。

知识蒸馏[Hinton 等人,2015 (https://arxiv.org/html/2605.29089#bib.bib3);Bengio 等人,2015 (https://arxiv.org/html/2605.29089#bib.bib7)] 沿着推理轨迹提供了密集的 token 级监督,补充了稀疏的结果奖励。在经典的离线策略蒸馏中,学生从固定的教师生成轨迹中学习,这会产生分布不匹配,因为训练轨迹来自教师而非学生自身的生成[Agarwal 等人,2024 (https://arxiv.org/html/2605.29089#bib.bib13);Xu 等人,2025 (https://arxiv.org/html/2605.29089#bib.bib42)]。在线策略蒸馏通过直接监督学生生成的样本[13]减少了这种差距,但大多数方法仍然依赖于外部教师模型。近期的同策略自蒸馏方法通过用特权轨迹、反馈或辅助提示增强同一个模型,移除了外部教师[Zhao 等人,2026 (https://arxiv.org/html/2605.29089#bib.bib23)]。然而,监督仍然依赖于外部注入的特权信息,而非模型自身的内部计算。

现有的离线策略和在线策略蒸馏方法主要侧重于提高模型准确性和推理能力,对利用模型自身的内部机制进行蒸馏的探索有限。然而,先前的研究[Yang 等人,2025b (https://arxiv.org/html/2605.29089#bib.bib14);Zhu 等人,2025 (https://arxiv.org/html/2605.29089#bib.bib11)] 表明,推理表示在各层之间系统性地演变:早期层保留更广泛的候选预测,中间层整合上下文和任务相关信息,而深层则逐渐收敛到最终的 next-token 决策。这种层次组织表明,不同层在推理过程中扮演着不同的功能角色。

同时,logit lens 框架[nostalgebraist,2020 (https://arxiv.org/html/2605.29089#bib.bib9)] 通过使用 unembedding 矩阵将中间表示投影到 token 空间,提供了初步的见解。这揭示了 Transformer 架构中的残差流包含丰富的预测信息,这些信息在各层和模块之间逐步演变[Dai 等人,2022 (https://arxiv.org/html/2605.29089#bib.bib34);Gupta 等人,2025 (https://arxiv.org/html/2605.29089#bib.bib33)]。所有这些见解都激励我们利用模型自身的内部计算作为蒸馏的内在监督源。

在本文中,我们提出了 **OISD**,一个用于 LLM 推理的同策略内部自蒸馏框架。在训练期间,**OISD** 保持最终层作为数据生成和 GRPO 优化的唯一行动策略,同时将其用作分离的内部教师。遵循 logit lens 框架[Belrose 等人,2023 (https://arxiv.org/html/2605.29089#bib.bib10)],选定的中间层被投影到 token 空间,形成同一个策略样本上的辅助内部学生。学生从两个互补信号中学习:logit 对齐(传递高级推理行为,即“如何思考”),以及注意力对齐(强制最终层到所选中间层的注意力模式一致,即“看哪里”),两者均无需外部特权信息。这样的公式化鼓励模型学习更强的中间表示。在推理时,教师模型被用作最终模型。**OISD** 引入了密集的内部监督,无需外部教师、特权提示或多个行动策略,因此得名“同策略内部自蒸馏”。我们的 OISD 与其他方法的比较如图 1 (https://arxiv.org/html/2605.29089#S1.F1) 所示。

### 1.1 贡献

我们的贡献总结如下:

- • 我们引入了一种称为“同策略内部自蒸馏”的新范式,用于推理强化学习,该范式利用模型自身的内部计算作为监督,同时保留用于数据生成和优化的单一行动策略。
- • 我们提出了 **OISD**,一个同策略内部自蒸馏框架,该框架使用最终层作为分离的内部教师,并使用中间层作为辅助内部学生,将同策略预测信号从最终层传递到中间表示。
- • 我们开发了两个互补的内部对齐目标:用于“如何思考”的 logit 对齐(通过转移预测信念),以及用于“看哪里”的注意力对齐(通过转移证据路由行为)。
- • 我们在四个数学推理基准上证明了相比强基线有显著且一致的改进。

我们希望这项研究能够开启对同策略内部自蒸馏的探索,并鼓励未来在这个方向上的研究。

## 2 相关工作

**用于推理的知识蒸馏。** 近期基于 PPO 的从人类反馈中强化学习(RLHF)方法为偏好优化提供了一个标准框架[Schulman 等人,2017 (https://arxiv.org/html/2605.29089#bib.bib47);Ouyang 等人,2022 (https://arxiv.org/html/2605.29089#bib.bib48)]。随后,面向推理的强化学习方法,如 GRPO、DeepSeek-R1、DAPO、GSPO、Qwen3 和 RAIF,利用可验证奖励和基于组的策略优化来改进长上下文推理[Shao 等人,2024 (https://arxiv.org/html/2605.29089#bib.bib16);Guo 等人,2025 (https://arxiv.org/html/2605.29089#bib.bib17);Zheng 等人,2025a (https://arxiv.org/html/2605.29089#bib.bib41);Yang 等人,2025a (https://arxiv.org/html/2605.29089#bib.bib21);Qin 等人,2026 (https://arxiv.org/html/2605.29089#bib.bib20);Yu 等人,2026 (https://arxiv.org/html/2605.29089#bib.bib40)]。然而,监督仍然在很大程度上是稀疏的、结果级别的,并且集中在最终的行动策略上。

为了提供更密集的推理监督,先前的工作[Hinton 等人,2015 (https://arxiv.org/html/2605.29089#bib.bib3);Bengio 等人,2015 (https://arxiv.org/html/2605.29089#bib.bib7)] 经常将教师分布或推理轨迹蒸馏到学生中。然而,离线策略蒸馏引入了训练-推理不匹配,因为学生在教师生成的轨迹上训练,但在推理时必须从其自身的内部状态进行推理。

在线策略蒸馏通过从学生生成样本并使用教师的每个 token log 概率作为学生实际访问状态的密集监督,减少了这种不匹配[Agarwal 等人,2024 (https://arxiv.org/html/2605.29089#bib.bib13);Xu 等人,2025 (https://arxiv.org/html/2605.29089#bib.bib42)]。更近期的自蒸馏方法通过使模型能够使用特权信息或更丰富的反馈进行自我监督,减少了对外部教师的依赖[Agarwal 等人,2024 (https://arxiv.org/html/2605.29089#bib.bib13);Hübotter 等人,2026 (https://arxiv.org/html/2605.29089#bib.bib5);Shenfeld 等人,2026 (https://arxiv.org/html/2605.29089#bib.bib4)]。例如,OPSD [Zhao 等人,2026 (https://arxiv.org/html/2605.29089#bib.bib23)] 利用特权解决方案,SDPO [Hübotter 等人,2026 (https://arxiv.org/html/2605.29089#bib.bib5)] 利用环境反馈,OPCD [Ye 等人,2026 (https://arxiv.org/html/2605.29089#bib.bib24)] 将上下文条件行为蒸馏到模型权重中,而 CODI [Shen 等人,2025 (https://arxiv.org/html/2605.29089#bib.bib29)] 将显式推理压缩到连续表示中。

我们的方法将教师完全置于模型自身内部:分离的最终层在同一个策略样本上监督一个中间层,同时仍然是唯一的行动策略。

**用于推理的内部信号。** 早期研究工作,如深度监督网络(DSN)[Lee 等人,2015 (https://arxiv.org/html/2605.29089#bib.bib43)],为将辅助损失附加到中间表示上设定了更广泛的优化先例。最近,逐层读出表明,在到达最终层之前,中间 Transformer 表示已经包含了有意义的预测结构。Logit-lens 和 tuned-lens 分析将隐藏状态解码为词汇分布,使得跨层的预测轨迹直接可见[nostalgebraist,2020 (https://arxiv.org/html/2605.29089#bib.bib9);Belrose 等人,2023 (https://arxiv.org/html/2605.29089#bib.bib10)]。

后来,DoLa 在解码期间使用层间对比,而 LayerSkip 训练中间退出以实现高效推理和自推测解码[Chuang 等人,2024 (https://arxiv.org/html/2605.29089#bib.bib27);Elhoushi 等人,2024 (https://arxiv.org/html/2605.29089#bib.bib28)]。特别地,BuPO 首先研究了在初始训练阶段优化细粒度内部层策略以引导语言模型的整体策略[Tan 等人,2025 (https://arxiv.org/html/2605.29089#bib.bib15)]。然而,作者观察到过度对齐可能导致性能崩溃,表明中间策略优化应仅应用于有限的训练迭代。其他方法将逐层信号用于分析、解码、提前退出或推理效率,而非作为辅助监督[Kapadia 等人,2026 (https://arxiv.org/html/2605.29089#bib.bib6)]。

推理监督也可以来自推理过程本身,而不仅仅是最终答案。CODI 通过自蒸馏将显式推理压缩到连续的内部表示中[Shen 等人,2025 (https://arxiv.org/html/2605.29089#bib.bib29)],而步骤级验证方法则研究对中间推理步骤的监督[Lightman 等人,2024 (https://arxiv.org/html/2605.29089#bib.bib35);Zheng 等人,2025b (https://arxiv.org/html/2605.29089#bib.bib46);Song 等人,2025 (https://arxiv.org/html/2605.29089#bib.bib45)]。注意力通过揭示信息在 token 预测期间如何路由,提供了另一种过程级信号。例如,MoLSAKI 在推理蒸馏期间将逐步注意力从教师转移到学生[Chen 等人,2025 (https://arxiv.org/html/2605.29089#bib.bib30)],而 RAL 通过策略梯度训练直接优化内部注意力分布[Li 等人,2026 (https://arxiv.org/html/2605.29089#bib.bib44)]。

我们的 **OISD** 在同策略设置中开发了两个互补的内部对齐目标:logit 对齐监督“如何思考”,注意力对齐监督“看哪里”,并且两者都在同一个样本上从最终层蒸馏到中间层。

## 3 同策略内部自蒸馏

在本节中,我们介绍 **OISD**,一个同策略内部自蒸馏框架,该框架在推理强化学习期间将最终层用作分离的内部教师,并将中间层用作辅助内部学生。我们进一步引入两个互补的内部对齐目标,用于监督中间推理表示。

### 3.1 Transformer 深度作为垂直推理轴

思维链(CoT)推理通过额外的推理 token 暴露了中间计算,沿时间维度*水平*扩展了推理[Wei 等人,2022 (https://arxiv.org/html/2605.29089#bib.bib31)]。更近期的隐式推理方法通过跨推理步骤演化隐藏状态或循环激活,而不是完全用自然语言表达推理,进一步扩展了这一范式[Hao 等人,2024 (https://arxiv.org/html/2605.29089#bib.bib49);Shen 等人,2025 (https://arxiv.org/html/2605.29089#bib.bib29)]。遵循 [Zhu 等人,2025 (https://arxiv.org/html/2605.29089#bib.bib11)],这些方法主要通过隐藏状态随时间演化来扩展推理能力。

**OISD** 则在一个由 L 个堆叠层组成的标准仅解码器 Transformer 中,探索了一种互补的*垂直*推理范式。与其通过额外的 token、循环激活或隐式展开步骤来扩展推理,**OISD** 利用了模型深度中已经形成的中间计算。最终层监督

相似文章

OPRD:在策略表示蒸馏

Hugging Face Daily Papers

OPRD提出了一种新的知识蒸馏方法,该方法在策略部署期间跨层对齐学生和教师的隐藏状态,消除了来自词空间KL估计的采样方差。实验表明,OPRD在数学推理基准(AIME 2024/2025、AIMO)上优于输出空间基线,同时速度快1.44倍,内存使用减少54%。

OPID: 同策略技能蒸馏用于智能体强化学习

Hugging Face Daily Papers

OPID提出了一种同策略技能蒸馏框架,从完成的轨迹中提取密集后见监督,将基于结果的强化学习与词元级自蒸馏相结合,以提高语言智能体在多轮任务上的训练效率和性能。