近年來，大型語言模型 (LLM) 在解決複雜推論任務方面取得了顯著進展，特別是在結合強化學習 (RL) 進行對齊訓練之後。然而，傳統的 RL 方法 (如 PPO、GRPO) 存在兩個根本性挑戰：一是運算效率低，每個查詢需要獨立生成多條軌跡，導致大量的重複計算；二是探索能力有限，模型難以有效探索多樣化的推論路徑，尤其是在獎勵稀疏、延遲高的環境中。
論文：TreePO: Bridging the Gap of Policy Optimization and Efficacy and Inference Efficiency with Heuristic Tree-based Modeling
連結：https://arxiv.org/pdf/2508.17445
針對這些問題，字節研究團隊聯合高校提出了一種全新的方法——TreePO（Tree-based Policy Optimization）。該方法將序列生成過程重新建模為樹狀搜尋結構，透過共享前綴、動態分支和早期剪枝，顯著提升了取樣效率和探索多樣性。同時，TreePO 還引入了一種層次化的優勢估計函數，能更精細地分配獎勵訊號，提升訓練穩定性。
本文將帶你深入解讀這篇高品質論文，理解 TreePO 如何在不犧牲性能的前提下，實現運算效率的大幅提升，以及它為何被認為是推動 RL 訓練規模化的重要一步。
研究動機與核心思想
傳統 RL 方法在訓練 LLM 時，通常會對同一個提示 (prompt) 獨立取樣多條完整軌跡 (trajectories)。這種方式雖然簡單，但存在明顯缺陷：
運算冗餘：不同軌跡往往在開頭部分共享相同的前綴 (如問題解析、變數定義)，但每條軌跡都要獨立計算這些部分，導致 KV 快取無法共享，運算資源浪費嚴重。
探索效率低：模型可能會反覆生成相似或錯誤的路徑，缺乏智能的分支與剪枝機制。
TreePO 的靈感來源於一個關鍵觀察：推論路徑天然呈樹狀結構。即不同答案往往在前期共享相同的推論步驟，後期才分岔。因此，如果能將生成過程組織成樹，就可以：
共享運算：只需計算一次公共前綴，後續分支復用其 KV 快取。
動態探索：在不確定性高的節點處主動分支，探索更多可能性；在確定錯誤或無效的路徑上提前停止。
TreePO 的核心思想正是將傳統的「平坦式」取樣轉變為「樹狀」取樣，並配以相應的優勢估計機制，從而實現高效且穩定的訓練。
TreePO 方法詳解
樹狀取樣機制 (Tree-based Rollout)
TreePO 將生成長序列的過程拆分為多個段 (segment)，每段包含固定數量的 token (如 256、512 等)。取樣過程如下：
段級解碼：模型每次生成一個段，而不是一個 token。
分支控制：在每個段結束後，系統根據啟發式策略 (如機率高低) 決定是否分支 (即生成多個後續段) 或停止。
回溯機制：如果某些路徑提前終止 (如生成 [EOS] 或錯誤模式)，系統會從其他活躍路徑中分配更多預算，保持總體取樣數不變。
這種設計極大地提高了 GPU 利用率，尤其是在生成長序列時，透過批次處理和快取共享，顯著減少了計算時間。
左圖和中圖：顯示其訓練穩定性優於傳統方法；右圖：TreePO 採樣的樹狀結構
啟發式取樣控制
TreePO 允許在以下環節注入啟發式策略：
早期停止：檢測到重複子字串或錯誤模式時立即停止目前路徑。
分支分配：根據目前段的生成機率動態分配分支數，鼓勵探索低機率但可能有效的路徑。
回溯選擇：僅對已生成完整答案或合理終止的路徑進行回溯，避免無效探索。
樹狀優勢估計 (Tree-based Advantage Estimation)
傳統的優勢函數 (如 GRPO) 通常在軌跡層級計算，容易因響應相似而導致訊號模糊。TreePO 利用樹結構，提出了一種層次化優勢估計方法。
假設一條軌跡 $\tau_i = (s_0, a_0, \dots, s_L, a_L)$ 由多個段組成：
定義子群 $\mathcal{S}_d(\tau_i)$ 為在深度 $d$ 處共享同一前綴的所有軌跡集合。則優勢函數可寫為：
$A(x_t, \tau_i) = V(x_t, \mathcal{S}_d(\tau_i))$
$\tau_i$：token $x_t$ 在軌跡 $\tau_i$ 中的最終優勢值。
$V(x_t, \mathcal{S}_d(\tau_i))$：在子群 $\mathcal{S}_d(\tau_i)$ 中的局部優勢。
$R(\tau_i)$：軌跡 $\tau_i$ 的獎勵。
$R(\mathcal{S}_d(\tau_i))$：子群 $\mathcal{S}_d(\tau_i)$ 的平均獎勵。
透過多層子群聚合，優勢估計能更精細地反映每個 token 對最終獎勵的貢獻，尤其當不同路徑在中期開始分岔時，能更好地區分各路徑的價值。
如何基於樹結構計算子群優勢
實驗設計與主要結果
實驗設置
模型：Qwen2.5-7B 系列模型 (基礎版、指令版、數學版)。
數據集：MATH、AIME、AMC、MINERVA、Olympiad Bench 等數學推論基準。
評估指標：多數投票準確率 (Major@16)、每秒 Token 數 (TokenPS)、每秒軌跡數 (TrajPS)、GPU 小時。
主要結果
TreePO 在所有數據集上均顯著超越 GRPO 基準線，總體準確率從 46.63% 提升至 58.21%。
TreePO 在保持性能的同時，GPU 小時節省 12%~43%，效率提升顯著。
效率分析
不同深度-段長配置下的效率對比，說明存在最佳權衡點。
隨著取樣數增加，TreePO 保持近乎線性的擴展性，而傳統方法增益有限。
討論與分析
效率-性能權衡
TreePO 在多數設置下都能實現更高的訓練效率和穩定性，儘管在某些情況下峰值準確率略低於傳統方法 (如 58.06% vs. 58.21%)，但其節省的運算成本 (22%~43%) 使得這種取捨非常值得。
樹結構參數的影響
深度 vs. 段長：更深的樹 (更多分支) 適合解碼平行化，但段長短會導致更多重複計算；淺層的樹 (段長長) 適合預填充優化，但探索多樣性受限。
最佳配置：不同模型有不同偏好，例如：
Qwen2.5-7B-Instruct：深度 28 (段長短)
Qwen2.5-Math-7B：深度 14 (段長長)
啟發式分支策略的局限性
(單純鼓勵低機率分支會導致性能下降，說明探索需有方向性，不能盲目。)
運算擴展性
TreePO 提供一族擴展曲線，用戶可根據運算預算選擇最佳分支因子 (d=2,4,8)，實現「運算最佳推論」。
結論
TreePO 透過將強化學習中的序列生成過程重新建模為樹狀搜尋，成功解決了傳統方法在運算效率和探索多樣性方面的瓶頸。實驗表明，TreePO 不僅能顯著提升訓練效率 (節省 22%~43% 的 GPU 時間)，還能保持甚至提升模型性能。
未來，TreePO 的框架可進一步擴展至多輪對話、工具呼叫、多代理人協作等更複雜的任務中，為實現更高效、更智能的大模型訓練奠定基礎。