透過演算法層面的創新，未來大語言模型做數學題的水平會不斷地提高。
這幾天，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 資料集中不同複雜度級別問題的可擴充套件性和適應性。
下表 3 為 MCT-Self-refne 演算法在奧數競賽的三個資料集上進行了測試：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 演算法在未見過的數學問題上的適用性，表明其在奧林匹克等競爭性學術環境中具有優勢。
如表 4 所示。與當前閉源大模型進行比較時，MCTSr 可以有效提升小引數開源模型（如 LLaMa-3）的數學推理能力到相當的水平。
更多技術細節和實驗結果請參閱原論文。