像AlphaZero這樣的遊戲神經網絡透過在測試時結合搜索機制，並將增強後的策略蒸餾回網絡，達到超越人類的棋盤遊戲表現。那麼，為什麼類似技術今天還沒被廣泛用於語言模型呢？DeepSeek-R1的作者提到他們在使用蒙地卡羅樹狀搜索（MCTS）時成效有限；Finbarr Timbers則在一篇優秀的文章中分析，問題可能出在他們選擇了UCT而非pUCT。

本文旨在探討兩個問題：

為此，我將MCTS應用於Qwen-2.5-1.5B-Instruct的推理步驟中，尋找更強的推理軌跡，並透過線上PPO迴圈將其蒸餾回模型。在組合算術遊戲Countdown的任務中，蒸餾後的模型（無搜索輔助下評估）達到11.3%的漸近mean@16評分，優於CISPO的8.4%和best-of-N的7.7%。相較於未經強化學習的指令模型（3.1%），提升了8.2個百分點。

分數偏低反映出這是基於1.5B模型的小規模實驗。我希望這篇文章是系列的第一篇，未來隨著使用更大模型和更多計算資源，分數能持續提升。

我最初嘗試用GSM8K作為測試環境，但GRPO與MCTS間差異不大，難以得出明確結論。於是改用Countdown遊戲作為環境。遊戲規則簡單：給定N個正整數，使用加減乘除計算目標數字。選擇Countdown的假設是，組合問題更能受益於樹狀搜索所帶來的並行自適應推理，而不像GSM8K那樣，序列推理也能有效解決問題。我們用2萬筆樣本訓練，測試集有820筆，每筆包含四個1到13之間的整數。

我發現使用稀疏獎勵（正確為1，錯誤為0）訓練時不穩定，改用密集獎勵函數：

若格式正確，獎勵為 $1.0 - 2 \cdot \min\left(\frac{|t - p|}{t}, 1.0\right)$，否則為-1.0。

其中$t$為真實目標，$p$為預測目標。

但評估仍使用稀疏獎勵，方便直觀理解分數（如通過率）。

MCTS算法已有詳細介紹，這裡不贅述；本文重點在於傳統MCTS與我嘗試方法的差異。簡言之，MCTS透過價值函數引導，迭代構建搜索樹以智能探索行動空間。

棋盤遊戲的行動空間相對有意義，每步棋對勝負影響大。相比之下，語言模型中許多推理過程的標記充當填充或語法糖，從top-k概率分支（或基於熵閾值）不一定帶來多樣化搜索。想像一個狀態，下一個可能標記是“but”、“however”、“yet”等，我們可能浪費大量計算資源建構龐大搜索樹，但每個標記的收益有限。

我偏好Tree-of-Thoughts（Yao等，2023）提出的方法，搜尋可能的下一推理步驟。此方法將每個節點狀態定義為連續標記序列。

為探索更多擴展參數，我的實作使用並行MCTS，N個代理共享同一樣本的搜索樹，並利用虛擬損失促進搜索多樣性。

從每個葉節點開始，我們生成K個完成序列，直到遇到停止標籤</step>。這K個序列構成該節點的行動空間。

由於pUCT需要行動層級先驗，我們計算序列級別的累積對數概率並套用softmax，得到相對先驗。這樣做可避免原始序列概率過小導致數值不穩定。

MCTS通常使用一個價值頭$V(s_t)$，透過訓練提升，幫助引導搜索至更佳軌跡。此價值頭由MLP和tanh函數組成，作用於transformer的最終隱藏狀態。

此方法與TS-LLM（Feng等，2023）類似，後者也結合AlphaZero風格樹狀搜索與句子層行動的學習價值函數。

棋盤遊戲MCTS的訓練信號通常來自最小化根節點搜索策略與模型原始策略間的KL散度。但由於我們的行動空間粒度（推理步驟vs標記）不匹配，需採用其他方法。我使用的做法是，所有工作者完成M次迭代後，進行貪婪選擇。

被指定為“訓練者”的工作者非同步從共享緩衝區拉取樣本，使用AdamW優化器，對每批B樣本執行單次PPO內部步驟，損失類型為CISPO。

訓練目標是最小化總損失$L_{total}$：

$L_{total} = c_{ppo} L_{ppo} + c_{value} L_{value} + c_{KL}\, \mathbb{D}_{KL}(\pi_\theta \mid\mid \pi_{ref})$

其中

$L_{cispo} = -\mathbb{E}[sg(\min(\frac{\pi_\theta(a_t \mid s_t)}{\pi_{old}(a_t \mid s_t)}), \epsilon) \cdot A_t \cdot \log \pi_\theta(a_t \mid s_t)]$

$A_t = r_{terminal} - sg(V_{old}(s_t))$為標記層優勢（同一終端獎勵分配給每個標記）。我未使用GAE，因為推理軌跡可達數千標記，且終端獎勵會使早期標記的折扣值極小。

$L_{value} = \mathbb{E}[(V(s_t) - r)^2]$

$\mathbb{D}_{KL}(\pi_\theta \mid\mid \pi_{ref}) = \frac{\pi_\theta(a_t \mid s_t)}{\pi_{ref}(a_t \mid s_t)} - \log \frac{\pi_\theta(a_t \mid s_t)}{\pi _{ref}(a_t \mid s_t)} - 1$（取自DeepSeek-R1論文）

訓練持續至評估分數趨於穩定。

所有實驗在Andromeda的8xH100節點上進行。MCTS中，六個GPU作為生成器，兩個作為訓練者。Rust工作者從數據集中抽樣問題，透過gRPC向生成器池發送推理請求，將選擇的軌跡寫入Redis流；訓練者從中拉取樣本。生成器與訓練者每8次梯度步驟同步權重，使用Redis pub/sub。

我以全局批次大小128、組大小16運行CISPO基線，有效批次大小2048。logits以float32計算，與ScaleRL一致。訓練至評估分數穩定，所有八個GPU均用於CISPO訓練，無訓練者/生成器分離。

為隔離樹結構的價值，我進行了軌跡以“best-of-N”（N=64）選擇而非樹搜索的實驗。

評估使用mean@16，即每個評估提示生成16次結果，以稀疏0/1獎勵評分並平均。評估時，MCTS蒸餾策略（無搜索輔助）達11.3%漸近mean@16分數，CISPO為8.4%，Best-of-N最差，停留在7.7%。

令人驚訝的是，儘管訓練獎勵顯著較高，Best-of-N蒸餾在評估中表現不如CISPO和MCTS。可能原因是：若模型在推理過程中有98%機率至少犯一次錯誤，64次選擇中仍有約72.6%機率選中至少一條正確軌跡。但若無激勵每次都產生穩健推理，模型不太可能學會提升單次推理表現的策略。就像我中學時用多種技巧追蹤數學中間步驟，減少考試中「愚蠢錯誤」的機率；若能多次考試，我可能不會用這些技巧！

所有代碼均為開源，詳見原文鏈接。

這意味著什麼？令我興奮的是，我們可以調整更多推理參數，如每棵樹的並行工作者數量、MCTS迭代次數。雖未充分調參，初步實驗顯示增加這些值能顯著提升性能。我期待進一步探索此方向，並在更大模型與計算資源下驗證其潛力。歡迎合作！

當然，這可能是“小模型現象”，方法在大模型上不如GRPO等強。是否能調整GRPO基線匹配MCTS？或許，但ScaleRL發現GRPO大多數超參數影響計算效率，非最終獎勵上限。

有人會說MCTS每樣本推理計算量大於GRPO，表現好是理所當然！但目標非做完全公平的計算比較；MCTS雖用更多推理計算，卻提供更多調節與擴展計算的槓桿，提升獎勵上限。反觀GRPO，投入百倍計算不見得能突破瓶頸。

感謝Andromeda團隊與Molly Mielke McCarthy贊助計算資源，Tom McCarthy與Joe Melkonian審閱初稿並提供寶貴意見。也感謝Finbarr Timbers的文章激發本研究。

訓練過程中各次運行的日誌頻率不同，故x軸限制略有差異。模型訓練至評估分數趨於穩定為止。