Distilling LLM Feedback for Lean Theorem Proving¶

会议: NeurIPS 2026
arXiv: 2605.30861
领域: 模型压缩
关键词: 反馈蒸馏, 形式化定理证明, 强化学习, 策略熵, Lean 4

一句话总结¶

针对传统强化学习（如 GRPO）在形式化定理证明中面临的稀疏奖励与模式坍塌难题，本文提出反馈蒸馏（Feedback Distillation），让模型在自身采样的轨迹上逐 token 拟合条件化特权大模型反馈的教师分布，在保持高策略多样性的同时实现高效策略热启动。

研究背景与动机¶

形式化定理证明（如 Lean 4 环境）凭借其庞大的数学引理库（Mathlib 拥有超过 30 万条声明）与确定性的编译器验证反馈，已成为评估大语言模型复杂逻辑推理能力的理想试验场。当前面向复杂推理后训练的标准范式普遍采用“专家解监督微调（SFT）+ 基于可验证奖励的强化学习（RLVR，通常使用 GRPO）”。然而，这一两阶段流程在形式化定理证明中遭遇了严峻的瓶颈：离线构造的 SFT 数据集代价高昂且存在严重的策略外（off-policy）分布偏移，极易诱发分布外能力的灾难性遗忘；而后续的 GRPO 依赖全轨迹层面的稀疏结果奖励，无法提供细粒度的信用分配（credit assignment），尤其当单 prompt 采样的所有候选轨迹均失败时完全丧失学习信号。

更深层次的矛盾在于强化学习阶段的探索匮乏与模式坍塌。近期证据表明，GRPO 并不会扩展模型的解空间多样性，反而会快速收敛并坍塌到训练前已存在的少数高奖励局部模态上。对于形式化数学而言，探索 Mathlib 中极其多样化的策略分支与冷门辅助引理是突破高难度命题的前提，模式坍塌直接导致 pass@$k$ 随采样预算扩展的边际收益急剧递减。

为了突破纯稀疏奖励 RL 的探索局限并规避 SFT 的策略外失配，研究团队探寻一种兼具在线自监督与高频细粒度引导的后训练机制。本文的核心 idea 是提出反馈蒸馏（Feedback Distillation），在学生模型自身采样的在线轨迹上，通过语言模型反馈构建特权教师，以逐 token 的交叉熵损失引导学生拟合该教师分布，实现高策略熵的在线知识迁移并为 GRPO 提供优质初始化。

方法详解¶

整体框架¶

反馈蒸馏的核心在于构建一种具有“特权信息条件化”的自身增强教师。在每个训练步骤中，学生模型针对无证明的 Lean 定理陈述生成交互证明轨迹（包含 Lean 源码写入、编译检查与 Mathlib 正则检索等工具调用）。随后，特权反馈模型（如冻结的自身副本或更强的外部大模型）结合当前尝试轨迹与编译器错误输出，提炼出 1 至 3 条独立、无主语、命令式的改进指令（例如指定具体引理或拆分策略）。将该反馈拼入上下文作为特权信息输入给教师模型，教师模型计算该轨迹上每个生成 token 的前向概率分布。学生模型通过最小化与教师分布之间的逐 token KL 散度进行参数更新，而教师权重则通过指数移动平均（EMA）缓慢跟踪学生参数，防止训练发散。

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
    A["定理陈述输入 $x \sim \mathcal{D}$"] --> B["学生在线采样轨迹<br/>$y \sim \pi_{\theta}(\cdot \mid x)$（含交互工具）"]
    B --> C["外生/自省指令反馈提取<br/>$\mathcal{F}(y)$：基于尝试与编译结果提炼指令"]
    C --> D["特权条件化教师推理<br/>$\pi_{\mu}(\cdot \mid x, \mathcal{F}(y))$ 计算逐 token 概率"]
    D --> E["逐 Token 截断散度蒸馏<br/>Top-$K$ 交叉熵更新学生权重 $\theta$"]
    E --> F["指数移动平均权重更新<br/>$\mu \leftarrow \alpha \mu + (1-\alpha) \theta$ 平滑跟踪"]
    F -.->|循环闭环| D

关键设计¶

1. 交互式 Lean 工具调用与在线采样执行 为了突破单步生成的表达上限，模型被接入一个包含三大操作的 Lean 交互式验证环境：写入源码的 lean_write_file、编译代码并捕获报错/成功状态的 lean_check_file 以及在拥有 30 万条声明的 Mathlib 中检索关键引理的 rg_in_mathlib。在采样阶段，学生模型 $\pi_\theta(\cdot \mid x)$ 与 Lean 编译器展开长达 8,192 token 的多轮试错。这一设计确保了蒸馏所用的数据分布严格来自当前策略自身的生成轨迹（on-policy），从根本上避免了离线专家 SFT 引入的分布漂移，使得学习信号直接作用于模型自身暴露出的具体薄弱环节。

2. 外生/自省自然语言特权反馈构建机制 与依赖完整真实解答的传统特权信息不同，本方法将特权信息 $\mathcal{F}(y)$ 约束为具有普适指导意义的高层批评建议。反馈模型接收问题陈述 $x$、学生尝试轨迹 $y$（含所有工具执行输出）以及最终 Lean 编译器输出，提炼出 1 至 3 条自包含的指令性建议（格式为“Do X”或“Don't do X”，限制每条 10–50 词，且剥离“你/学生”等主观代词）。这种反馈既可以来源于更强大的第三方模型（如 Claude Opus 4.6），也可以来源于冻结的模型自身（Self-Feedback Distillation）。通过将反馈约束在抽象策略与引理提示层级，模型不会直接复制现成证明，而是学会将高级认知反思内化至自身权重，显著保留了探索过程中的认知不确定性表达（如 wait、perhaps、check 等词汇）。

3. 逐 Token 概率匹配与 Top-$K$ 截断蒸馏 为了克服轨迹级离散奖励导致的信用分配模糊问题，反馈蒸馏在序列每个位置施加逐 token 的 KL 散度监督。由于教师模型分布 $\pi_\mu$ 不依赖于待优化参数 $\theta$，最小化前向 KL 等价于最小化交叉熵损失。为了在词表大小高达数万的现代大语言模型上控制计算开销，团队将全词表求和截断为教师预测概率最高的前 $K$ 个 token（实践中设 $K=25$）： $$\hat{\mathcal{L}}_{\theta} = -\sum_{t=1}^{|y|} \sum_{a \in \mathrm{Top}K(\pi_{\mu})} \pi_{\mu}(a \mid x, \mathcal{F}(y), y_{<t}) \log \pi_{\theta}(a \mid x, y_{<t})$$ 在梯度反向传播时，采样操作使用 stop-gradient，梯度仅流经学生模型对轨迹中非工具输出 token 的条件概率对数。这一公式使得即便整条轨迹最终未能闭合证明，教师在关键推理分支位置给出的细粒度概率指引依然能够提供强有力的梯度推进。

4. 渐进式指数移动平均（EMA）教师权重跟踪 在自蒸馏循环中，若教师模型权重始终冻结，随着学生不断吸收反馈并拓展知识边界，教师将无法识别早期已学到的有效模式，甚至会对学生习得的新技能给出极低概率从而产生负向惩罚。为此，设计采用指数移动平均更新教师权重 $\mu$：每隔 5 个优化步，按以下规则平滑更新： $$\mu \leftarrow \alpha \mu + (1-\alpha) \theta$$ 平滑系数 $\alpha \in [0, 1]$ 严格调节适应速度与训练稳定性之间的平衡。经验验证表明，较激进的更新（较低的 $\alpha$）虽能带来初期的快速性能跃升，但极易造成后期的灾难性振荡；而完全冻结（$\alpha=1$）则会导致学习极度迟缓。选定 $\alpha=0.9$ 能在维持稳定策略演化的同时实现持续的知识内化。

损失函数 / 训练策略¶

整个训练过程分为两个互补阶段： 1. 反馈蒸馏预热阶段：使用 $1\times 10^{-5}$ 的学习率，每个批次包含 128 个 prompt，每 prompt 仅需单条 rollout 采样（有效批大小为 128），经过 50 步线性预热后使用截断交叉熵损失优化 100 至 400 步。此阶段的核心任务是维持高策略熵并建立对 Lean 环境交互与 Mathlib 引理引用的初步能力。 2. GRPO 强化精调阶段：以反馈蒸馏阶段的 checkpoint 作为热启动初始权重，切换为 GRPO 算法。为了精确对比优势函数，GRPO 每 prompt 采样 5 条 rollout（有效批大小达到 640），学习率下调至 $1\times 10^{-7}$，并辅以权重为 0.001 的 KL 散度正则惩罚，在保持已习得多模态探索能力的同时最大化最终通过率。

实验关键数据¶

主实验¶

主实验在 LeanWorkbook 子集上进行训练，并在包含 244 道竞赛级数学定理的 MiniF2F 测试集以及包含 256 道命题的 LeanWorkbook 测试集上评估 pass@1 准确率与策略熵表现。对比基线包括未经训练的原始基座模型、纯 GRPO 算法、从不同反馈蒸馏检查点继续训练的 GRPO 以及标准的离线 SFT+GRPO 流水线。

模型与训练方案	训练步数 / 检查点	MiniF2F Pass@1 (%)	LeanWorkbook Test Pass@1 (%)	策略熵 (Policy Entropy)	采样数 / Prompt
Qwen3-8B (Base Zero-shot)	0	2.0%	~3.1%	较高 (未对齐)	—
Qwen3-8B + 纯 GRPO	400 步	23.0%	46.5%	低 (发生模式坍塌)	5
Qwen3-8B + 反馈蒸馏 (Claude FD)	400 步	22.8%	44.8%	显著高于 GRPO	1
Qwen3-8B + GRPO over FD-100	FD 100 步 + GRPO	27.5%	51.2%	中高	5 (RL期)
Qwen3-8B + GRPO over FD-200	FD 200 步 + GRPO	29.8%	54.0%	中高	5 (RL期)
Qwen3-8B + GRPO over FD-300	FD 300 步 + GRPO	30.5%	55.4%	中	5 (RL期)
Qwen3-8B + GRPO over FD-400	FD 400 步 + GRPO	31.0%	56.8%	均衡	5 (RL期)
Qwen3.5-9B (Base Zero-shot)	0	11.0%	~14.5%	较高	—
Qwen3.5-9B + GRPO over SFT	SFT 200 步 + GRPO	~38.5%	~63.0%	中低	5 (RL期)
Qwen3.5-9B + GRPO over FD-200	FD 200 步 + GRPO	39.2%	64.5%	保持良好	5 (RL期)

消融实验与反馈源对比¶

下表汇总了在 Qwen3.5-9B 上针对四种不同特权反馈信息源的对比效果，以及在 200 步前置训练下纯反馈蒸馏与后续 GRPO 的表现差异；第二部分展示了 EMA 衰减系数 $\alpha$ 与工具调用接口对最终 MiniF2F 准确率的影响。

实验维度	具体配置	MiniF2F 准确率 (%)	相对表现与核心观察
特权反馈源对比	自身冻结反思 (Self-Feedback)	22.5%	无外部信息，但仍产生自省提升
(Qwen3.5-9B, 200步FD)	原始编译器输出 (Lean Output)	27.0%	信号较生硬，缺乏高层反思提示
	真实参考证明 (Ground Truth Proof)	36.8%	信息密度极高，但易抑制认知不确定性表达
	Claude Opus 4.6 自然语言建议	38.0%	抽象策略指导最佳，有效避免机械照抄
EMA 平滑系数 $\alpha$ 消融	$\alpha = 0.5$ (激进更新)	16.2%	前期学习极快，但 150 步后剧烈震荡崩盘
(Qwen3-8B, 300步)	$\alpha = 0.9$ (基准设置)	22.8%	早期稳定提升，持续保持策略多样性
	$\alpha = 1.0$ (完全冻结教师)	11.5%	教师无法识别已学知识，学习极其缓慢
交互工具机制消融	完整交互工具 (`write/check/rg`)	22.8%	多轮编译器交互驱动闭环验证修正
(Qwen3-8B Claude FD)	移除工具交互 (单步端到端生成)	9.4%	缺少符号反馈，无法定位 Lean 编译断点

计算资源与采样开销核算（累计前 200 步）¶

资源核算维度	Claude 反馈蒸馏 (FD)	纯 GRPO 算法	GRPO / FD 比值
每步 Prompt 的采样数 (Rollouts)	128	640	5.0×
前 200 步总采样数 (Total Rollouts)	25,600	128,000	5.0×
墙上时钟训练总时间 (Wall-clock)	15.85 小时	27.74 小时	1.8×
生成阶段耗时 (Generation Time)	5.15 小时	11.39 小时	2.2×
外部反馈获取耗时 (Feedback API)	1.39 小时	0 小时	—
生成响应总 Token 数 (Response Tokens)	137.1 M	801.0 M	5.8×
Lean 工具调用总次数 (Tool Calls)	423.6 k	2,064.5 k	4.9×
单轨迹平均工具调用次数	16.5 次	16.1 次	1.0×

关键发现¶

策略熵衰减抑制与 Pass@$k$ 扩展优势：在 400 步训练点处，虽然纯反馈蒸馏与纯 GRPO 的 pass@1 几乎一致（22.8% vs 23.0%），但随着推断时采样次数 $k$ 的增加，反馈蒸馏的 pass@$k$ 曲线展现出显著更陡峭的上升斜率。这与训练过程中持续记录到的更高策略熵和 Mathlib 新引理发现曲线（图 3）完全吻合，证明反馈蒸馏没有像 GRPO 那样出现模式坍塌，而是真实保留了多样化的解题探索空间。
阶段互补性与最佳热启动：直接使用 GRPO 从头训练容易陷入局部最优，但以 FD-400 作为起点时，GRPO 的 MiniF2F 成绩从 23% 跃升至 31%，在 LeanWorkbook 测试集上从 46.5% 跃升至 56.8%。这表明反馈蒸馏非常适于担当第一阶段的“在线预热器”，在保留多样性的同时完成对 Lean 语法的对齐。
认知词汇（Epistemic Verbalization）的保留机制：消融分析显示，若直接提供 Ground Truth 答案作为特权信息，模型在思维链中使用的审慎词汇（如 wait, perhaps, check）会在几十步内雪崩式消失，导致过度自信与分布外退化；而基于批评建议的自然语言反馈则有效保全了这种认知审慎性，且生成的平均思维链长度显著更充实。

亮点与洞察¶

在线无偏差的特权知识蒸馏范式：传统 SFT 依赖离线构造的专家解答，与模型自身采样存在严峻的分布偏移。反馈蒸馏将外部专家知识降维为针对模型自身错误行为的高层指令，通过自蒸馏机制施加逐 token 的稠密监督，在保持严格 on-policy 的同时实现了与离线 SFT 匹敌的知识迁移效率。
规避模式坍塌的熵保持特性：强化学习中基于优势函数的二值过滤天然惩罚所有失败尝试，极易驱动模型收敛至狭隘的确定性策略。反馈蒸馏将失败轨迹通过特权信息转化为局部的正向概率校准，显著延缓了策略熵的坍塌，为下游测试时计算扩展（test-time compute scaling）铺平了道路。
解耦的自然语言跨模型通信：与需要共享 Tokenizer、对齐隐层维度或传输庞大 Logits 张量的经典知识蒸馏不同，反馈蒸馏仅通过纯文本交互沟通教师模型。这使得前沿闭源旗舰大模型（如 Claude Opus）能以极低的 API 成本充当导师，极大降低了系统架构门槛。

局限与展望¶

长期训练过程中的 EMA 稳定性挑战：尽管设定 $\alpha=0.9$ 能够保障数百步内的平稳收敛，但若进一步延长训练步数，学生与教师之间的反馈自激仍可能引发晚期的策略震荡乃至性能退化。未来需要更鲁棒的动态正则化或自适应步长机制来支撑长周期训练。
外部 API 带来的额外经济成本：虽然反馈蒸馏在本地 GPU 算力消耗上仅为 GRPO 的约 $1/1.8$（节省了大量的轨迹生成计算），但调用高阶外部模型生成 2.5 万次反馈消耗了约 1,100 美元的 API 账单。如何通过高性价比小模型或优化自省（Self-Feedback）策略抹平性能差距是落地关键。
任务领域的单一性：目前实验主要集中在 Lean 4 形式化定理证明这单一验证闭环环境中。该范式未来能否无缝拓展到代码生成、非形式化竞赛数学乃至通用 Agent 推理任务，仍有待更广泛的跨领域验证。

评分¶

新颖性: ⭐⭐⭐⭐⭐ 首次将自然语言批判性反馈转化为特权自蒸馏目标，优雅打通了在线强化学习与外部高层知识注入的边界。
实验充分度: ⭐⭐⭐⭐⭐ 兼顾 LeanWorkbook 与 MiniF2F 测试集，详尽提供了策略熵演化、Mathlib 引理发现率、算力账单核算及认知词汇衰减分析。
写作质量: ⭐⭐⭐⭐⭐ 动机阐述深刻，对 SFT、GRPO 与自蒸馏三者的优劣分析切中肯綮，公式推导严谨自洽。
价值: ⭐⭐⭐⭐⭐ 为大模型复杂推理领域的后训练范式（尤其是解决 RLVR 阶段的探索缺乏与模式坍塌难题）开辟了极具实用价值的新方向。