让语言模型内部自我教学,提升推理能力
OISD: On-Policy Internal Self-Distillation of Language Models

- 用最终层作教师,指导中间层学习思维和注意力模式
- 在4个数学推理任务上显著优于现有强化学习基线
- 无需外部数据,适合提升大模型推理性能的研究者
近期基于强化学习的后训练方法主要依赖稀疏的输出层面奖励优化最终策略,却忽视了中间表示中蕴含的预测信号。本文提出一种新范式——就地策略内部自蒸馏(OISD),通过将最终层的预测信号迁移至中间层,提升模型推理能力。在滚动和组相对策略优化(GRPO)过程中,最终层同时充当策略与独立内部教师,指导选定中间层通过两种互补机制对齐:逻辑值对齐(转移高层推理行为)和注意力对齐(强制一致的关注模式),且无需外部特权信息。OISD与GRPO结合,采用带符号优势加权的Jensen–Shannon对齐,以蒸馏有信息量的中间表示,同时保持统一行动策略的一致性。实验表明,OISD在4个数学推理任务上均实现显著且一致的性能提升,优于强基线。代码将于 https://github.com/THE-MALT-LAB/OISD 公开。
原文摘要 · Abstract (English)
Recent reinforcement learning (RL) post-training approaches primarily optimize the final output policy using sparse outcome-level rewards, while largely overlooking predictive signals encoded in intermediate representations. In this paper, we introduce a new paradigm called on-policy internal self-distillation and propose the OISD framework, which improves reasoning by transferring on-policy predictive signals from the final layer to intermediate representations. During rollout and Group Relative Policy Optimization (GRPO) optimization, the final layer acts as both the policy and a detached internal teacher for selected intermediate layers, which are guided to align with it through two complementary mechanisms: logit alignment, which transfers high-level reasoning behaviors (how to think), and attention alignment, which enforces consistent attention patterns (where to look) from the final layer to the selected intermediate layer, both without requiring external privileged information. Our OISD, together with GRPO, employs signed advantage-weighted Jensen--Shannon alignment to distill informative intermediate representations while preserving policy consistency under a unified acting policy. Experimental results demonstrate the effectiveness of OISD, with substantial and consistent improvements over strong reasoning RL baselines across four mathematical reasoning tasks. The code will be released at https://github.com/THE-MALT-LAB/OISD
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。