萬梓煜是上海交通大學的三年級在讀博士生，導師為溫穎教授和張偉楠教授，主要研究興趣為強化學習與大語言模型、決策大模型。馮熙棟是倫敦大學學院四年級博士生，導師為汪軍老師。同時目前也是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 式的蒙特卡洛樹搜索的高效性與進一步迭代優化大語言模型本身的可能。