這幾天，17 歲中專生姜萍在 2024 阿里巴巴全球數(shù)學競賽預選賽中取得全球第 12 名的新聞刷了屏。而同時，AI 挑戰(zhàn)賽的成績顯示，在所有 563 支 AI 參賽隊伍中，最高分 34 分，平均分 18 分，趕上了人類選手平均水平。

AI 參與數(shù)學競賽的主要短板是邏輯推理能力弱，證明題很難拿到完整得分點。這也是 GPT-4、LLaMA 等當前大語言模型（LLM）在需要策略和邏輯推理的任務中面臨的重大挑戰(zhàn)。

其中的一大障礙是輸出的準確性和可信度，尤其是在需要保證精度的數(shù)學上下文中，LLM 在推理時往往容易產(chǎn)生幻覺。輸出結果表面上看似合理，但實際上不相關或事實不正確，最終導致不合理的推理過程。

雖然像 Self-Refine 這樣的重寫技術有助于緩解這種傾向，但依然可能導致現(xiàn)實世界復雜的數(shù)學問題產(chǎn)生誤導性或錯誤的結果。

因此，為了應對這些挑戰(zhàn)，來自復旦大學、上海 AI Lab 的研究者提出了 MCT Self-Refine（MCTSr），將 LLM 與蒙特卡洛樹搜索（MCTS）算法相結合，并重點提高 LLM 在復雜數(shù)學推理任務（比如奧數(shù)競賽題）中的表現(xiàn)。

作為一種決策工具，MCTS 廣泛應用于人工智能中需要戰(zhàn)略規(guī)劃的場景，通常用于游戲和復雜的問題解決環(huán)境。本文通過將 MCTS 的系統(tǒng)探索能力與 LLM 的 Self-Refine 和 Self-Evaluation 能力相結合， 旨在創(chuàng)建一個更強大的框架來應對當前 LLM 難以解決的復雜推理任務。

• 論文地址：https://arxiv.org/pdf/2406.07394
• 項目地址：https://github.com/trotsky1997/MathBlackBox

不過，在將 MCTS 與 LLM 集成過程中存在一些技術挑戰(zhàn)。傳統(tǒng)的 MCTS 策略可能與 LLM 輸出的隨機性和生成性不太吻合，后者通常涉及無限、連續(xù)的潛在動作空間。這種不一致需要在 MCTS 框架內(nèi)采用定制的期望計算和反向傳播方法，以更好地適應 LLM 的特有屬性。

此外，研究者還引入了一種動態(tài)剪枝策略，它結合了改進的置信上限（UCB）公式，以優(yōu)化高風險任務中有效決策制定所需要的探索 - 利用平衡。 

可以說，這項研究推進了 LLM 在復雜推理挑戰(zhàn)中的應用，為未來整合 AI 相關的技術創(chuàng)新奠定了基礎，從而使得 LLM 驅動的應用擁有了更強大的決策制定、推理準確性和可靠性。

方法概覽

MCTSr 架構圖如圖 1 所示：

MCTSr 工作流包括：

• 初始化：使用模型生成的答案和虛擬響應建立根節(jié)點，以最大限度地減少模型過度擬合趨勢；
• 選擇：該算法采用值函數(shù) Q 對所有未完全展開的答案進行排序，并采用貪心策略選擇值最高的節(jié)點進行進一步的探索和優(yōu)化；
• Self-Refine ：選擇好的答案 a 使用 Self-Refine 框架進行優(yōu)化。最初，模型生成反饋 m，指導優(yōu)化過程以產(chǎn)生增強的答案 a ′；
• Self-Evaluation：精煉后的答案經(jīng)過評分從而采樣一個獎勵值，并計算其 Q 值。這涉及模型自我獎勵反饋和約束，如嚴格的評分標準和抑制滿分，以確保評分的可靠性和公平性；
• 反向傳播：將精煉答案的值反向傳播到其父節(jié)點和其他相關節(jié)點，以更新樹的值信息。如果任何子節(jié)點的 Q 值發(fā)生變化，則更新父節(jié)點的 Q；
• UCT 更新：在所有節(jié)點的 Q 值更新完成后，確定一個候選節(jié)點集合 C，用于進一步擴展或選擇，然后使用 UCT 更新公式更新所有節(jié)點的 UCT 值，以備下一步的選擇階段。

迭代上述階段，直到滿足終止條件 T 為止。

Self-Refine 

在 self-refine 階段， 模型通過多輪對話完善提示來優(yōu)化針對問題 P 的答案 a。首先，模型生成一個關于答案 a 的反思性或批判性評論 m。隨后，在 m 的指導下，模型修改答案 a，產(chǎn)生一個改進版本 a'，這種迭代的精煉方式提高了模型響應質量。

自評估

在數(shù)學問題 P 的答案精煉過程中，一個答案 a 的 Q 值被定義為將 a 進一步精煉成更優(yōu)答案的預期質量。這個定義是基于從 a 到其重寫形式的轉換具有馬爾可夫性質，即下一個狀態(tài)（即改寫后的答案）僅依賴于當前狀態(tài)（即當前的答案 a），而與之前的狀態(tài)無關。

此外，研究者還設計了三個約束：提示約束、滿分抑制、重復采樣。采樣后，計算 a 的 Q 值。

反向傳播

在所有葉節(jié)點的獎勵值經(jīng)過采樣和 Q 值更新完成后，然后將這些變化傳播至其父節(jié)點和祖節(jié)點。在這個更新過程中，如果節(jié)點 a 的子節(jié)點集合 Children (a) 中任何元素的 Q 函數(shù)值發(fā)生變化，那么節(jié)點 a 的 Q 函數(shù)值也將進行更新。這樣的傳播確保了節(jié)點的 Q 值能夠反映其所有可能子節(jié)點的最新狀態(tài)和評估。 

更新 UCT 和選擇

在更新了樹中所有節(jié)點的 Q 值之后，會進入下一輪選擇階段。這個過程包括以下步驟：

• 候選節(jié)點選擇：在選擇節(jié)點時，研究者無需從根節(jié)點開始，而是按層次順序遍歷樹中的節(jié)點。
• UCT 更新：借鑒 AlphaGo，該研究使用 UCT 和 UCB-1 方法來平衡節(jié)點的探索和利用；對于候選集 C 中的節(jié)點 a，其 UCT_a 值為：

終止函數(shù)

提前終止：當搜索結果的改進開始減少或連續(xù)搜索產(chǎn)生重復結果時，終止發(fā)生。

搜索約束：一旦展開次數(shù)達到預定限制或樹中的一個或多個節(jié)點滿足最大深度約束，搜索就會終止。

實驗結果

為了評估 MCTSr 算法在解決數(shù)學問題中的有效性，研究者將 LLaMA3-8B 作為基礎模型，并使用 MCTSr 進行增強。他們在 Zero-Shot CoT、Self-Refine、4-rollouts MCTSr 和 8-rollouts MCTSr 等幾種設置中，將 LLaMA3-8B 與 GPT-4、Claude 3 和 Gemini 1.5-Pro 等進行了比較。

研究者在 GSM8K 和 GSM-hard 測試集（它們分別包含了典型和具有挑戰(zhàn)性的數(shù)學問題）上評估了上述方法，結果如下表 1 所示。

可以發(fā)現(xiàn)，MCTSr 的 rollout 次數(shù)與成功率之間存在著直接相關性，并隨著迭代次數(shù)增加而顯著提升，在不太復雜的 GSM8K 中尤為明顯。不過對于更復雜的 GSM-Hard 測試集，即使 rollout 次數(shù)更高也會達到性能上限，表明當前策略在解決復雜問題時存在局限性。

這些結果強調了 MCT-Self-refine 算法的穩(wěn)健性和潛在邊界，以及持續(xù)改進的必要性，從而有效應對更復雜的挑戰(zhàn)。

下表 2 展示了在 MATH 數(shù)據(jù)集上應用不同復雜度級別的 MCT-Self-refine 算法的結果。數(shù)據(jù)集分為五個難度級別，從 Level 1（最簡單）到 Level 5（最具挑戰(zhàn)性）。

結果顯示，Level 1 的成功率最高，8 次 rollout 后，MCTSr 實現(xiàn)了 90.16% 的成功率，解決了 437 個問題中的 394 個。隨著 rollout 次數(shù)的增加，這一級別的成功率顯著提高。

在最具挑戰(zhàn)性的 Level 5 難度，8 次 rollout 后，MCTSr 的成功率為 34.06%，解決了 1324 個問題中的 451 個。這說明了隨著難度不斷增加，該算法在高度復雜的場景中性能受到限制。

所有級別的整體性能顯示，8 次 rollout 后，MCTSr 的累計成功率為 58.24%，解決了 5000 個問題中的 2912 個。這一成功率相較于 Zero-Shot CoT 的初始成功率 24.36% 有了顯著提高。這表明了，rollout 次數(shù)的增加與成功率的提高呈現(xiàn)出一致性，強調了 MCT-Self-refine 算法在提升不同數(shù)學復雜度級別的問題解決能力方面的有效性。

這些結果還驗證了 MCT-Self-refine 算法在學術和問題解決上下文中的潛力，并強調了其對 MATH 數(shù)據(jù)集中不同復雜度級別問題的可擴展性和適應性。

下表 3 為 MCT-Self-refne 算法在奧數(shù)競賽的三個數(shù)據(jù)集上進行了測試：AlME、GAIC Math Odyssey 和 OlympiadBench。

AIME：從 Zero-Shot CoT 的 2.36%（解決 22 個問題）到 MCTSr 的 11.79%（解決 110 個問題）。

GAIC Math Odyssey：成功率從 17.22%（解決 67 個問題）上升至 49.36%（解決 192 個問題）。

OlympiadBench：從 Zero-Shot CoT 的 1.25%（解決 16 個問題）提高到 MCTSr 的 7.76%（解決 99 個問題）。

這些結果證實了 MCT-Self-refine 算法在未見過的數(shù)學問題上的適用性，表明其在奧林匹克等競爭性學術環(huán)境中具有優(yōu)勢。

如表 4 所示。與當前閉源大模型進行比較時，MCTSr 可以有效提升小參數(shù)開源模型（如 LLaMa-3）的數(shù)學推理能力到相當?shù)乃健?/span>

更多技術細節(jié)和實驗結果請參閱原論文。