本文作者來自中山大學、香港科技大學、上海交通大學和華為諾亞方舟實驗室等機構。第一作者項鯤為中山大學博士生，劉智立為香港科技大學博士生，姜子昊為上海交通大學碩士研究生。

AtomThink 是一個包括 CoT 注釋引擎、原子步驟指令微調、政策搜索推理的全流程框架，旨在通過將 “慢思考 “能力融入多模態(tài)大語言模型來解決高階數(shù)學推理問題。量化結果顯示其在兩個基準數(shù)學測試中取得了大幅的性能增長，并能夠輕易遷移至不同的多模態(tài)大模型當中。

• 論文：https://arxiv.org/abs/2411.11930
• 主頁（即將開源）：https://github.com/Quinn777/AtomThink

背景與挑戰(zhàn)

在人工智能研究領域，高階數(shù)學推理一直是一項極具挑戰(zhàn)的任務。之前的一些工作通過精心設計 prompt 來激發(fā)模型生成思維鏈（CoT），而最近大熱的 OpenAI o1 的推出更是標志著強人工智能的巨大飛躍，它擅長通過利用擴展推理鏈和放大測試時間來解決復雜問題，即 “慢思考”。

然而，在視覺數(shù)學任務中由于信息建模的數(shù)據(jù)和計算資源需求顯著增加，將慢思考技術應用于多模態(tài)大語言模型（MLLM）具有挑戰(zhàn)性。盡管之前的一些工作說明激發(fā)模型固有 CoT 能力能夠改善推理，但他們并未考慮推理鏈中間步驟的質量，也缺乏對多模態(tài) CoT 每個節(jié)點對應能力的細粒度分析，因此很難應用測試時間縮放策略來進行針對性的性能提升。

原子步驟質量評估

因此，本研究率先提出了一種原子步驟質量評估策略（如上圖所示），為慢思考能力提供新的分析視角。原子步驟是語義維度的最小推理步驟?？紤]到人類可能利用不同的認知能力來解決數(shù)學問題，首先從 GPT-4o 的推理行為分布中構建一個規(guī)范的推理能力集合，代表高級智能模型在進行思考時所使用的不同能力。最后，將候選模型輸出應用結果監(jiān)督和重映射來估計其對應能力的綜合得分。

下圖所示的結果表明當前開源模型的平均原子步驟質量較差，特別是在圖像識別、變量定義和計算等能力項上存在明顯缺陷。這一發(fā)現(xiàn)進一步促使我們關注現(xiàn)有模型執(zhí)行單步推理的能力，并通過提高原子推理步驟的質量來引入更強大的慢思考。

AtomThink 慢思考框架

因此，為了解決開源 MLLM 原子步驟質量較差的問題，本文提出了 AtomThink 慢思考框架，通過引入多模態(tài) CoT 注釋引擎、原子步驟微調策略和政策搜索策略，旨在通過仔細訓練來增強 MLLM 的解碼能力，并結合采樣后搜索來識別最佳預測節(jié)點，以逐步生成一條高質量的推理路徑。

1. 多模態(tài)注釋引擎

首先文中從數(shù)據(jù)角度引入了動態(tài)提示和短 CoT 增強策略來構建數(shù)據(jù)引擎。動態(tài)提示策略促使已有 LLM 迭代地構建狀態(tài)推理路徑。每個路徑節(jié)點代表一個推理步驟，包括前一階段、當前狀態(tài)和可能的動作。可能的行動包括繼續(xù)推理、驗證和得出結論，這由 LLM 自身決定。為充分利用現(xiàn)有 VQA 數(shù)據(jù)集的短 CoT 注釋和答案，本文還使用 GPT-4o 來執(zhí)行原子化分割和注釋增強，這允許我們通過簡單的提示來將原始推理過程從語義上劃分為多個離散的步驟。

通過上述方式，本文從公開數(shù)據(jù)集中采樣數(shù)學問題并生成長思維鏈，通過人工篩查和后處理去除不正確的節(jié)點，最終制作了 AtomMATH 數(shù)據(jù)集。它包括 AMATH-SFT 和 AMATH-PRM 兩個子集，分別用于監(jiān)督指令微調和過程監(jiān)督訓練。

上表 2 中顯示了該數(shù)據(jù)集的質量較高，甚至在 GPT-4o 的評分中超越了人工標注的 PRM800k 數(shù)據(jù)集。

2. 原子步驟微調

該部分包括對現(xiàn)有 MLLM 的指令微調和對 PRM 的對齊訓練。首先通過將輸入數(shù)據(jù)重構為獨立的歷史狀態(tài)和當前動作，讓多模態(tài)大模型學習近似馬爾可夫決策的輸出格式。而在 PRM 的訓練則基于已有 LLM 來執(zhí)行后訓練對齊，通過最小化下面的交叉熵損失函數(shù)來進行優(yōu)化：

3. 策略搜索

由于擴大測試時間會產(chǎn)生多個候選步驟，文中將現(xiàn)有的策略分為路徑維度搜索和步驟維度搜索。路徑維度搜索包括：

1）多數(shù)投票：通過選擇多個推理路徑中最常見的結果來選擇最優(yōu)。

2）Best-of-N：計算每個候選路徑的所有節(jié)點得分，通過不同聚合策略來將密集分數(shù)映射到整條路徑。聚合策略包括最壞動作、最后動作和平均得分。

步驟維度搜索包括：

1）貪心算法：側重于在推理過程的每一步做出局部最優(yōu)選擇。它根據(jù)當前狀態(tài)選擇最佳的即時行動（步驟），而不考慮未來的后果。

2）Beam Search：在每個推理動作中探索多個分支，并為每個推理階段維護固定數(shù)量的較優(yōu)候選原子步驟。它在探索不同的推理鏈和當前最優(yōu)步驟之間取得了平衡。

實驗結果

實驗測試了四種不同的推理范式，包括：

1）直接輸出：直接輸出答案。

2）CoT 輸出：設計 CoT 提示來一次性產(chǎn)生中間推理步驟和答案。

3）QuickThink：逐步產(chǎn)生一條原子推理路徑而不進行搜索，推理時間較短。

4）SlowThink：使用 Beam Search 來搜索和剪枝一棵搜索樹，推理時間較長。

上表展示了本文框架的性能。首先在只使用 QuickThink 時，對于 LLaVA-Llama3-8B 其 AtomThink 版本就大大超越了基線模型，在 MathVista 上實現(xiàn)了約 10% 的改進，在 MathVerse 上實現(xiàn)了 19% 的提升。這表明，當模型具有原子推理能力時，它已經(jīng)可以利用快速思考進行更準確的數(shù)學推理。

實驗結果還證明了純語言模型也能對多模態(tài)推理提供有效的過程監(jiān)督。在和 LLaVA-Llama3-8B 與 EMOVA-8B 的對比中 SlowThink 都獲得了一致的性能提升，特別是對比 CoT 輸出的性能甚至提升了一倍，表明該框架具有強大的可擴展性。

對于不同的搜索策略本文也做出了細致比較，實驗證明 Best-of-N 和 Beam Search 比多數(shù)投票和貪心算法的提升明顯，特別是基于平均得分聚合的 Best-of-N 在 MathVista-Math 任務中獲得了最高的 58.7% 的精度。

通過改變候選步驟數(shù)量，文章還探討了多模態(tài)數(shù)學推理任務中 Test-time scaling law 的存在。隨著候選步驟的增加，兩個模型都呈現(xiàn)出不斷增強的趨勢，其中較弱的 LLaVA 模型即使推理時間增加十倍也沒有呈現(xiàn)出邊際效應。

總結

本文將原子思維能力引入多模態(tài)大語言模型，以解決困難的數(shù)學推理問題。文中發(fā)布了一個高質量的長 CoT 數(shù)據(jù)集以及一個關注于提升原子步驟質量的慢思考框架。該方法在解決問題的每一步都始終保持著較高的推理質量，從而提高了各種多模態(tài)數(shù)學任務的推理性能，為開發(fā)廣義慢思考模型鋪平了道路。