提炼大语言模型反馈用于 Lean 定理证明
推理模型的后训练通常结合监督微调和基于可验证奖励的强化学习,最常用的是 GRPO。然而,GRPO 存在奖励稀疏、探索受限和模式坍缩等问题。 基于近期自蒸馏工作,我们提出 反馈蒸馏,一种训练方法:模型在 token 级 匹配其自身在语言模型提供的特权反馈条件下的分布。反馈蒸馏提供 token 级监督,并能注入外部知识。 在 Lean4 定理证明任务上评估,我们发现反馈蒸馏比 GRPO 保持了更大的生成轨迹多样性,产生了更高的策略熵和更好的 pass@k 缩放。两种方法互补:从反馈蒸馏检查点初始化 GRPO 优于单独使用任一方法。总的来说,我们的结果表明了改进复杂推理后训练的一个有前景的方向。
论文精读
TL;DR Feedback Distillation 将语言模型的特权反馈蒸馏为 token 级分布匹配,解决了 GRPO 的稀疏奖励与探索不足问题,在 Lean 定理证明中提升了生成多样性与 pass@k,并能与 GRPO 互补叠加。
问题
问题背景
大语言模型后训练(post-training)在复杂推理任务(如数学定理证明)中,普遍采用监督微调(SFT) 与可验证奖励强化学习(RL) 的组合。其中 GRPO(Group Relative Policy Optimization)是当前主流 RL 算法,常用于提升代码、数学等领域的推理能力。
现有方法局限
然而,GRPO 存在三个关键瓶颈:
- 稀疏奖励:在定理证明等任务中,成功证明的轨迹极少,奖励信号几乎只在最终状态出现,导致有效学习样本稀缺;
- 探索不足:优化目标依赖当前策略的采样分布,缺乏对未知行为的系统鼓励,策略容易陷入局部最优;
- 模式坍塌(mode collapse):长期训练后,策略生成的证明路径趋于单一,丧失多样性,损害 pass@k 指标(k 次尝试内至少一次成功的概率)的随 k 增长。
这些局限使得纯 GRPO 训练难以充分利用验证器反馈,且训练过程不稳定,样本效率低。
难点与重要性
复杂推理(如 Lean4 定理证明)的本质挑战在于:
- 搜索空间巨大,正确路径依赖精细的逻辑组合,模型需同时探索多方向;
- 奖励信号二元且稀疏,单纯强化学习难以有效分配信用(credit assignment);
- 业界愈发关注推理模型的多样性与长期可扩展性,而非仅在固定采样预算下的一次性准确率。
从工程视角看,解决这些问题能提升 AI 在自动化证明、代码合成等安全敏感领域的可靠性,降低对海量人工标注的依赖。
行业类比
这类似于用代码大模型生成复杂算法:若只基于最终测试是否通过(稀疏奖励)进行强化学习,模型会反复生成同一类解法;而引入逐行注释或中间状态反馈,则可保持解法多样性,提升多轮尝试的成功率。
核心洞察
- Feedback Distillation 通过 token 级自蒸馏注入语言模型反馈,将传统 RL 的稀疏奖励转化为密集监督,从根本上缓解了 GRPO 在复杂推理中的探索不足与模式坍塌。与直接微调或 RL 不同,它让模型在保持自身分布一致性的同时吸收外部知识,训练信号更稳定,策略熵更高,为需要长序列多步决策的定理证明等任务提供了新的优化思路。
- Feedback Distillation 与 GRPO 并非替代关系,而是互补:先蒸馏知识再强化校正的两阶段范式能同时获得策略多样性和证明正确性。这种顺序组合(FD→GRPO)比任一单独方法显著提升了 pass@k,表明后训练中知识注入与奖励优化可以解耦,为设计更优的训练管线提供了工程上可借鉴的模式。
方法
方法流程
Feedback Distillation 是一种针对复杂推理任务的后训练方法,通过在 token 级别注入外部语言模型的特权反馈,引导策略模型进行自蒸馏式学习。
输入为定理陈述(如 Lean 4 代码片段)及初始证明目标。主流程分为三部分:
轨迹生成与反馈采集:
- 当前策略模型对给定问题采样多条证明轨迹(tactic 序列)。
- 同时,一个冻结的反馈模型(强语言模型)观察每条生成轨迹,在每一步提供 token 级别的特权反馈——包括正确性判断、策略建议或对 epectedic 状态的评价。反馈采用 Top-K 截断 以保证有效性。
条件自蒸馏训练:
- 策略模型被训练以匹配自身在受反馈条件约束下的 token 分布。具体而言,对于每一步,将反馈 token 拼接到上下文后,计算策略模型输出概率与目标分布(即策略模型在已包含反馈的条件下应产生的分布)之间的 KL 散度,并最小化该损失。
- 这一过程相当于将外部反馈的知识“蒸馏”进模型权重,使模型在无反馈时也能隐式编码类似推理模式。
密集监督与多样性保持:
- 与 GRPO 仅依赖最终可验证奖励(稀疏信号)不同,Feedback Distillation 提供逐 token 的监督,极大缓解了稀疏奖励带来的搜索效率问题。
- 由于训练目标是维持条件分布的熵,而非单纯强化高奖励行为,策略的策略熵(policy entropy)得以保持,避免模式坍缩,在后续的 pass@k 评估中展现更好的多样性。
输出是一个能够生成更多样证明路径、且单轮证明成功率更高的策略模型。该方法可与 GRPO 互补:先用 Feedback Distillation 预训练策略,再以 GRPO 微调,可获得最优效果。
与同类自蒸馏方法的差异:Feedback Distillation 不依赖自我合成数据或验证器奖励,而是直接利用外部 LM 的 token 级特权反馈进行密度更高的知识迁移,有效克服了 GRPO 的稀疏奖励与探索不足问题。
实验
实验设计
本文在 Lean4 定理证明 环境中评估 Feedback Distillation。基线包括标准 GRPO(群体相对策略优化)和仅使用监督微调的方法。核心指标为 策略熵(衡量生成轨迹的多样性)与 pass@k(k 次尝试内的证明成功率),并测试从 Feedback Distillation 检查点初始化 GRPO 的组合策略。训练数据由定理陈述和来自语言模型的 特权反馈(privileged feedback)构成,反馈提供 token 级监督信号,用于蒸馏自身分布。
关键发现
Feedback Distillation 相比 GRPO 产生更高策略熵,表明其对探索更充分,避免模式坍塌。在 pass@k 指标上,Feedback Distillation 展现出更好的 scaling 特性——随着采样次数增加,证明成功率提升更显著。更重要的是,先用 Feedback Distillation 预训练,再启用 GRPO 进行强化学习,所得模型性能超过任一单独方法,证实二者的互补性:Feedback Distillation 提供多样化初始化,GRPO 进行精细化验证。
与基线对比的深度解读
标准 GRPO 依赖稀疏奖励,容易陷入有限模式,而 Feedback Distillation 通过 token 级监督与外部知识注入(特权反馈)实现密集训练信号,本质上是一种 在策略知识迁移(on-policy knowledge transfer)。与纯自蒸馏方法不同,此处反馈来自更强大的语言模型,使得模型能内化“专家”的判断。结果表明,在复杂推理任务中,密集、细粒度的监督 可以缓解稀疏奖励的局限,为 post-training 流程提供新路径:先使用 token 级蒸馏建立基础探索能力,再用 RL 进行终端优化。
行业影响
落地场景
Feedback Distillation 适用于需要复杂多步推理的产品,尤其当推理过程可被可验证奖励函数评估时。最直接的场景包括:
- 形式化数学证明助手(如基于 Lean 4 的交互式证明环境):产品可生成更多样、成功率更高的证明路径。
- 代码生成与验证工具:在生成代码的同时注入可执行的测试或规范反馈,提升生成代码的正确性。
- 研究辅助 AI:帮助研究者探索定理证明、自动形式化等任务,降低对昂贵人类标注的依赖。
商业价值
- 降本:Feedback Distillation 提供 token 级别的密集监督,比稀疏奖励的 GRPO 样本效率更高,可减少训练时需要的交互步数和环境调用成本。在 Lean 证明这类昂贵环境中,这一点能显著节约计算资源。
- 增收与体验提升:更高的 pass@k 和策略熵意味着模型在多次采样时成功概率更高、输出更多样,可直接提升自动证明产品的可用性,扩展用户群体(如学生、研究人员、工程师)。产品差异化能力增强,可支撑付费订阅或企业服务。
- 风险控制:在需要高可靠性的场景(如金融合约形式化验证),联合使用 Feedback Distillation 与 GRPO 可降低模式坍塌风险,提供更稳健的推理能力。
与现有产品/工作流的接口
该方法可作为现有监督微调 + 强化学习流程的增强插件:
- 在 SFT 之后,插入 Feedback Distillation 阶段:利用 frozen 的“反馈模型”(如经过工具交互强化的 LLM)提供逐 token 的 KL 蒸馏目标,将外部知识与内部分布对齐。
- 后续可选继续 GRPO 微调:论文表明,从 Feedback Distillation 检查点启动 GRPO 可取得最优效果。
- 反馈源可灵活替换:除了外部强大模型,也可使用自我反馈(Self-Feedback Distillation),便于集成到现有持续训练或 RLHF 管道中。
具体落地用例
- 教育科技 – Lean 交互式习题平台
一家在线数学教育平台集成 Lean 4 环境为学生提供自动证明提示。使用 Feedback Distillation 训练的模型能给出多条证明路径,教师可审核并用于生成分级作业。pass@k 的提升意味着学生每次求助时,系统有更高概率返回至少一个有效方案,减少学生挫败感。 - 企业服务 – 智能合约安全审计
一家区块链安全公司用形式化方法验证智能合约。在自动证明生成流程中使用经 Feedback Distillation + GRPO 训练的模型,能更可靠地发现边界攻击并生成修复证明。token 级监督来自带工具交互的验证器反馈,确保模型既保持探索能力(高熵)又不偏离正确性约束。
局限
- **依赖外部反馈质量**:Feedback Distillation 需要由外部语言模型产生特权反馈,该反馈的质量直接影响训练效果。若反馈模型能力不足或产生噪声,蒸馏得到的 token 级分布可能误导策略优化。此外,在每次训练前获取反馈引入了额外的推理成本,尤其是在需要大量候选轨迹的场景下,可能成为计算瓶颈。论文主要使用 Llama 3.3 70B 作为反馈模型,未系统探讨不同规模或能力水平的反馈模型的影响。
- **实验领域局限**:方法只在 Lean 4 定理证明这一特定复杂推理任务上验证,数学定理证明虽有代表性,但自然语言推理、代码生成、科学问答等任务的特征(如动作空间、奖励稀疏性)差异显著。论文未在其他推理基准(如 MATH、HumanEval、ARC)上评估,方法是否具备领域泛化能力尚不清楚。定理证明的特定工具交互(Lean 编译器)也可能使结论难以直接迁移至通用推理场景。
- **对初始策略的依赖**:Feedback Distillation 是一种在线策略蒸馏方法,要求初始模型能生成足够多样化的候选轨迹,以便构建有意义的反馈条件分布。若初始策略过弱(如未经过监督微调),轨迹质量低下或模式单一,反馈蒸馏可能无法提供有效学习信号,反而加剧模式坍塌。该方法本质上是在现有推理能力基础上的改进手段,而非从零开始的训练方案,这限制了它对低资源任务的适用性。