利用反馈蒸馏提升Lean4定理证明的推理能力
速览
针对GRPO在推理模型后训练中面临的奖励稀疏和探索受限问题,研究者提出反馈蒸馏方法。该方法让模型在词元层面匹配由语言模型生成的特权反馈分布,从而提供更细粒度的监督和注入外部知识。在Lean4定理证明任务中,该方法生成的轨迹多样性更高,策略熵更大,且与GRPO结合使用效果更佳。
AI 深度解读
Distilling LLM Feedback for Lean Theorem Proving:深度解读
背景
在大型语言模型(LLM)的推理能力训练中,后训练(Post-training)阶段通常采用监督微调(Supervised Fine-Tuning, SFT)结合来自可验证奖励的强化学习(Reinforcement Learning from Verifiable Rewards, RLVR)的组合策略。其中,基于组相对策略优化(Group Relative Policy Optimization, GRPO)的算法因其高效性和稳定性,成为目前最主流的选择之一。
然而,尽管 GRPO 表现优异,它在处理复杂推理任务时仍面临三个显著瓶颈:
- 奖励稀疏(Sparse Rewards):在长序列生成中,只有最终答案正确才能获得奖励,中间步骤的错误难以被有效惩罚或引导。
- 探索受限(Limited Exploration):模型容易陷入局部最优,难以探索更优的解题路径。
- 模式崩溃(Mode Collapse):模型生成的轨迹多样性降低,倾向于重复固定的推理模式,限制了泛化能力。
针对这些问题,近期关于自蒸馏(Self-distillation)的研究为改进后训练提供了新思路。本文提出了一种名为 Feedback Distillation(反馈蒸馏) 的新训练方法,旨在通过引入语言模型生成的特权反馈(Privileged Feedback),在 token 级别对模型进行监督,从而克服上述缺陷。
核心内容
本文提出的 Feedback Distillation 是一种创新的训练范式。其核心思想是:让模型学习匹配其自身分布,但这种分布是基于由另一个语言模型(或同一模型的不同阶段)生成的“特权反馈”条件化后的结果。
1. 方法论机制
传统的强化学习依赖稀疏的标量奖励信号,而 Feedback Distillation 引入了细粒度的 token 级监督。具体而言:
- 特权反馈生成:利用一个具备更强推理能力或特定知识的语言模型,为当前模型的生成过程提供详细的反馈信号。这些反馈不仅包含最终的对错判断,还包含对推理步骤的修正建议或知识注入。
- 分布匹配:主模型被训练去拟合一个目标分布,该分布是在给定上述特权反馈的条件下生成的。这意味着模型不仅学习“答案是什么”,还学习“在特定反馈指导下,每个 token 应该如何生成”。
- 外部知识注入:由于反馈可以由外部模型生成,该方法天然具备注入外部领域知识的能力,弥补了模型内部知识的不足。
2. 在 Lean4 定理证明中的评估
为了验证该方法的有效性,研究团队将其应用于 Lean4 定理证明场景。Lean4 是一个形式化数学证明助手,对逻辑严密性和推理深度要求极高,是测试复杂推理能力的理想基准。
实验结果对比了 Feedback Distillation 与标准的 GRPO 方法:
- 轨迹多样性:Feedback Distillation 生成的推理轨迹保持了更高的多样性。
- 策略熵(Policy Entropy):该方法产生的策略熵更高,表明模型在决策时保留了更多的不确定性探索空间,避免了过早收敛到单一模式。
- Pass@k 缩放性能:在多次采样(Pass@k)场景下,Feedback Distillation 表现出更好的性能缩放趋势,即随着采样次数增加,找到正确证明的概率提升更显著。
3. 方法的互补性
研究发现,Feedback Distillation 与 GRPO 并非互斥,而是高度互补的:
- 单独使用 Feedback Distillation 可以提供丰富的初始分布和多样性。
- 单独使用 GRPO 可以通过强化学习进一步优化策略。
- 组合策略:使用 Feedback Distillation 的检查点(Checkpoint)来初始化 GRPO 的训练,其最终效果优于单独使用任何一种方法。这表明,先用蒸馏方法建立多样化的推理基础,再用强化学习进行精细优化,是提升复杂推理能力的有效路径。
关键要点
- 解决稀疏奖励问题:Feedback Distillation 通过 token 级别的监督信号,解决了传统 RLVR 方法中奖励信号稀疏、难以指导中间推理步骤的问题。
- 提升探索能力:相比 GRPO,该方法能维持更高的策略熵和生成轨迹多样性,有效缓解模式崩溃(Mode Collapse),使模型在复杂推理空间中探索更广泛的路径。
- 知识注入机制:通过引入外部语言模型生成的特权反馈,该方法能够动态注入外部知识,增强模型在特定领域(如形式化数学)的推理能力。
- 与 GRPO 互补:Feedback Distillation 适合作为 GRPO 的前置训练阶段。以 Feedback Distillation 的结果初始化 GRPO,能结合两者的优势,实现超越单一方法的性能提升。
- 适用场景:该方法特别适用于需要长链条推理、高逻辑严密性的任务,如 Lean4 定理证明,为复杂推理模型的后训练提供了新的可行方向。
意义与影响
Feedback Distillation 的提出,为大型语言模型的推理能力训练开辟了一条有前景的新途径。
首先,它挑战了当前以 GRPO 为主导的强化学习范式,证明了细粒度的监督信号(来自蒸馏的反馈)在复杂推理任务中可能比稀疏的强化学习奖励更有效。这对于解决大模型在数学、代码生成等需要严格逻辑链条任务中的“幻觉”和“死胡同”问题具有重要意义。
其次,该方法强调了多样性在推理训练中的价值。传统观点往往追求单一最优解的快速收敛,而本文结果表明,保持策略的多样性和高熵状态,有助于模型在面对未知或复杂问题时展现出更强的泛化能力和鲁棒性。
最后,这种“蒸馏+强化学习”的两阶段互补策略,为构建下一代推理模型提供了实用的工程范式。通过结合监督学习的稳定性和强化学习的优化能力,研究者可以更系统地提升模型在形式化验证、科学发现等高门槛领域的表现。随着形式化验证工具(如 Lean4, Coq)与 LLM 结合的深入,此类基于反馈蒸馏的训练方法有望成为标准流程的一部分。
