最弱环节:通过最坏情况约束强化学习蒸馏LLM推理能力
The Weakest Link: Distilling LLM Reasoning with Worst-Case Constrained Reinforcement Learning
精选理由
提出最弱环节约束蒸馏,兼顾答案正确性与推理逻辑可信度,助力小模型高质量部署。
AI 摘要
arXiv:2610.00332v1 公布类型:新 摘要:将大型语言模型(LLMs)的推理能力蒸馏到更小的学生模型中,是实现高效部署的核心挑战。当前方法面临一个根本性矛盾:单纯优化可验证的任务奖励(例如通过GRPO)会导致奖励劫持,即学生模型虽然得出正确的最终答案,却依赖存在缺陷的中间逻辑;而通过针对教师模型的软散度惩罚进行正则化(例如基于KL的蒸馏)则会削弱任务性能,更关键的是,允许学生模型在某一步出现严重逻辑违规,却用其他步骤中与教师模型的高度一致性加以补偿。我们认为,这种平均化方式从根本上偏离了推理的本质:思维链的有效性取决于其中最薄弱的环节。基于这一观察,我们将推理蒸馏 formulates 为一个约束强化学习问题:在对轨迹每个前缀上的教师模型对数似然施加最坏情况约束的条件下,最大化任务奖励。为避免对偶拉格朗日求解器的高昂成本,以及Saute等状态增广方法在测试时对教师模型的依赖,我们推导出一种未增广的约束MDP,其奖励变换能够保持硬约束语义,允许将策略梯度低方差地分解为单步项和长期项,并且能够证明在惩罚项极限下几乎必然满足最坏情况约束。通过在数学推理和代码生成任务上的大量实验,我们证明,该方法显著拓展了准确性—保真度帕累托前沿。在达到纯RL方法所具有的高最终答案正确性的同时,大幅减少教师约束违规,我们最终在所有评估设置中实现了最高的严格推理成功率。