大模型+蒙特卡洛树搜索，一招让LLaMa-3 8B奥数水平直逼GPT-4
大模型+蒙特卡洛树搜索，一招让LLaMa-3 8B奥数水平直逼GPT-4
最近，17岁中专生姜萍在2024阿里巴巴全球数学竞赛预选赛中取得全球第12名的新闻刷屏。与此同时，AI在数学竞赛中的表现也引人注目：在所有563支AI参赛队伍中，最高分达到34分，平均分18分，已接近人类选手平均水平。然而，AI在逻辑推理能力方面仍存在短板，尤其是在需要策略和逻辑推理的任务中。为应对这一挑战，复旦大学和上海AI实验室的研究者提出了MCT Self-Refine（MCTSr）算法，通过将大语言模型（LLM）与蒙特卡洛树搜索（MCTS）算法相结合，显著提升了LLM在复杂数学推理任务中的表现。
MCTSr算法框架
作为一种决策工具，MCTS广泛应用于人工智能中需要战略规划的场景，通常用于游戏和复杂的问题解决环境。本文通过将MCTS的系统探索能力与LLM的Self-Refine和Self-Evaluation能力相结合，旨在创建一个更强大的框架来应对当前LLM难以解决的复杂推理任务。
在将MCTS与LLM集成过程中存在一些技术挑战。传统的MCTS策略可能与LLM输出的随机性和生成性不太吻合，后者通常涉及无限、连续的潜在动作空间。这种不一致需要在MCTS框架内采用定制的期望计算和反向传播方法，以更好地适应LLM的特有属性。
此外，研究者还引入了一种动态剪枝策略，它结合了改进的置信上限（UCB）公式，以优化高风险任务中有效决策制定所需要的探索-利用平衡。
MCTSr工作流程
MCTSr工作流包括以下几个阶段：
初始化：使用模型生成的答案和虚拟响应建立根节点，以最大限度地减少模型过度拟合趋势。
选择：该算法采用值函数Q对所有未完全展开的答案进行排序，并采用贪心策略选择值最高的节点进行进一步的探索和优化。
Self-Refine：选择好的答案a使用Self-Refine框架进行优化。最初，模型生成反馈m，指导优化过程以产生增强的答案a'。
Self-Evaluation：精炼后的答案经过评分从而采样一个奖励值，并计算其Q值。这涉及模型自我奖励反馈和约束，如严格的评分标准和抑制满分，以确保评分的可靠性和公平性。
反向传播：将精炼答案的值反向传播到其父节点和其他相关节点，以更新树的值信息。如果任何子节点的Q值发生变化，则更新父节点的Q。
UCT更新：在所有节点的Q值更新完成后，确定一个候选节点集合C，用于进一步扩展或选择，然后使用UCT更新公式更新所有节点的UCT值，以备下一步的选择阶段。
迭代上述阶段，直到满足终止条件T为止。
实验结果
为了评估MCTSr算法在解决数学问题中的有效性，研究者将LLaMA3-8B作为基础模型，并使用MCTSr进行增强。他们在Zero-Shot CoT、Self-Refine、4-rollouts MCTSr和8-rollouts MCTSr等几种设置中，将LLaMA3-8B与GPT-4、Claude 3和Gemini 1.5-Pro等进行了比较。
研究者在GSM8K和GSM-hard测试集（它们分别包含了典型和具有挑战性的数学问题）上评估了上述方法，结果如下表1所示。
可以发现，MCTSr的rollout次数与成功率之间存在着直接相关性，并随着迭代次数增加而显著提升，在不太复杂的GSM8K中尤为明显。不过对于更复杂的GSM-Hard测试集，即使rollout次数更高也会达到性能上限，表明当前策略在解决复杂问题时存在局限性。
这些结果强调了MCT-Self-refine算法的稳健性和潜在边界，以及持续改进的必要性，从而有效应对更复杂的挑战。
下表2展示了在MATH数据集上应用不同复杂度级别的MCT-Self-refine算法的结果。数据集分为五个难度级别，从Level 1（最简单）到Level 5（最具挑战性）。
结果显示，Level 1的成功率最高，8次rollout后，MCTSr实现了90.16%的成功率，解决了437个问题中的394个。随着rollout次数的增加，这一级别的成功率显著提高。
在最具挑战性的Level 5难度，8次rollout后，MCTSr的成功率为34.06%，解决了1324个问题中的451个。这说明了随着难度不断增加，该算法在高度复杂的场景中性能受到限制。
所有级别的整体性能显示，8次rollout后，MCTSr的累计成功率为58.24%，解决了5000个问题中的2912个。这一成功率相较于Zero-Shot CoT的初始成功率24.36%有了显著提高。这表明了，rollout次数的增加与成功率的提高呈现出一致性，强调了MCT-Self-refine算法在提升不同数学复杂度级别的问题解决能力方面的有效性。
这些结果还验证了MCT-Self-refine算法在学术和问题解决上下文中的潜力，并强调了其对MATH数据集中不同复杂度级别问题的可扩展性和适应性。
下表3为MCT-Self-refne算法在奥数竞赛的三个数据集上进行了测试：AlME、GAIC Math Odyssey和OlympiadBench。
- AIME：从Zero-Shot CoT的2.36%（解决22个问题）到MCTSr的11.79%（解决110个问题）。
- GAIC Math Odyssey：成功率从17.22%（解决67个问题）上升至49.36%（解决192个问题）。
- OlympiadBench：从Zero-Shot CoT的1.25%（解决16个问题）提高到MCTSr的7.76%（解决99个问题）。
这些结果证实了MCT-Self-refine算法在未见过的数学问题上的适用性，表明其在奥林匹克等竞争性学术环境中具有优势。
如表4所示。与当前闭源大模型进行比较时，MCTSr可以有效提升小参数开源模型（如LLaMa-3）的数学推理能力到相当的水平。
更多技术细节和实验结果请参阅原论文。