沒想到！AlphaZero式樹搜尋也能用來增強大語言模型推理與訓練

AIxiv專欄是機器之心釋出學術、技術內容的欄目。過去數年，機器之心AIxiv專欄接收報道了2000多篇內容，覆蓋全球各大高校與企業的頂級實驗室，有效促進了學術交流與傳播。如果您有優秀的工作想要分享，歡迎投稿或者聯絡報道。投稿郵箱：liyazhou@jiqizhixin.com；zhaoyunfeng@jiqizhixin.com

萬梓煜是上海交通大學的三年級在讀博士生，導師為溫穎教授和張偉楠教授，主要研究興趣為強化學習與大語言模型、決策大模型。馮熙棟是倫敦大學學院四年級博士生，導師為汪軍老師。同時目前也是Google DeepMind的student researcher。主要研究方向是強化學習與大語言模型，多智慧體以及元強化學習。

2016年 DeepMind 的 AlphaZero 展示了強大的學習和適應能力，登上《自然》雜誌封面，並在之後通過自我對弈不斷提升自身水平，最終戰勝了人類冠軍，而這也為之後學者在大語言模型與樹搜尋的結構化結合奠定了基礎。

大語言模型樹搜尋

大語言模型與思維鏈（Chain-of-Thought, CoT）的結合增強了其複雜推理能力，使其在數學和邏輯推理等任務上表現更佳。然而，語言模型仍存在誤差：一方面，受資料數量和質量的影響，大語言模型在複雜任務上仍與專家系統和求解器有差距；另一方面，僅依靠大語言模型難以解決長程規劃（long-horizon planning）任務。

為解決這些問題，研究者提出了將結構化的樹/圖搜尋與大語言模型結合的方式。思維樹（Tree of Thought, ToT）模仿人類認知中的慢系統，利用深度/廣度優先搜尋顯著提升大語言模型的規劃能力。Reasoning via Planning (RAP) 則將大語言模型的思維鏈過程視作規劃（planning），使用其內在知識進行狀態評估，並結合傳統蒙特卡洛樹搜尋（MCTS），從而增強語言模型的效能。這些方法利用大語言模型的多工能力，通過提示工程（prompt engineering）對中間結果進行價值判斷。

然而，這種方法並不普遍適用。評估多步推理問題的中間狀態本身也是一個推理子問題，依賴CoT生成評估，無法保證評估的可靠性。此外，語言模型的自我評估能力和逆轉詛咒問題，以及子問題難度降低不顯著等因素，限制了這類方法的應用效果，尤其是對於較小規模、易部署的模型。

一個潛在的解決方案是參考AlphaZero。2016年，DeepMind通過AlphaZero在複雜多步推理問題如圍棋上取得突破性進展。AlphaZero結合了傳統MCTS和深度神經網路的優勢，使用價值函式學習簡化了MCTS中的Simulation/Rollout過程，並通過蒸餾樹搜尋增強策略的迭代優化，為大語言模型在樹搜尋評估和長程規劃效率問題上提供了方向。

基於此，來自倫敦大學學院，上海交通大學，卡耐基梅隆大學的合作團隊將 AlphaZero 方法精髓與大語言模型的文字生成結合，提出了大語言模型樹搜尋訓練增強框架 TSLLM。通過將這一任務建模為多步決策問題，引入強化學習中價值函式學習的概念，以訓練的方式微調一個價值函式以提供更為魯棒可靠的搜尋中間價值評估。與此同時在 TSLLM 中也實現了不同樹搜尋演算法的對比，尤其是探究了價值函式結合的簡化 MCTS 在不同型別任務上的優缺點。最後，團隊探究了由樹搜尋引導的迭代優化方式對大語言模型進一步優化的可能性。目前，該論文已被ICML 2024接收。

• 論文名稱：AlphaZero-Like Tree-Search can Guide Large Language Model Decoding and Training
• 論文連結：https://arxiv.org/abs/2309.17179
• 程式碼連結：https://github.com/waterhorse1/LLM_Tree_Search

TSLLM的基本框架如下：

                           ^{圖一：TSLLM 基本框架一覽}

TSLLM 有如下特點：

• TSLLM 是一個普遍適用和可擴充套件的框架，通過學習價值函式可應用於幾乎任何任務，以及任何大小的語言模型。
• TSLLM 在不同問題上驗證了樹搜尋能增強大語言模型推理階段表現的同時，也進一步驗證了其迭代增強語言模型作為一個語言模型訓練新正規化的潛力。
• 在設計上，TSLLM 支援逐句/詞元細粒度的搜尋。
• 使用可靠魯棒的價值函式作為狀態評估，TSLLM 支援包括簡單的 BFS/DFS，傳統 MCTS，AlphaZero 式的 MCTS-α，MCTS-Rollout 等演算法。
• TSLLM進行了全面且公平的對比。例如，爲了實現與非搜尋演算法（如 CoT/CoT-SC）的公平對比，TSLLM 通過統計總體計算量的方式衡量不同演算法的效果與效率。

AlphaZero 式的樹搜尋增強的大語言模型

研究團隊將大語言模型的自迴歸生成過程建模為一個多步決策問題，定義詞元/句級的語言生成過程的概率建模。對於給定的自然語言任務，他們通過學習的價值函式估計與最終獎勵估計，建模自然語言任務生成過程中的期望回報與稀疏獎勵。並通過樹搜尋的方式在推理與訓練階段增強大語言模型的能力。

大語言模型推理階段增強：

基於學習得到的價值函式，TSLLM 中實現了不同的樹搜尋演算法，從簡單的價值函式引導的廣度/深度搜索（DFS/BFS-V）到傳統 MCTS，該團隊基於 AlphaZero 的中間價值回傳思想實現了 MCTS-α，以及提出離線搜尋變種MCTS-Rollout。下圖比較了傳統 MCTS 與 AlphaZero 式的 MCTS 的主要區別，如圖所示，傳統 MCTS 需要通過模擬（Simulation）達到停止節點，纔會開始價值回傳。另外，在 TSLLM 中，他們還討論並實現了多條搜尋路徑的聚合形式，以及提出了考慮計算量的公平比較方式。

                             ^{圖二：傳統蒙特卡洛樹搜尋（左）與AlphaZero 式的蒙特卡洛樹搜尋（右）對比。}

大語言模型迭代訓練增強：

最後，在 TSLLM 中，研究團隊指出樹搜尋還能夠進一步強化大語言模型本身。他們將結構化搜尋作為一個策略增強運算元（ Policy Improvement Operator），利用這個運算元本身，可以迭代式的優化語言模型策略與價值/獎勵估計函式。在 TSLLM 中，研究團隊類比 AlphaZero/Expert Iteration 中的迭代式優化方法：一方面通過 supervised finetuning 蒸餾這一更好表現的策略分佈，而另一方面持續微調價值函式估計。通過這樣的方式不斷持續增強大語言模型策略本身的任務解決能力與價值函式指導的搜尋增強策略的能力。

實驗結果

在實驗過程中，團隊非常重視樹搜尋演算法和基線演算法的合理對比。針對於一些演算法評估的不合理現象，團隊強調了：

• 演算法的合理對比。樹搜尋演算法天生會帶來更高的計算複雜度，合理的演算法對比應在相似的計算量上進行。
• 選擇合理基線和設定。例如，團隊發現一個經常被忽略的基線演算法：Majority-Vote + Outcome Reward Model。實驗中團隊發現其可以作為簡單卻非常強大的基線，在GSM8K上甚至可以超過樹搜尋演算法。同時團隊嚴格避免了不合理的實驗設定：如利用測試集的真值進行樹搜尋回溯。

在實驗中，團隊在數學推理/規劃，邏輯推理任務，價值對齊的文字生成以及文字化的決策推理任務上進行了豐富的實驗與對比，在搜尋深度上最大達到了 64，顯著深於之前的演算法。

微調學習的價值函式有更可靠、更好的泛化能力：

團隊首先就基於學習的價值函式與大語言模型自我評估的方式進行對比。下表實驗結果表明基於學習的價值函式無論是在 GSM8k還是 Game24 問題上都優於 ChatGPT，即使在評估 ChatGPT 本身作為策略的中間狀態時，表現依舊優於 ChatGPT 本身，充分說明了基於學習的價值函式的可靠性與魯棒性。

不同的搜尋演算法具有其不同適應性：

團隊比較TSLLM 中不同樹搜尋演算法在類似計算量下的結果， 發現 MCTS-α與 MCTS-Rollout的搜尋演算法主要在長程推理（搜尋樹深度較大，如 Alignment，Endgame）問題上，顯著優於其他搜尋演算法。但對於輕量級或搜尋深度較低的問題，BFS/DFS也具有較強的效能。

樹搜尋演算法的Scaling（擴充套件）性質受限：

同時，團隊也嘗試了對樹搜尋的取樣次數進行scaling。結果發現，絕大部分樹搜算演算法都可以隨著路徑搜尋次數的增加而獲得性能提升。但同時團隊也發現，一些簡單的baseline （如COT-SC+ORM）具有更好的scaling性質。例如在GSM8K中， COT-SC+ORM的baseline可以取得比樹搜尋更優越的效能與擴充套件屬性。基於這個結果，團隊認為如何獲得更為優越的擴充套件效能將成為樹搜尋演算法未來的重要研究方向。 

團隊的實驗也驗證了迭代優化可以進一步提升大語言模型的能力：

下文左圖展示了 TSLLM 通過 MCTS-α在訓練問題集上搜索迭代的結果，在 GSM8k 上相比於初始的策略模型 ，通過一輪樹搜尋演算法的搜尋迭代優化後的策略 的表現優於在 5、10 倍取樣資料上進行 Rejection Sampling 優化的結果；然後也可以發現，在RLHF 資料集上， 的表現依舊不如 PPO，這主要是由於 PPO 對語言模型引數進行了多次線上迭代優化。當對比迭代後的價值函式 我們也可以發現，其能夠進一步增強 的樹搜索結果。

總結

總結來說，在這一工作中，研究團隊提出了大語言模型的樹搜尋推理與訓練增強框架 TSLLM，在經驗結果上強調了可靠中間價值評估的重要性與樹搜尋演算法在不同問題上的效果與效率，驗證了 AlphaZero 式的蒙特卡洛樹搜尋的高效性與進一步迭代優化大語言模型本身的可能。