多模态慢思考：分解原子步骤以解决复杂数学推理

多模态慢思考：分解原子步骤以解决复杂数学推理

在人工智能领域，高阶数学推理一直是一个极具挑战性的任务。最近，来自中山大学、香港科技大学、上海交通大学和华为诺亚方舟实验室的研究团队提出了一种名为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数据集以及一个关注于提升原子步骤质量的慢思考框架。该方法在解决问题的每一步都始终保持着较高的推理质量，从而提高了各种多模态数学任务的推理性能，为开发广义慢思考模型铺平了道路。

论文链接：https://arxiv.org/abs/2411.11930
项目主页（即将开源）：https://github.com/Quinn777/AtomThink