【论文解读】LLaMA-Berry: 将“refine”当作动作解决数学推理

1st author

• trotsky1997 (Di Zhang)
• ‪Jianbo Wu‬ - ‪Google Scholar‬
• Personal Homepage
• Tong Che

paper: [2410.02884] LLaMA-Berry: Pairwise Optimization for O1-like Olympiad-Level Mathematical Reasoning

code: trotsky1997/MathBlackBox

5. 总结

LLaMA-Berry无疑是LLM在高级数学推理领域的一个重要进展。它巧妙地将MCTS的强大搜索能力与针对LLM特性设计的Self-Refine机制和PPRM奖励模型相结合，形成了一套端到端的解决方案。尤其是PPRM和EBC的设计，为如何从LLM中提取和利用相对偏好信息提供了有益的启示。

从更宏观的视角看:

• 结构化推理的趋势: LLaMA-Berry代表了从简单的“提示-生成”模式向更结构化、更深思熟虑的推理过程的转变。MCTS提供了一个自然的框架来组织这种多轮次的思考和优化。OpenAI的Q*和O1项目可能也在探索类似的思路。
• 偏好学习的重要性: 对于复杂的认知任务，让模型学习相对偏好可能比学习绝对价值更容易，也更符合人类的判断方式。DPO等技术的成熟为这一方向铺平了道路。
• 计算开销与效率: MCTS虽然强大，但计算开销不容忽视。未来的工作需要在效果和效率之间取得更好的平衡，例如通过更智能的剪枝策略、学习的启发式函数，或者将部分搜索能力内化到LLM本身。
• 通用性: 虽然本文聚焦于数学推理，但其核心思想——将复杂问题分解为“状态（完整方案）-动作（优化方案）-评估（偏好比较）”的MCTS框架——有潜力推广到其他需要深度探索和规划的领域，如代码生成与优化、科学发现、复杂规划等。

LLaMA-Berry的探索非常有价值。它不仅在奥林匹克数学这一极具挑战性的任务上取得了令人印象深刻的结果，更重要的是，它为我们展示了一条提升LLM深度推理能力的、颇具前景的技术路径。期待未来能看到这一框架在更大规模模型上的应用，以及在降低计算复杂度和拓展应用领域方面的进一步突破。

1. 思想

这篇论文试图解决一个核心难题：如何让大型语言模型（LLMs）在复杂、多步骤的数学问题（特别是奥林匹克竞赛级别）上实现更高效、更准确的推理？

传统的思路，如思维链（CoT）或简单的自洽性（Self-Consistency），在面对这类问题时，往往因为搜索空间的巨大和缺乏有效反馈而陷入困境，容易产生次优甚至错误的解。即便是引入了自我反思（Self-Refine）等机制，也可能因为反馈的偏差而导致优化方向错误。

LLaMA-Berry的核心思想可以概括为：

1. 全局探索替代局部贪心：将整个解的生成和优化过程视为一个探索过程，而不是简单的逐步生成。这自然地导向了如蒙特卡洛树搜索（MCTS）这样的方法。
2. 以“优化”为行动：在MCTS的框架中，每一步的“行动”不再是生成下一个token或步骤，而是对一个完整的候选解进行“自我优化”（Self-Refine），产生一个可能更好的新解。
3. 相对评估优于绝对打分：人类判断两个解的优劣，往往比给一个解打绝对分数更容易、更准确。因此，论文提出用一个成对偏好奖励模型（Pairwise Preference Reward Model, PPRM） 来判断两个候选解中哪一个更优，而不是给每个解打一个标量分数。
4. 从局部偏好到全局排序：有了成对的偏好关系，还需要一个机制将其汇总成全局的解质量排序，以便MCTS做出更明智的决策。

这套组合方案，借鉴了AlphaZero/MuZero中MCTS的强大搜索能力，并针对LLM的特性进行了适配。

2. 方法

LLaMA-Berry的框架主要由三个组件构成：SR-MCTS、PPRM 和 EBC。

2.1 SR-MCTS (Self-Refine Monte Carlo Tree Search)

这是整个系统的骨架，将“自我优化”过程融入MCTS的节点扩展中。

• 基本框架: 将问题求解形式化为马尔可夫决策过程 (MDP)。

  • $s$: 状态 (State), 代表一个完整的解方案 (a complete solution)。
  • $a$: 动作 (Action), 代表对当前解方案的重写或优化操作 (a rewriting action)。
  • $Q(s,a)$: 在状态 $s$ 执行动作 $a$ 的预期回报 (expected reward)。

• 核心循环:

  1. 选择 (Selection): 从根节点开始，递归地使用 UCT (Upper Confidence Bound applied to Trees) 规则选择子节点，直到达到一个叶节点 $s_i$。UCT平衡了对已知高回报节点的利用（exploitation）和对未充分探索节点的探索（exploration）。
     $$a^{\ast }=\arg \underset{a^{\prime }\in A(s)}{\max }(Q(s,a^{\prime })+c\cdot \sqrt{\frac{\ln N(s)}{N(s,a^{\prime })}})$$
     其中：

     • $A(s)$ 是状态 $s$ 下所有可能的动作集合。
     • $Q(s,a^{\prime })$ 是从状态 $s$ 执行动作 $a^{\prime }$ 的估计价值。
     • $N(s)$ 是父节点 $s$ 被访问的次数。
     • $N(s,a^{\prime })$ 是从状态 $s$ 执行动作 $a^{\prime }$ 后达到的子节点被访问的次数。
     • $c$ 是探索常数，用于调整探索的权重。

  2. 扩展 (Expansion) 与 自我优化 (Self-Refine): 对选定的叶节点（即一个候选解 $s_i$），执行自我优化过程。

     • LLM首先生成一个“批判” (Critique) $C_i=\mathcal{C}(s_i)$，识别当前解 $s_i$ 中的数学错误、逻辑缺陷等。
     • 然后，LLM根据原始解 $s_i$ 和批判 $C_i$ 生成一个“重写” (Rewrite) 的新解 $s_{i+1}=\mathcal{R}(s_i,C_i)$。
     • 这个新解 $s_{i+1}$ 作为搜索树中的一个新子节点被添加。

  3. 评估 (Evaluation): 使用后续将介绍的 PPRM 和 EBC 方法，评估新生成的解 $s_{i+1}$ 的质量，得到其价值估计 $Q(s_{i+1})$。

  4. 反向传播 (Backpropagation): 将新节点 $s_{i+1}$ 的评估价值 $Q(s_{i+1})$ 从该节点逐层向上反馈，更新其所有祖先节点的访问次数 $N(\cdot )$ 和价值估计 $Q(\cdot ,\cdot )$。

2.2 PPRM (Pairwise Preference Reward Model)

PPRM 的思想是学习判断两个解方案中哪一个更优，而不是给单个解方案打绝对分数。

• 输入: 给定一个数学问题 $Ques$ 和两个候选解方案 $a_1,a_2$。
• 输出: $P(a_1\succ a_2|\varphi )$，即 $a_1$ 优于 $a_2$ 的概率，其中 $\succ$ 代表“更优”，$\varphi$ 是模型的参数。论文中使用了类似 Bradley-Terry 模型的形式：
  $$P(a_1\succ a_2|\varphi )=\frac{e^{\varphi (a_1)}}{e^{\varphi (a_1)}+e^{\varphi (a_2)}}$$
  其中 $\varphi (a)$ 可以理解为一个隐藏的“质量”函数。
• 训练:
  • 将此任务构建为一个问答（QA）任务。例如，PPRM被提问：“对于问题 $Ques$，解 $a_1$ 是否比解 $a_2$ 更好？” PPRM需要回答“是”或“否”。
  • 利用人类反馈强化学习 (RLHF) 的技术，和直接偏好优化 (Direct Preference Optimization, DPO) 方法进行训练。DPO 的优势在于它直接从偏好数据中学习，而无需显式地拟合一个奖励模型。
  • 训练数据来源于 PRM800K 和 OpenMathInstruct-1 等数学解题数据集，通过比较步骤质量或最终答案正确性来生成偏好对。

2.3 EBC (Enhanced Borda Count)

EBC 负责将PPRM产生的局部成对偏好聚合成全局的量化分数，用于指导MCTS。

1. 构建偏好矩阵 (Preference Matrix): 对所有 $n$ 个候选解，使用PPRM得到一个 $n\times n$ 的偏好矩阵 $M$。如果 $P(a_i\succ a_j)\ge 0.5$，则 $M_{ij}=1$，否则为 $0$。

2. 计算传递闭包 (Transitive Closure): 数学解的优劣通常具有传递性 (如果 $a_i\succ a_k$ 且 $a_k\succ a_j$，那么 $a_i\succ a_j$)。论文使用 Floyd-Warshall 算法计算偏好矩阵 $M$ 的传递闭包 $C$。这有助于处理潜在的噪声和不一致性，并强化偏好关系。

   Floyd-Warshall 算法是一种用于在加权有向图中查找所有节点对之间最短路径的动态规划算法。

   一文足矣——动态规划经典之Floyd（弗洛伊德）算法_floyd算法负权回路-CSDN博客

3. Borda计数与排序 (Borda Count and Ranking): 对于传递闭包矩阵 $C$，每个解 $a_i$ 的Borda分数定义为其在该矩阵中“胜过”的其他解的数量，即 ${\sum }_{j\ne i}{C}_{ij}$。根据Borda分数对所有解进行排序。

   Borda 计数是一种投票或排序聚合方法，它根据每个选项（或候选人，或在这里的解方案）在所有比较中“击败”其他选项的次数来给它们打分。得分越高的选项排名越靠前。

   例如，在一场循环赛中，每个选手（解方案）赢一场得1分，输了不得分。最后看每个选手的总得分来排名。

4. 全局量化分数 (Global Quantile Score): 将排序名次转换为全局分位数 $Q_g(v)$。例如，对于解 $v$，其分位数为 $Q_g(v)=1-\frac{\text{rank}(v)-1}{|V|-1}$，其中 $|V|$ 是解的总数，$\text{rank}(v)$ 是 $v$ 的排名。

5. 最终评估值: 一个解 $s^{\prime }$ 在MCTS中的最终评估值 $Q(s^{\prime })$ 是其全局量化分数 $Q_g(s^{\prime })$ 和基于其在搜索树中与邻近节点比较得出的局部优势分数 $Q_l(s^{\prime })$ 的加权组合：
   $$Q(s^{\prime })=\alpha \cdot Q_g(s^{\prime })+(1-\alpha )\cdot Q_l(s^{\prime })$$
   其中 $\alpha$ 是一个超参数，平衡全局和局部信息。

通过这三个组件的协同工作，LLaMA-Berry 能够在复杂的数学问题上进行更有效的搜索和评估。

3. 优势

相较于现有方法，LLaMA-Berry展现出以下主要优势：

1. 更强的探索能力：SR-MCTS通过MCTS的系统性探索，结合Self-Refine作为节点扩展方式，比传统的逐步生成或贪婪搜索方法更能跳出局部最优，探索更广阔的解空间。
2. 更鲁棒的奖励信号：PPRM通过学习成对偏好而非绝对分数，降低了对奖励模型精确标度（scaling）的依赖。判断“A比B好”通常比“A值多少分，B值多少分”更容易学习和泛化，尤其是在多样化的数学问题中。DPO的引入也使得训练更直接。
3. 更全面的解评估：EBC通过聚合局部偏好并计算传递闭包，生成了更全局和一致的解质量排序，有效缓解了成对比较可能带来的局部性和潜在冲突。
4. 对基础模型无额外训练：核心的推理LLM (如LLaMA-3.1-8B) 无需针对此框架进行额外的微调。PPRM是一个独立的、相对较小的模型 (Gemma2-2B)。
5. 针对复杂问题的潜力：整个框架设计天然适合处理需要深度推理和多轮优化的复杂问题，如奥林匹克数学竞赛题。

4. 实验

• 实验设置:
  • 基础LLM: Meta-LLaMA-3.1-8B-Instruct。
  • PPRM模型: Google Gemma2-2B-Instruct。
  • 基准数据集: GSM8K (基础数学), MATH (高中竞赛级), OlympiadBench (奥赛级), AIME24, AMC23 (美国数学竞赛), College Math, MMLU STEM, GPQA (研究生水平问答)。
  • 评价指标:
    • major@k: 在$k$个采样中，多数解正确的问题比例。
    • rm@k (本文中等同于 solved rate@k): 在$k$个采样中，至少有一个解正确的问题比例 (pass@k 或 best-of-N)。
    • 零样本CoT的贪婪性能 (greedy performance)。

• 主要结论:
  1. 显著性能提升: LLaMA-Berry在所有测试的基准上都显著提升了基础LLaMA-3.1-8B模型的性能。例如，在OlympiadBench和College Math上，16次推演 (@16) 的解决率提升超过35%。
  2. 小模型匹敌大模型: 尽管基础模型仅为8B参数，LLaMA-Berry在多个基准上（如OlympiadBench, College Math）的性能超越了参数量远大于己的Qwen2-Math-72B。在最具挑战性的AIME2024上，LLaMA-Berry (@16) 取得了26.7%的解决率 (8/30)，远超其他开源模型，并接近如GPT-4 Turbo这样的闭源模型。
  3. 消融实验:
     • SR-MCTS优于迭代式Self-Refine：MCTS的引入有效缓解了迭代优化中因错误批判导致的性能退化。
     • PPRM至关重要：与使用LLM自身评估作为奖励信号相比，PPRM显著提升了在GSM8K和AIME2024等任务上的性能，尤其是在复杂问题上，PPRM提供的更准确的引导信号作用巨大。
  4. 性能随推演次数扩展: 增加MCTS的推演次数 (rollouts) 通常能带来性能的持续提升，直到达到模型本身能力或搜索效率的瓶颈。