新智元报道
编辑：乔杨耳朵
LLM 数学水平不及小学生怎么办？CMU 清华团队提出了 Lean-STaR 训练框架，在语言模型进行推理的每一步中都植入 CoT，提升了模型的定理证明能力，成为 miniF2F 上的新 SOTA。
如果想训练 LLM 证明定理的能力，你会怎么做？
既然模型可以通过海量语料学会生成文本，那如果我们能喂给它足够数量的形式证明数据，定理证明能力自然水到渠成？
然而，我们看到的事实是，无论用符号形式还是自然语言，GPT 等大模型的推理能力都不如人意。
就像一样，AI 再聪明却依旧会在简单的算术上犯蠢。
然而，LLM 的数学能力弱，不代表自动化的定理证明器对数学没用。
前段时间
陶哲轩也曾在采访中强调，使用。这是一股不可小觑的力量。
最近，CMU 和清华的一项研究就致力于让 LLM 的「自然语言思维链」和 Lean 的形式化证明结合在一起。
论文地址：https://arxiv.org/abs/2407.10040
论文提出，Lean、Coq、Isabelle 等基于形式语言（代码）的自动化证明方法，忽略了大量可能对推理过程有用的「非形式化信息」。
比如，每个证明步骤之前的潜在思维过程是必不可少的，但却不会形式化地体现在最终的公式和代码中。
比如，图 1 中右侧的推理思路，在左侧的证明步骤中完全「无处安放」。
因此，作者提出了 Lean-STaR 训练框架，让语言模型既学会逐步推理的思维，也学会形式化的证明方式。
这意味着，需要将自然语言和形式语言交织在一起，也将「思考」和「证明」的过程交织在一起。
方法：Lean-STaR
顾名思义，Lean-STaR 这个方法同时结合了之前的两项成果——Lean 和 STaR。
Lean 是一种函数式编程语言，可以用作交互式定理证明器（Interactive Theorem Prover）。
项目主页：https://lean-lang.org/
这是由 Leonardo de Moura 在微软研究院期间发起的开源项目，目前已经更新到 Lean 4。
比如，要想形式化证明，能从n≤m推断出n+k≤m+k，就可以用 Lean 写为如下形式（图6）：
首先给出一种高级的「策略」（tactic，图中所示为归纳策略k），将当前要证明的目标状态简化为多个子目标（下图中的 case 0 和 case ih）。
这些子目标又会形成新的「状态」（state）。当所有子目标都得到证明时，我们就给出了定理的完整证明。
STaR 则是来源于斯坦福和谷歌研究院在 2022 年发表的一篇论文，全称是「自学推理器」（Self-Taught Reasoner）。
论文地址：https://arxiv.org/abs/2203.14465
其基本思想就是用到了「自举法」（bootstrapping）。
首先根据训练数据中的问题和答案，提示语言模型，生成能解释答案的「原理」（rationale）。
之后，再用这个包含了问题、答案和原理的混合数据集对 LM 进行微调，提升模型的推理能力（图1）。
Lean-STaR 模型的微调也是采用了「渐进优化」的思路，逐步将以上两个相关工作的成果融合在一起，完善底层的策略预测模型。模型构建的流水线如图 4 所示。
直接策略预测（Direct Tactic Prediction）
首先，将定理证明问题简单地定义为马尔科夫决策过程（MDP）。
从这个角度来看，证明过程是状态 si、策略 ai 和奖励 ri∈R等 3 个变量组成的轨迹(s1，a1，r1) (s2，a2，r2)⋯其中，ITP（比如 Lean）用于提供每个新状态 si+1。
在这种经典设置中，证明定理的过程包括向 LM 提供状态s，让模型M生成策略() 。
因此，可以使用仅包含成功证明轨迹的基本数据集
对基本模型进行监督微调，得到 SFT 模型。
思维增强策略预测（Thought-augmented Tactic Prediction）
结合之前所述的研究动机，我们假设「潜在想法」可以提高模型的策略预测能力，因此引入一个表示「思维」的隐变量 ti，然后将模型扩展为：
此时，根据状态预测下一个策略的分布可以表示为：
如果用这种方式预测，我们就需要一个全新的数据集
用于训练模型M，然而 Lean 给出的证明步骤只包含 si 和 ai。
论文的解决方法是：借助一个强大的语言模型G（如 GPT-4）作为「预言家」，让它在给定当前状态 si 和真实策略 ai 的情况下生成 ti，从而创建出新的数据集 DT（即图 4 中的 CoT Dataset）。
作者将这种方法称为「回顾性原理生成」（retrospective rationale generation）。
将 SFT 模型在 DT 数据集上再进行一次微调后，就得到了第一个思维增强的策略预测模型 Lean-CoT。
自举思维增强定理证明（Bootstrapping Thought-augmented Theorem Proving）
在 Lean-CoT 模型的基础上，作者提出，可以应用「专家迭代」（expert iteration）方法进一步提升性能。
具体来说，从初始的 Lean-CoT 模型 M0 以及初始数据集D开始，让 M0 对每个问题进行K次采样，每次采样都会产生一个证明轨迹 [(s0,t0,a0)，(s1,t1,a1)，⋯，(sn,tn,an)]，之后过滤出成功的证明轨迹并去重，得到新数据集 D1。
接下来，在数据集 DT∪D1 上进一步微调 M0 模型以得到 Lean-STaR 模型 M1。
将上述过程进行多次迭代，即可不断更新 Lean-STaR 模型。
评估实验
为了测试 Lean-STaR 的具体性能，研究使用了可用的最佳开放语言模型 Lean 语料库 (InternLM2-Math-base-7b) 上进行预训练，并遵循 Lean 的 Mathlib 作为底层训练集的标准实践。
首先以 LeanDojo Benchmark 4 v9 作为监督微调(SFT)数据集，包含超过 23.1 万个示例，进行 1 轮微调以获得 SFT 模型。
之后从数据集中随机选择 17256 个不同的成功证明轨迹，并使用 GPT-4-0125 模型注释出 52438 个想法，并且执行两次专家迭代。
实验在 MiniF2F 基准上评估 Lean-STaR，使用了与之前的实验工作类似的评估设置，但主要使用的是采样方法（sampling）而不是最佳优先搜索（best-first search）来进行评估。
实验结果表明，回顾性原理生成和专家迭代都显著提高了模型的定理证明能力。
实验结果
实验的主要结果如下表所示，Lean-STaR 比之前基于 Lean 的 SOTA 模型有了显著的改进。
例如，在类似的推理预算下，同样使用 best-first search，Lean-STaR 从 InternLM2 的 30.3% 提升至 34.8%，也同样高于使用 GPT-4 的 COPRA（30.7%）。
随着计算预算的增加，Lean-STAR 的性能进一步提升至 36.1%。
思维增强改进定理证明
Lean-STaR 的第一阶段在思维增强的合成数据集上进行微调，训练模型来交替生成思维和策略。
此阶段的微调模型（在表 1 中表示为 Lean-CoT）达到了 32.8% 的通过率，高于此阶段之前的模型（表示为 SFT，29.5%）。
可以证明，第一阶段的思维增强能提高语言模型的定理证明能力，即使对于已经专门用于生成 Lean 策略的语言模型（例如 SFT）也依旧成立。
自举法（Bootstrapping）进一步改进
Lean-STaR 的第二阶段包括使用当前语言模型生成新的思维和策略，保存正确结果，并结合初始数据集进行训练。
从表 1 结果来看，每次迭代都会提高模型的定理证明性能，从 32.8%（初始模型）到 34%（迭代 1 次后的L-STR）再到 34.8%（迭代 2 次后的L-STR）。
此外，我们发现该模型可以通过额外采样进一步改进，将采样的K值加倍后，分数能进一步提升至 36.1%。
无 CoT 的专家迭代实验
表 5 显示了没有 CoT 的专家迭代结果（即仅使用状态和策略，没有思维增强），对比 Lean-CoT 和 Lean-STaR 的表现。
仅用专家迭代时，准确率就达到了 43.0%，低于 Lean-STaR (45.5%)。
这表明 Lean-STaR 的性能提升不仅仅来自于专家迭代的使用，思维增强也有不可忽略的效果。
问题类型与难度
MiniF2F-test 中的问题有多个来源，包括 AIME、AMC、IMO 等数学竞赛以及 MATH 数据集，并进行了手动形式化处理。
这些问题可能有不同的难度和类型。表 2 展示了成功证明的问题数量，按类型和难度划分。
Lean-CoT 提高了解决所有类别难题的表现，尤其是数学竞赛中的难题。
除了这些改进之外，Lean-STAR 的改进主要集中在数论方面。
搜索和抽样预算
表 4 说明了问题通过率与搜索规模或抽样预算S×K的关系。
实验发现，Lean-STAR 性能与K值的大小成正比，特别是当K值相对较大时。
对比前两列和 Lean-STaR 可以发现，附带思维的额外采样能提高定理证明性能，而没有思维的额外采样可能会饱和。
作者猜测，可能是因为「思维」增加了输出的多样性，并有助于对定理证明空间进行探索。
因此，Lean-STaR 更具可扩展性（就推理阶段算力而言），并且可以通过额外的专家迭代进一步改进。
更强基础模型和更多数据实验
实验还使用了更强的语言模型 InternLM2-Math-plus-7b 训练 LeanSTaR，来测试不同语言模型性能的影响。
不仅基座模型更强，为数据集注释「思维」的模型也从 GPT-4 升级到 GPT-4o，生成了 1.4 万条想法。
实验只执行一次专家迭代，收集了大约 6 万条（证明状态、思维、下一步策略）正确的数据，命名为「STaR 数据集」。
在 STaR 数据集上进一步微调得到 Lean-STAR 模型，其测评结果如表 3 所示，可以看到 Lean-STaR 仍然比基线有了显著的改进。
结论和局限性
研究团队提出了 Lean-STaR，这是一种新颖的方法，通过将思维链 (CoT) 原理集成到每个证明步骤中，显著增强了语言模型用形式化数学语言进行定理证明的能力。
方法首先根据 ground truth 回顾性地为证明步骤生成「原理」，然后微调语言模型，训练模型学会生成「原理」并预测后续策略，从而得到 Lean-CoT 模型。
然后使用专家迭代进一步改进该模型，根据被证明为正确的采样结果进行微调，并使用 Lean solver 进行验证。
研究的贡献包括引入第一个思维增强的定理证明数据集，并证明专家迭代可以进一步提高性能。得到的模型在 miniF2F 测试上取得最新 SOTA，将通过率从 30.3% 提高到 36.1%。
这些进步不仅提高了自动化定理证明的准确性，而且还提供了一个可扩展且高效的框架来促进对数学的理解，这可能会对教育、科学发现和程序验证产生重大影响。
方法的主要限制在于，其性能可能受限于计算可扩展性，实验中用于微调 Lean-CoT 和 Lean-STaR 模型的数据集都不是很大。
需要注意的是，专家迭代的速度也存在严重瓶颈，会受限于 Lean ITP 的缓慢进程。
此外，使用 GPT-4 生成合成数据成本较大，并可能引入偏差。
参考资料：
