本文提出RLKD,一个基于强化学习的知识蒸馏框架,通过生成结构奖励模型(GSRM)将教师模型推理中的隐式多分支结构传递给学生模型,实验表明其在数学和问答任务上显著优于SFT和传统RL方法。
Reinforcement Learning, Large Language Model, Reasoning, Pre-training, Fine-tuning
Shicheng Xu, Liang Pang, Yunchang Zhu, Jia Gu, Zihao Wei, Jingcheng Deng, Feiyang Pan, Huawei Shen, Xueqi Cheng
State Key Laboratory of AI Safety, Institute of Computing Technology, CAS, University of Chinese Academy of Sciences, Huawei Inc.
Generated by grok-3
Background Problem
大型语言模型(LLMs)在复杂推理任务中表现出色,但其推理能力的训练成本高昂,资源受限的团队难以直接开发具备类似能力的模型。监督微调(SFT)通过蒸馏教师模型的推理路径为较小模型提供了一种快捷方式,然而,SFT仅能让学生模型模仿教师模型表面的推理路径,而无法传递真实推理中隐含的多分支结构(由元推理和求解交织构成),导致学生模型陷入僵硬的模仿而非真正的推理。本文旨在解决这一问题,通过提出一种新的蒸馏方法,使学生模型能够学习教师模型推理中的隐式多分支结构,从而实现更真实的推理能力。
Method
本文提出了RLKD,一个基于强化学习(RL)的知识蒸馏框架,旨在将教师模型推理中的隐式多分支结构传递给学生模型。其核心方法包括以下步骤:
- 生成结构奖励模型(GSRM):首先,通过GPT-4o的上下文学习(ICL)和监督微调数据集,训练一个生成模型(基于Qwen2.5-7B-Instruct),将推理路径分解为多个元推理-求解对(meta-reasoning-solving pairs),其中元推理负责选择子问题,求解负责回答子问题;然后,设计一个结构化奖励机制,通过比较学生和教师模型在每个步骤的元推理和求解内容的对齐度,计算奖励值。
- RL训练:结合GSRM的奖励与任务结果奖励(如数学任务的准确率),采用Group-based Relative Policy Optimization(GRPO)进行强化学习训练,指导学生模型在每个推理步骤中从多个潜在路径中进行更好的采样,而非仅仅模仿教师的表面输出。 批判性思考:GSRM的分解过程依赖于预训练模型和人工设计的指令,可能引入分解不准确或主观性问题,尤其是在复杂推理任务中,元推理和求解的界限可能并不清晰。此外,奖励机制的线性比较和提前终止策略可能过于简单,无法完全捕捉多分支结构的复杂性。
Experiment
实验在数学推理(AIME24, MATH-500)和研究生级问答(GPQA-Diamond)数据集上进行,采用pass@1和pass@k(k>1)作为评估指标,旨在测试模型的推理能力和多路径探索能力。实验设置包括三种对比:1)RLKD与SFT-RL流水线在Qwen2.5-Math-7B上的对比,显示RLKD仅用0.1%数据即可超越SFT-RL;2)RLKD与PPO、GRPO等RL基线在DeepSeek-R1-Distill-Qwen-7B上的对比,表明RLKD能进一步提升SFT蒸馏模型的性能;3)RLKD与SFT在领域偏移和跨领域设置下的对比,显示RLKD在训练过程中持续提升性能,而SFT性能下降。结果表明,RLKD在pass@k指标上提升显著,表明其能帮助学生模型学习多分支推理结构,而非单纯记忆教师路径。此外,推理路径多样性分析也支持这一结论。 批判性思考:虽然实验结果看似令人印象深刻,但训练数据量极小(仅0.1%),可能导致结果的统计显著性不足,难以推广到更大规模数据集或模型。此外,实验未涉及更大参数规模模型(如70B+),其方法的可扩展性存疑。数据集选择和评估指标虽然合理,但部分数据集(如MATH-500)可能已出现在某些基线模型的训练集中(如Qwen2.5-Math-7B-Instruct),可能对结果产生偏差。
Further Thoughts
RLKD的思路为LLM推理能力的蒸馏提供了一个新颖的视角,尤其是在资源受限场景下,通过强化学习捕捉隐式结构可能比传统的SFT更有效。然而,我认为其方法可能面临实际应用中的挑战:首先,GSRM的奖励计算依赖于元推理-求解对的准确分解,这在不同领域或任务中可能需要大量定制化工作,限制了通用性。其次,强化学习本身的不稳定性和高计算成本可能在更大规模模型上成为瓶颈。未来可以探索是否能结合自监督学习或对比学习方法,减少对奖励模型的依赖,同时提高训练效率。此外,RLKD与最近的一些工作(如基于思维树或图结构的推理方法)可能存在互补性,是否能将显式结构与隐式多分支结构结合,进一步提升推理能力,是一个值得深入研究的方向。