大模型+蒙特卡洛樹搜索，一招讓LLaMa-3 8B奧數水平直逼GPT-4

作者：機器之心 2024-06-17 12:39:37

通過算法層面的創新，未來大語言模型做數學題的水平會不斷地提高。

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

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

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

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

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

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

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

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

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

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

方法概覽

MCTSr 架構圖如圖 1 所示：

MCTSr 工作流包括：

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

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

Self-Refine

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

自評估

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

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

反向傳播

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

更新 UCT 和選擇

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

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

終止函數

提前終止：當搜索結果的改進開始減少或連續搜索產生重復結果時，終止發生。

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

實驗結果

為了評估 MCTSr 算法在解決數學問題中的有效性，研究者將 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 測試集（它們分別包含了典型和具有挑戰性的數學問題）上評估了上述方法，結果如下表 1 所示。

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

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

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

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

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

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

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