论文
最薄弱的环节:用最坏情况约束强化学习提炼 LLM 推理
The Weakest Link: Distilling LLM Reasoning with Worst-Case Constrained Reinforcement Learning
摘要
将大语言模型 (LLM) 的推理能力提炼给较小的学生是高效部署的核心挑战。当前的方法面临着一种根本性的张力:纯粹针对可验证的任务奖励(例如通过 GRPO)进行优化会导致奖励作弊行为,学生通过有缺陷的中间逻辑得出正确的最终答案,同时对教师进行软发散惩罚(例如基于 KL 的蒸馏)进行规范化会削弱任务绩效,更重要的是,允许学生在一个步骤中补偿严重的逻辑违规,而在其他步骤中与教师的高度一致。我们认为,这种平均从根本上与推理的本质不一致:思想链的有效性取决于其最薄弱的环节。受这一观察的启发,我们将推理蒸馏表述为一个受约束的强化学习问题,其中任务奖励最大化,受到沿轨迹每个前缀的教师对数似然的最坏情况约束。为了避免双拉格朗日求解器的高昂成本和测试时教师对状态增强方法(例如 Saute)的依赖性,我们推导了一种未增强的约束 MDP,其奖励变换保留了硬约束语义,允许将低方差策略梯度分解为单步和长期项,并且几乎可以肯定在惩罚限制中满足最坏情况约束。通过对数学推理和代码生成任务的大量实验,我们证明了我们的方法显着扩展了准确性保真度帕累托前沿。通过匹配纯强化学习的高最终答案正确性并大幅减少教师约束违规,我们最终在所有评估的设置中实现了最高的严格推理成功率。