本論文作者來自華為蒙特利爾諾亞方舟實(shí)驗(yàn)室的康計(jì)堃，李信擇，陳熙， Amirreza Kazemi，陳博興。

人工智能（AI）在過去十年里取得了長(zhǎng)足進(jìn)步，特別是在自然語(yǔ)言處理和計(jì)算機(jī)視覺領(lǐng)域。然而，如何提升 AI 的認(rèn)知能力和推理能力，仍然是一個(gè)巨大的挑戰(zhàn)。

近期，一篇題為《MindStar: Enhancing Math Reasoning in Pre-trained LLMs at Inference Time》的論文提出了基于樹搜索的推理時(shí)間能力提升方法 MindStar [1]，該方法在開源模型 Llama-13-B 與 Mistral-7B 上達(dá)到了近似閉源大模型 GPT-3.5 與 Grok-1 在數(shù)學(xué)問題上的推理能力。

• 論文標(biāo)題：MindStar: Enhancing Math Reasoning in Pre-trained LLMs at Inference Time
• 論文地址：https://arxiv.org/abs/2405.16265v2

MindStar 在數(shù)學(xué)問題上的應(yīng)用效果：

圖 1 ：不同大型語(yǔ)言模型的數(shù)學(xué)準(zhǔn)確率。LLaMA-2-13B 在數(shù)學(xué)性能上與 GPT-3.5 (4-shot) 類似，但節(jié)省了大約 200 倍的計(jì)算資源。

1. 引言

隨著模型規(guī)模的快速增長(zhǎng)，基于 Transformer 的大型語(yǔ)言模型（LLMs）在指令遵循 [1,2]、編碼輔助 [3,4] 和創(chuàng)意寫作 [5] 等領(lǐng)域展示了令人印象深刻的成果。然而，解鎖 LLMs 解決復(fù)雜推理任務(wù)的能力仍然是一大挑戰(zhàn)。最近的一些研究 [6,7] 嘗試通過監(jiān)督微調(diào)（Supervised Fine-Tuning, SFT）來解決，通過將新的推理數(shù)據(jù)樣本與原始數(shù)據(jù)集混合，使 LLMs 學(xué)習(xí)這些樣本的底層分布，并嘗試模仿所學(xué)邏輯來解決未見過的推理任務(wù)。盡管這種方法有性能提升，但它嚴(yán)重依賴于大量的訓(xùn)練和額外的數(shù)據(jù)準(zhǔn)備 [8,9]。

Llama-3 報(bào)告 [10] 強(qiáng)調(diào)了一個(gè)重要的觀察：當(dāng)面對(duì)一個(gè)具有挑戰(zhàn)性的推理問題時(shí)，模型有時(shí)會(huì)生成正確的推理軌跡。這表明模型知道如何產(chǎn)生正確答案，但在選擇上存在困難?；谶@一發(fā)現(xiàn)，我們提出了一個(gè)簡(jiǎn)單的問題：我們能否通過幫助 LLMs 選擇正確的輸出來增強(qiáng)它們的推理能力？為探索這一點(diǎn)，我們進(jìn)行了一項(xiàng)實(shí)驗(yàn)，利用不同的獎(jiǎng)勵(lì)模型進(jìn)行 LLMs 輸出選擇。實(shí)驗(yàn)結(jié)果表明，步驟級(jí)選擇顯著優(yōu)于傳統(tǒng)的 CoT 方法。

2. MindStar 方法

圖 2 MindStar 的算法架構(gòu)圖

我們引入了一種新的推理搜索框架 ——MindStar（M*），通過將推理任務(wù)視為搜索問題，并利用過程監(jiān)督的獎(jiǎng)勵(lì)模型（Process-supervised Reward Model, PRM），M * 在推理樹空間中有效導(dǎo)航，識(shí)別近似最優(yōu)路徑。結(jié)合束搜索（Beam Search, BS）和 Levin 樹搜索（Levin Tree Search, LevinTS）的思想，進(jìn)一步增強(qiáng)了搜索效率，并保證在有限計(jì)算復(fù)雜度內(nèi)找到最佳推理路徑。

2.1 過程監(jiān)督獎(jiǎng)勵(lì)模型

過程監(jiān)督獎(jiǎng)勵(lì)模型 (PRM) 的設(shè)計(jì)目的是評(píng)估大語(yǔ)言模型 (LLM) 生成的中間步驟，以幫助選擇正確的推理路徑。這種方法借鑒了其他應(yīng)用中 PRM 的成功經(jīng)驗(yàn)。具體而言，PRM 以當(dāng)前推理路徑和潛在的下一步作為輸入，并返回獎(jiǎng)勵(lì)值。

PRM 通過考慮整個(gè)當(dāng)前推理軌跡來評(píng)估新步驟，鼓勵(lì)與整體路徑的一致性和忠實(shí)性。高獎(jiǎng)勵(lì)值表明，新的步驟）對(duì)于給定的推理路徑可能是正確的，從而使擴(kuò)展路徑值得進(jìn)一步探索。相反，低獎(jiǎng)勵(lì)值則表示新步驟可能不正確，這意味著遵循此路徑的解決方案也可能不正確。

M* 算法包含兩個(gè)主要步驟，迭代直到找到正確的解決方案：

1. 推理路徑擴(kuò)展：在每次迭代中，基礎(chǔ) LLM 生成當(dāng)前推理路徑的下一步。

2. 評(píng)估和選擇：使用 PRM 評(píng)估生成的步驟，并根據(jù)這些評(píng)估選擇下一次迭代的推理路徑。

2.2 推理路徑擴(kuò)展

在選擇要擴(kuò)展的推理路徑后，我們?cè)O(shè)計(jì)了一個(gè)提示模板（Example 3.1），以從 LLM 中收集下一步。正如示例所示，LLM 將原始問題作為 {question}，將當(dāng)前推理路徑作為 {answer}。注意，在算法的第一次迭代中，所選擇的節(jié)點(diǎn)是僅包含問題的根節(jié)點(diǎn)，因此 {answer} 為空。對(duì)于推理路徑，LLM 生成 N 個(gè)中間步驟，并將它們作為當(dāng)前節(jié)點(diǎn)的子節(jié)點(diǎn)附加。在算法的下一步中，將評(píng)估這些新生成的子節(jié)點(diǎn)，并選擇一個(gè)新的節(jié)點(diǎn)進(jìn)行進(jìn)一步擴(kuò)展。我們還意識(shí)到，生成步驟的另一種方法是使用步驟標(biāo)記對(duì) LLM 進(jìn)行微調(diào)。然而，這可能會(huì)降低 LLM 的推理能力，更重要的是，這與本文的重點(diǎn) —— 在不修改權(quán)重的情況下增強(qiáng) LLM 推理能力相悖。

2.3 推理路徑選擇

在擴(kuò)展推理樹后，我們使用預(yù)訓(xùn)練的過程監(jiān)督獎(jiǎng)勵(lì)模型（PRM）來評(píng)估每個(gè)新生成的步驟。正如前面提到的，PRM 采用路徑和步驟 ，并返回相應(yīng)的獎(jiǎng)勵(lì)值。在評(píng)估之后，我們需要一種樹搜索算法來選擇下一個(gè)要擴(kuò)展的節(jié)點(diǎn)。我們的框架不依賴于特定的搜索算法，在這項(xiàng)工作中，我們實(shí)例化了兩種最佳優(yōu)先搜索方法，即 Beam Search 和 Levin Tree Search。

3. 結(jié)果與討論

在 GSM8K 和 MATH 數(shù)據(jù)集上的廣泛評(píng)估顯示，M * 顯著提升了開源模型（如 LLaMA-2）的推理能力，其表現(xiàn)可與更大規(guī)模的閉源模型（如 GPT-3.5 和 Grok-1）媲美，同時(shí)大幅減少了模型規(guī)模和計(jì)算成本。這些發(fā)現(xiàn)突顯了將計(jì)算資源從微調(diào)轉(zhuǎn)移到推理時(shí)間搜索的潛力，為未來高效推理增強(qiáng)技術(shù)的研究開辟了新途徑。

表 1 展示了各種方案在 GSM8K 和 MATH 推理基準(zhǔn)上的對(duì)比結(jié)果。每個(gè)條目的數(shù)字表示問題解決的百分比。符號(hào) SC@32 表示在 32 個(gè)候選結(jié)果中的自一致性，而 n-shot 表示少樣本例子的結(jié)果。CoT-SC@16 指的是在 16 個(gè)思維鏈（CoT）候選結(jié)果中的自一致性。BS@16 代表束搜索方法，即在每個(gè)步驟級(jí)別涉及 16 個(gè)候選結(jié)果，而 LevinTS@16 詳細(xì)說明了使用相同數(shù)量候選結(jié)果的 Levin 樹搜索方法。值得注意的是，MATH 數(shù)據(jù)集上 GPT-4 的最新結(jié)果為 GPT-4-turbo-0409，我們特別強(qiáng)調(diào)這一點(diǎn)，因?yàn)樗砹?GPT-4 家族中的最佳性能。

圖 3 我們研究了 M * 性能如何隨著步驟級(jí)別候選數(shù)量的變化而變化。我們選擇 Llama-2-13B 作為基礎(chǔ)模型，并分別選擇束搜索（BS）作為搜索算法。

圖 4 Llama-2 和 Llama-3 模型家族在 MATH 數(shù)據(jù)集上的尺度定律。所有結(jié)果均來自它們的原始資源。我們使用 Scipy 工具和對(duì)數(shù)函數(shù)來計(jì)算擬合曲線。

表 2 不同方法在回答問題時(shí)的平均 token 生產(chǎn)數(shù)量

4. 結(jié)論

本文介紹了 MindStar（M*），一種新穎的基于搜索的推理框架，用于增強(qiáng)預(yù)訓(xùn)練大型語(yǔ)言模型的推理能力。通過將推理任務(wù)視為搜索問題并利用過程監(jiān)督的獎(jiǎng)勵(lì)模型，M* 在推理樹空間中有效導(dǎo)航，識(shí)別近似最優(yōu)路徑。結(jié)合束搜索和 Levin 樹搜索的思想，進(jìn)一步增強(qiáng)了搜索效率，并保證在有限計(jì)算復(fù)雜度內(nèi)找到最佳推理路徑。廣泛的實(shí)驗(yàn)結(jié)果表明，M* 顯著提升了開源模型的推理能力，其表現(xiàn)可與更大規(guī)模的閉源模型媲美，同時(shí)大幅減少了模型規(guī)模和計(jì)算成本。

這些研究成果表明，將計(jì)算資源從微調(diào)轉(zhuǎn)移到推理時(shí)間搜索具有巨大的潛力，為未來高效推理增強(qiáng)技術(shù)的研究開辟了新途徑。