AIxiv专栏是AI在线发布学术、技术内容的栏目。过去数年，AI在线AIxiv专栏接收报道了2000多篇内容，覆盖全球各大高校与企业的顶级实验室，有效促进了学术交流与传播。如果您有优秀的工作想要分享，欢迎投稿或者联系报道。投稿邮箱：[email protected]；[email protected]
本文作者来自中山大学、香港科技大学、上海交通大学和华为诺亚方舟实验室等机构。第一作者项鲲为中山大学博士生，刘智立为香港科技大学博士生，姜子昊为上海交通大学硕士研究生。
AtomThink 是一个包括 CoT 注释引擎、原子步骤指令微调、政策搜索推理的全流程框架，旨在通过将 “慢思考 “能力融入多模态大语言模型来解决高阶数学推理问题。量化结果显示其在两个基准数学测试中取得了大幅的性能增长，并能够轻易迁移至不同的多模态大模型当中。
论文：https://arxiv.org/abs/2411.11930
主页（即将开源）：https://github.com/Quinn777/AtomThink
背景与挑战
在人工智能研究领域，高阶数学推理一直是一项极具挑战的任务。之前的一些工作通过精心设计 prompt 来激发模型生成思维链（CoT），而最近大热的 OpenAI o1 的推出更是标志着强人工智能的巨大飞跃，它擅长通过利用扩展推理链和放大测试时间来解决复杂问题，即 “慢思考”。
然而，在视觉数学任务中由于信息建模的数据和计算资源需求显著增加，将慢思考技术应用于多模态大语言模型（MLLM）具有挑战性。尽管之前的一些工作说明激发模型固有 CoT 能力能够改善推理，但他们并未考虑推理链中间步骤的质量，也缺乏对多模态 CoT 每个节点对应能力的细粒度分析，因此很难应用测试时间缩放策略来进行针对性的性能提升。
原子步骤质量评估
因此，本研究率先提出了一种原子步骤质量评估策略（如上图所示），为慢思考能力提供新的分析视角。原子步骤是语义维度的最小推理步骤。考虑到人类可能利用不同的认知能力来解决数学问题，首先从 GPT-4o 的推理行为分布中构建一个规范的推理能力集合，代表高级智能模型在进行思考时所使用的不同能力。最后，将候选模型输出应用结果监督和重映射来估计其对应能力的综合得分。
下图所示的结果表明当前开源模型的平均原子步骤质量较差，特别是在图像识别、变量定义和计算等能力项上存在明显缺陷。这一发现进一步促使我们关注现有模型执行单步推理的能力，并通过提高原子推理步骤的质量来引入更强大的慢思考。
AtomThink 慢思考框架
因此，为了解决开源 MLLM 原子步骤质量较差的问题，本文提出了 AtomThink 慢思考框架，通过引入多模态 CoT 注释引擎、原子步骤微调策略和政策搜索策略，旨在通过仔细训练来增强 MLLM 的解码能力，并结合采样后搜索来识别最佳预测节点，以逐步生成一条高质量的推理路径。
1. 多模态注释引擎
首先文中从数据角度引入了动态提示和短 CoT 增强策略来构建数据引擎。动态提示策略促使已有 LLM 迭代地构建状态推理路径。每个路径节点代表一个推理步骤，包括前一阶段、当前状态和可能的动作。可能的行动包括继续推理、验证和得出结论，这由 LLM 自身决定。为充分利用现有 VQA 数据集的短 CoT 注释和答案，本文还使用 GPT-4o 来执行原子化分割和注释增强，这允许我们通过简单的提示来将原始推理过程从语义上划分为多个离散的步骤。
通过上述方式，本文从公开数据集中采样数学问题并生成长思维链，通过人工筛查和后处理去除不正确的节点，最终制作了 AtomMATH 数据集。它包括 AMATH-SFT 和 AMATH-PRM 两个子集，分别用于监督指令微调和过程监督训练。
上表 2 中显示了该数据集的质量较高，甚至在 GPT-4o 的评分中超越了人工标注的 PRM800k 数据集。
2. 原子步骤微调
该部分包括对现有 MLLM 的指令微调和对 PRM 的对齐训练。首先通过将输入数据重构为独立的历史状态和当前动作，让多模态大模型学习近似马尔可夫决策的输出格式。而在 PRM 的训练则基于已有 LLM 来执行后训练对齐，通过最小化下面的交叉熵损失函数来进行优化：
3. 策略搜索
由于扩大测试时间会产生多个候选步骤，文中将现有的策略分为路径维度搜索和步骤维度搜索。路径维度搜索包括：
1）多数投票：通过选择多个推理路径中最常见的结果来选择最优。
2）Best-of-N：计算每个候选路径的所有节点得分，通过不同聚合策略来将密集分数映射到整条路径。聚合策略包括最坏动作、最后动作和平均得分。
步骤维度搜索包括：
1）贪心算法：侧重于在推理过程的每一步做出局部最优选择。它根据当前状态选择最佳的即时行动（步骤），而不考虑未来的后果。
2）Beam Search：在每个推理动作中探索多个分支，并为每个推理阶段维护固定数量的较优候选原子步骤。它在探索不同的推理链和当前最优步骤之间取得了平衡。
实验结果
实验测试了四种不同的推理范式，包括：
1）直接输出：直接输出答案。
2）CoT 输出：设计 CoT 提示来一次性产生中间推理步骤和答案。
3）QuickThink：逐步产生一条原子推理路径而不进行搜索，推理时间较短。
4）SlowThink：使用 Beam Search 来搜索和剪枝一棵搜索树，推理时间较长。
上表展示了本文框架的性能。首先在只使用 QuickThink 时，对于 LLaVA-Llama3-8B 其 AtomThink 版本就大大超越了基线模型，在 MathVista 上实现了约 10% 的改进，在 MathVerse 上实现了 19% 的提升。这表明，当模型具有原子推理能力时，它已经可以利用快速思考进行更准确的数学推理。
实验结果还证明了纯语言模型也能对多模态推理提供有效的过程监督。在和 LLaVA-Llama3-8B 与 EMOVA-8B 的对比中 SlowThink 都获得了一致的性能提升，特别是对比 CoT 输出的性能甚至提升了一倍，表明该框架具有强大的可扩展性。
对于不同的搜索策略本文也做出了细致比较，实验证明 Best-of-N 和 Beam Search 比多数投票和贪心算法的提升明显，特别是基于平均得分聚合的 Best-of-N 在 MathVista-Math 任务中获得了最高的 58.7% 的精度。
通过改变候选步骤数量，文章还探讨了多模态数学推理任务中 Test-time scaling law 的存在。随着候选步骤的增加，两个模型都呈现出不断增强的趋势，其中较弱的 LLaVA 模型即使推理时间增加十倍也没有呈现出边际效应。
总结
本文将原子思维能力引入多模态大语言模型，以解决困难的数学推理问题。文中发布了一个高质量的长 CoT 数据集以及一个关注于提升原子步骤质量的慢思考框架。该方法在解决问题的每一步都始终保持着较高的推理质量，从而提高了各种多模态数学任务的推理性能，为开发广义慢思考模型铺平了道路。
参考文献
[1] Lightman, Hunter, et al. "Let's verify step by step." arXiv preprint arXiv:2305.20050 (2023).
[2] Snell, Charlie, et al. "Scaling llm test-time compute optimally can be more effective than scaling model parameters." arXiv preprint arXiv:2408.03314 (2024).
[3] Wei, Jason, et al. "Chain-of-thought prompting elicits reasoning in large language models." Advances in neural information processing systems 35 (2022): 24824-24837.
[4] Wang, Peiyi, et al. "Math-shepherd: Verify and reinforce llms step-by-step without human annotations." Proceedings of the 62nd Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers). 2024.
[5] Wang, Jun, et al. "OpenR: An Open Source Framework for Advanced Reasoning with Large Language Models." arXiv preprint arXiv:2410.09671 (2024).