用于 Lean 定理证明的大语言模型反馈蒸馏 Distilling LLM Feedback for Lean Theorem Proving
提出 Feedback Distillation 方法,改进定理证明训练
前置知识
Self-Distillation (自蒸馏)
自蒸馏是一种训练方法,通过最小化学生模型和教师模型之间的差异来优化模型,其中教师模型接收特权信息而学生模型不接收。特权信息可以是环境输出、正确答案或其他在推理时不可用但训练时可用的信号。训练损失通常是 KL 散度,让学习分布向教师分布靠拢。与传统蒸馏不同,教师和学生通常共享架构和初始化。
本文的 Feedback Distillation 是自蒸馏的一种变体,理解自蒸馏的基本原理对掌握本文方法至关重要。
GRPO (Group Relative Policy Optimization)
GRPO 是一种强化学习算法,通常用于可验证奖励的场景。它为每个提示采样多个输出,然后通过对比它们的奖励来优化策略。算法的核心思想是最大化高奖励样本的概率,同时惩罚低奖励样本。在定理证明中,GRPO 使用 Lean 编译器作为验证器,成功的证明获得正奖励,失败的证明获得零奖励或负奖励。
GRPO 是本文的主要对比基线,理解其工作原理和局限性(如模式坍缩、多样性减少)是理解本文创新的关键。
Lean 定理证明器
Lean 是一个交互式定理证明器,使用依赖类型理论来形式化数学证明。Lean 4 是最新版本,提供了更好的元编程能力和性能。Mathlib 是 Lean 的主要数学库,包含超过 30 万个数学声明(定义、定理、引理等)。在 Lean 中证明定理需要编写使用策略的证明脚本,策略告诉证明器如何操作证明状态。Lean 编译器可以验证证明脚本是否正确。
Lean 是本文的实验环境,理解其工作方式、Mathlib 的规模以及证明验证机制对理解实验设置和结果至关重要。
Pass@k 和 Policy Entropy
Pass@k 是代码生成和定理证明中的标准评估指标,表示当模型生成 k 个候选时至少有一个正确的概率。对于定理证明,如果 k 个尝试中至少有一个通过 Lean 编译器验证,则计为成功。Policy entropy 策略熵衡量模型输出分布的随机性或多样性,熵越高表示模型产生更多样化的输出,而熵越低表示模型集中在少数模式上。
本文的关键发现之一是 Feedback Distillation 比 GRPO 维持更高的策略熵,从而实现更好的 pass@k 扩展,理解这些指标对评估结果意义重要。
EMA (Exponential Moving Average)
指数移动平均是一种平滑技术,用于更新目标网络的权重。在本文中,教师模型的权重通过 EMA 更新,其中 mu 是教师权重,theta 是学生权重,alpha 在 0 到 1 之间是衰减系数。alpha 等于 1 等价于冻结教师,alpha 等于 0 等价于完全复制学生。EMA 让教师逐步跟踪学生的学习进度,防止训练不稳定。
EMA 在 Feedback Distillation 中起关键作用,alpha 的选择直接影响训练稳定性和性能,本文实验展示了不同 alpha 值的效果。
研究动机
现有用于推理模型的训练流程通常结合监督微调和强化学习,但在 Lean 定理证明等复杂推理任务中存在显著局限性。首先,监督微调阶段受限于生成微调数据集的难度和成本,而且因为它在离策略演示上训练,可能导致狭窄微调分布之外的灾难性遗忘。更重要的是,GRPO 等强化学习算法并不会放大模型答案的多样性或在训练期间促进探索,而是坍缩到 RL 训练之前已经存在的高奖励模式。这在形式化数学中特别成问题,因为探索大型引理库(如 Lean 的 Mathlib,包含超过 30 万个声明)对于成功至关重要。此外,GRPO 仅提供轨迹级奖励,当所有样本都失败时无法提供学习信号。实验数据显示,在 LeanWorkbook 测试集上,单独使用 GRPO 训练 Qwen3.5-9B 仅达到 59% 的 pass@1 准确率,而且训练过程中会出现不稳定性。
本文的目标是本文的目标是提出一种能够改进 Lean 定理证明训练的新方法,该方法应该能够在保持输出多样性的同时注入外部知识,并提供细粒度的信用分配。具体来说,作者希望设计一种训练机制,能够提供 token 级别的监督信号,而不是像 GRPO 那样将奖励均匀分配到整个轨迹。该方法还应该能够在不需要大量专家演示数据的情况下,帮助模型发现新的能力,如调用适当的策略和引用 Mathlib 中的相关引理。最终目标是在 LeanWorkbook 和 MiniF2F 等基准测试上超越现有的 GRPO 方法。
与已有工作不同的是,本文的独特切入角度是将反馈蒸馏引入定理证明训练。与现有的使用环境输出或真实答案作为特权信息的自蒸馏方法不同,本文使用由冻结的模型副本或第三方 LLM 产生的文本反馈作为特权信息。这种反馈模型分析学生的尝试、验证器的输出,然后提供可操作的指导,如识别证明策略错误、标记策略或引理的误用、鼓励更好的工具使用等。这种方法的一个关键优势是反馈可以是提示式的(如使用 Real.pi 而不是 Unicode pi 符号),而不是提供完整解决方案,从而保持模型的探索能力。本文还系统地研究了不同反馈源(Claude、真实答案、Lean 编译器输出)的效果,并分析了训练不稳定性与 EMA 参数的关系。
核心方法
Feedback Distillation 的整体思路是让模型学习匹配自己在接收到特权反馈条件下的分布。直观上,就像老师给学生反馈时学生能学到更多一样,模型也可以从对自己尝试的反馈中学习。技术路线上,本文使用两个相同架构的模型实例:学生模型和教师模型,两者都从相同的预训练 LLM 初始化。教师模型接收特权信息(即反馈)而学生模型不接收。训练损失是沿着学生采样的轨迹计算教师和学生分布之间的逐 token KL 散度。这种 on-policy 方法提供了细粒度的信用分配,因为每个 token 都有独立的监督信号,而不像 GRPO 那样将奖励均匀分配到整个轨迹。教师模型的权重通过指数移动平均更新,让它能够逐步跟踪学生的学习进度。
Feedback Distillation 的核心创新在于使用 LLM 生成的文本反馈作为自蒸馏中的特权信息,而不是环境信号或真实答案。这与已有方法的本质区别在于反馈可以是提示式的、部分的指导,而不是完整的解决方案。例如,当学生错误地使用 Unicode pi 符号而不是 Lean 的 Real.pi 时,反馈可以是使用 Real.pi 的建议,而不是给出整个证明。这种反馈方式的优势在于它既提供了有用的指导,又保持了模型的探索能力。另一个关键创新是 Top-K 截断,为了减少计算成本,本文将词汇表上的求和截断到教师分布下概率最高的 K 个 token(在实践中使用 K 等于 25),这大大提高了计算效率。
方法步骤详情
Feedback Distillation 的方法步骤完整描述如下。首先,从数据集 D 中采样一个问题 x(例如 Lean 语句),学生模型生成一个尝试 y。然后,反馈函数 F 接收尝试 y 并返回文本反馈信号。反馈模型可以是更强的 LLM(如 Claude Opus 4.6),也可以是学生模型的冻结副本(Self-Feedback Distillation)。教师模型定义为在问题和反馈条件下生成。训练损失是逐 token KL 散度的期望。在实践中,使用等价的交叉熵损失形式,并采用 Top-K 截断减少计算。教师权重每五个梯度步通过 EMA 更新,其中 alpha 是衰减系数。
技术新颖性
Feedback Distillation 的技术新颖性体现在多个方面。首先,它将 LLM 反馈引入自蒸馏框架,这种反馈可以作为间接的知识蒸馏渠道,强模型提供信息而不需要提供完整解决方案,这与标准蒸馏不同。其次,这种方法是 inherently on-policy 的,因为训练信号是在学生的生成上计算的,这种性质被证明可以改善泛化性并减少与 SFT 等离策略方法相关的灾难性遗忘。第三,知识通过自然语言反馈而不是 logit 匹配传递,这消除了托管反馈模型或共享其 tokenizer 的实际需求。第四,本文进行了系统的分析,比较了不同反馈源的效果,研究了 EMA 参数对训练稳定性的影响,分析了反馈表述的影响,并观察到批评形式的反馈比完整解决方案更少抑制认识论表达(模型使用不确定词语如 wait、hmm、perhaps)。
实验结果
本文的核心发现通过多个实验展示。首先,GRPO over Feedback Distillation 性能最佳,实验显示在 Feedback Distillation 训练 200、300 和 400 步后继续使用 GRPO 训练,性能优于 GRPO 单独训练。在 LeanWorkbook 测试集上,使用 Qwen3.5-9B 作为基座模型,GRPO 单独训练达到 59% 的 pass@1 准确率,而 Feedback Distillation+GRPO 达到 75%。其次,Feedback Distillation 在整个训练过程中维持显著高于 GRPO 的熵,表明策略保留了更多多样性。这转化为更好的测试时扩展:在 GRPO 和 Feedback Distillation 达到类似 pass@1 的检查点,Feedback Distillation 随着 k 增加扩展更好,确认更高的熵反映的是答案多样性而不是均匀噪声。第三,Feedback Distillation 导致模型在成功的证明中使用新的 Lean 策略和 Mathlib 引理的速度更快、更持续。第四,关于反馈源,实验显示 Claude 和真实答案作为反馈源优于 Lean 编译器输出,这是预期的,因为两种设置都提供对更强信息源的访问。Self-Feedback Distillation 仍然有效:当反馈模型与训练的模型相同时,将 GRPO 与 Self-Feedback Distillation 结合仍导致比 GRPO 单独更高的峰值性能。最后,关于 EMA 参数,实验显示更激进的更新(较低的 alpha 值)导致早期学习更快,而 alpha 等于 1 避免不稳定性但学习缓慢,在本文的训练预算下,alpha 等于 0.9 提供最佳权衡。
查看结构化数据
| 任务 | 指标 | 本文 | 基线 | 提升 |
|---|---|---|---|---|
| LeanWorkbook 定理证明 | pass@1 测试准确率 | 75% (Feedback Distillation+GRPO) | 59% (GRPO 单独) | 提升 27% |
| LeanWorkbook 定理证明 | 策略熵 | 显著高于 GRPO | GRPO | 保持输出多样性 |
| LeanWorkbook 定理证明 | Mathlib 引理发现率 | 更快更持续的新引理发现 | GRPO | 累积不同引理数量增长更快 |
| MiniF2F 基准测试 | 测试准确率 | 与 GRPO 相当或更高 | GRPO | 取决于训练步骤 |
局限与改进
本文的局限性分析包括作者承认的和观察到的。首先,作者承认当前结果尚未与大规模训练的最先进 SFT+GRPO 管道匹配,这表明还有改进空间。其次,训练稳定性仍然是 Feedback Distillation 的挑战,即使使用较大的 alpha 值,长期不稳定性仍然存在,开发更强大的稳定机制是扩展到更长训练运行的自然下一步。第三,反馈表述对训练性能有重要影响,虽然结构化指导的提示提高了稳定性和最终性能,但也减慢了早期学习,提示工程仍有改进空间。第四,本文实验主要在 Lean 定理证明环境,虽然在形式化数学中令人鼓舞,但将 Feedback Distillation 扩展到其他形式化和非形式化推理任务仍有待验证。最后,关于样本效率,虽然 Feedback Distillation 在训练早期阶段比 GRPO 更高效(GRPO 的批量大小是 Feedback Distillation 的五倍),但更彻底的样本效率评估留给未来工作。
独立分析的弱点
Feedback Distillation 的独立分析弱点包括:第一,训练不稳定性仍然是一个问题,即使在 alpha 等于 0.9 的设置下,长期训练仍可能出现性能波动或崩溃,改进方向可以是开发自适应 EMA 调度、梯度裁剪或其他稳定技术。第二,反馈模型的质量直接影响训练效果,使用较弱模型作为反馈源时性能下降,改进方向可以是开发反馈质量的评估指标和自动反馈优化方法。第三,Top-K 截断虽然提高了计算效率,但可能丢失一些低概率但重要的 token 信息,改进方向可以是自适应 K 值或重要性加权采样。第四,当前方法主要针对定理证明,在其他推理任务(如代码生成、问题求解)上的泛化能力有待验证,改进方向可以是跨任务的统一反馈框架。第五,反馈生成需要额外的计算开销,改进方向可以是缓存反馈、批量化处理或开发轻量级反馈模型。
未来方向
未来研究方向包括作者提出的和基于成果可延伸的。作者明确提出的方向是:开发更强大的稳定机制以扩展到更长训练运行;将 Feedback Distillation 扩展到其他形式化和非形式化推理任务;进行更彻底的样本效率评估。基于成果可延伸的方向包括:探索不同反馈表述的影响,如结构化反馈 vs. 自然语言反馈、即时反馈 vs. 延迟反馈;研究反馈蒸馏与其他训练方法的结合,如与课程学习、对比学习的集成;开发自动化反馈生成和质量评估系统;探索反馈蒸馏在多模态推理(如视觉-语言任务)中的应用;研究反馈蒸馏对不同规模模型的影响,从小模型到大模型的扩展规律;开发反馈蒸馏的理论框架,提供收敛性和泛化性保证;探索反馈蒸馏在持续学习和在线学习场景中的应用。
复现评估
复现评估方面,本文使用 verl 代码库进行实验,使用 Kimina Lean Server 评估 Lean 证明的正确性。实验在配备 8 块 NVIDIA H200 GPU 的单节点上运行,GRPO 实验平均每个训练步需要 550 秒,而 Feedback Distillation 需要每个步 300 秒。数据集使用 LeanWorkbook 的子集,包括 10000 个已知证明的语句用于训练,256 个语句的测试集用于评估,以及 MiniF2F 测试分割的 244 个竞赛级问题。基座模型是 Qwen3-8B 和 Qwen3.5-9B,在初始化时在 MiniF2F 上分别得分 2% 和 11%。反馈模型使用 Claude Opus 4.6。超参数包括 EMA 衰减系数 alpha 等于 0.9、Top-K 截断 K 等于 25、批量大小等(详细参数见附录 A.1)。虽然论文描述了实验设置,但代码是否开源未明确说明,数据集和基座模型是公开的,算力需求较高(8 块 H200 GPU),完整复现有一定难度。
论文图表