本论文作者来自华为蒙特利尔诺亚方舟实验室的康计堃,李信择,陈熙, Amirreza Kazemi,陈博兴。
人工智能(AI)在过去十年里取得了长足进步,特别是在自然语言处理和计算机视觉领域。然而,如何提升 AI 的认知能力和推理能力,仍然是一个巨大的挑战。
近期,一篇题为《MindStar: Enhancing Math Reasoning in Pre-trained LLMs at Inference Time》的论文提出了基于树搜索的推理时间能力提升方法 MindStar [1],该方法在开源模型 Llama-13-B 与 Mistral-7B 上达到了近似闭源大模型 GPT-3.5 与 Grok-1 在数学问题上的推理能力。
论文标题:MindStar: Enhancing Math Reasoning in Pre-trained LLMs at Inference Time
论文地址:https://arxiv.org/abs/2405.16265v2
MindStar 在数学问题上的应用效果:
图 1 :不同大型语言模型的数学准确率。LLaMA-2-13B 在数学性能上与 GPT-3.5 (4-shot) 类似,但节省了大约 200 倍的计算资源。
1. 引言
随着模型规模的快速增长,基于 Transformer 的大型语言模型(LLMs)在指令遵循 [1,2]、编码辅助 [3,4] 和创意写作 [5] 等领域展示了令人印象深刻的成果。然而,解锁 LLMs 解决复杂推理任务的能力仍然是一大挑战。最近的一些研究 [6,7] 尝试通过监督微调(Supervised Fine-Tuning, SFT)来解决,通过将新的推理数据样本与原始数据集混合,使 LLMs 学习这些样本的底层分布,并尝试模仿所学逻辑来解决未见过的推理任务。尽管这种方法有性能提升,但它严重依赖于大量的训练和额外的数据准备 [8,9]。
Llama-3 报告 [10] 强调了一个重要的观察:当面对一个具有挑战性的推理问题时,模型有时会生成正确的推理轨迹。这表明模型知道如何产生正确答案,但在选择上存在困难。基于这一发现,我们提出了一个简单的问题:我们能否通过帮助 LLMs 选择正确的输出来增强它们的推理能力?为探索这一点,我们进行了一项实验,利用不同的奖励模型进行 LLMs 输出选择。实验结果表明,步骤级选择显著优于传统的 CoT 方法。
2. MindStar 方法
图 2 MindStar 的算法架构图
我们引入了一种新的推理搜索框架 ——MindStar(M*),通过将推理任务视为搜索问题,并利用过程监督的奖励模型(Process-supervised Reward Model, PRM),M * 在推理树空间中有效导航,识别近似最优路径。结合束搜索(Beam Search, BS)和 Levin 树搜索(Levin Tree Search, LevinTS)的思想,进一步增强了搜索效率,并保证在有限计算复杂度内找到最佳推理路径。
2.1 过程监督奖励模型
过程监督奖励模型 (PRM) 的设计目的是评估大语言模型 (LLM) 生成的中间步骤,以帮助选择正确的推理路径。这种方法借鉴了其他应用中 PRM 的成功经验。具体而言,PRM 以当前推理路径和潜在的下一步作为输入,并返回奖励值。
PRM 通过考虑整个当前推理轨迹来评估新步骤,鼓励与整体路径的一致性和忠实性。高奖励值表明,新的步骤)对于给定的推理路径可能是正确的,从而使扩展路径值得进一步探索。相反,低奖励值则表示新步骤可能不正确,这意味着遵循此路径的解决方案也可能不正确。
M* 算法包含两个主要步骤,迭代直到找到正确的解决方案:
1. 推理路径扩展:在每次迭代中,基础 LLM 生成当前推理路径的下一步。
2. 评估和选择:使用 PRM 评估生成的步骤,并根据这些评估选择下一次迭代的推理路径。
2.2 推理路径扩展
在选择要扩展的推理路径后,我们设计了一个提示模板(Example 3.1),以从 LLM 中收集下一步。正如示例所示,LLM 将原始问题作为 {question},将当前推理路径作为 {answer}。注意,在算法的第一次迭代中,所选择的节点是仅包含问题的根节点,因此 {answer} 为空。对于推理路径,LLM 生成 N 个中间步骤,并将它们作为当前节点的子节点附加。在算法的下一步中,将评估这些新生成的子节点,并选择一个新的节点进行进一步扩展。我们还意识到,生成步骤的另一种方法是使用步骤标记对 LLM 进行微调。然而,这可能会降低 LLM 的推理能力,更重要的是,这与本文的重点 —— 在不修改权重的情况下增强 LLM 推理能力相悖。
2.3 推理路径选择
在扩展推理树后,我们使用预训练的过程监督奖励模型(PRM)来评估每个新生成的步骤。正如前面提到的,PRM 采用路径和步骤 ,并返回相应的奖励值。在评估之后,我们需要一种树搜索算法来选择下一个要扩展的节点。我们的框架不依赖于特定的搜索算法,在这项工作中,我们实例化了两种最佳优先搜索方法,即 Beam Search 和 Levin Tree Search。
3. 结果与讨论
在 GSM8K 和 MATH 数据集上的广泛评估显示,M * 显著提升了开源模型(如 LLaMA-2)的推理能力,其表现可与更大规模的闭源模型(如 GPT-3.5 和 Grok-1)媲美,同时大幅减少了模型规模和计算成本。这些发现突显了将计算资源从微调转移到推理时间搜索的潜力,为未来高效推理增强技术的研究开辟了新途径。
表 1 展示了各种方案在 GSM8K 和 MATH 推理基准上的对比结果。每个条目的数字表示问题解决的百分比。符号 SC@32 表示在 32 个候选结果中的自一致性,而 n-shot 表示少样本例子的结果。CoT-SC@16 指的是在 16 个思维链(CoT)候选结果中的自一致性。BS@16 代表束搜索方法,即在每个步骤级别涉及 16 个候选结果,而 LevinTS@16 详细说明了使用相同数量候选结果的 Levin 树搜索方法。值得注意的是,MATH 数据集上 GPT-4 的最新结果为 GPT-4-turbo-0409,我们特别强调这一点,因为它代表了 GPT-4 家族中的最佳性能。
图 3 我们研究了 M * 性能如何随着步骤级别候选数量的变化而变化。我们选择 Llama-2-13B 作为基础模型,并分别选择束搜索(BS)作为搜索算法。
图 4 Llama-2 和 Llama-3 模型家族在 MATH 数据集上的尺度定律。所有结果均来自它们的原始资源。我们使用 Scipy 工具和对数函数来计算拟合曲线。
表 2 不同方法在回答问题时的平均 token 生产数量
4. 结论
本文介绍了 MindStar(M*),一种新颖的基于搜索的推理框架,用于增强预训练大型语言模型的推理能力。通过将推理任务视为搜索问题并利用过程监督的奖励模型,M* 在推理树空间中有效导航,识别近似最优路径。结合束搜索和 Levin 树搜索的思想,进一步增强了搜索效率,并保证在有限计算复杂度内找到最佳推理路径。广泛的实验结果表明,M* 显著提升了开源模型的推理能力,其表现可与更大规模的闭源模型媲美,同时大幅减少了模型规模和计算成本。
这些研究成果表明,将计算资源从微调转移到推理时间搜索具有巨大的潜力,为未来高效推理增强技术的研究开辟了新途径。
还没有评论,来说两句吧...