大模型強化學習新突破——SPO新范式助力大模型推理能力提升!
當前,強化學習(RL)在提升大語言模型(LLM)推理能力方面展現出巨大潛力。DeepSeek R1、Kimi K1.5 和 Qwen 3 等模型充分證明了 RL 在增強 LLM 復雜推理能力方面的有效性。
然而,要實現有效的強化學習,需要解決一個根本性的挑戰,即信用分配問題(credit assignment):在大語言模型的場景下,如何將整個序列(LLM 的回復)最終的評估結果,歸因到序列中具體的決策動作(token)上。
這一問題的困難在于獎勵信號非常稀疏 — 只能在序列結束時才能獲得明確的成功或失敗反饋。
當前主要方法
在強化學習中,通常采用優勢值估計(advantage estimation)的方法來解決信用分配問題。目前針對大語言模型的強化學習方法主要分為兩類,它們之間的區別在于優勢值估計的粒度不同。
粗粒度的軌跡級 (trajectory-level) 方法,如 DeepSeek R1 使用的 GRPO,只根據最終的獎勵為整個序列計算一個優勢值。這種方法雖然高效但反饋信號過于粗糙,LLM 無法對錯誤回答中正確的部分進行獎勵,也無法對正確回答中冗余的部分進行懲罰。
另一種極端是細粒度的 token 級(token-level)方法,如經典的 PPO。這類方法為每個 token 估計優勢值,需要依賴額外的 critic 模型來預測每個 token 的狀態價值(V 值)。然而,在大語言模型的強化學習任務中,不同 prompt 對應的軌跡分布差異很大,而且在訓練過程中每個 prompt 采樣出來的模型回復數量非常有限,critic 模型難以訓練好,造成 token 級的優勢值估計誤差很大。
新的 SPO 框架
為突破這一瓶頸,來自中科院軟件所和香港城市大學的的研究團隊創新性提出了 Segment Policy Optimization (SPO) 框架。
- 論文題目:Segment Policy Optimization: Effective Segment-Level Credit Assignment in RL for Large Language Models
- 作者:Yiran Guo, Lijie Xu, Jie Liu, Dan Ye, Shuang Qiu
- 鏈接:https://arxiv.org/abs/2505.23564
- 代碼鏈接:https://github.com/AIFrameResearch/SPO
SPO 使用了一種中等粒度的段級(segment-level)優勢值估計方式。它不像軌跡級方法只在最后一步計算優勢,也不像 token 級方法每步都計算優勢,而是將生成的序列劃分為若干相連的段,計算每個段的優勢值。
這種段級的優勢值估計方式具有幾個明顯的優勢:
(1) 更優的信用分配:相比軌跡級方法,段級方法能夠提供更局部化的優勢反饋,讓模型能夠獎勵錯誤回答中仍然有價值的部分,同時也能懲罰正確回答中冗余和無效的片段。
(2) 更準確的優勢值估計:相比 token 級方法,段級方法所需的估計點數量更少,從而能夠有效利用蒙特卡洛(Monte Carlo, MC)采樣得到更加準確且無偏的優勢值估計,而無需再依賴額外且不穩定的 critic 模型。
(3) 更靈活、更易調整:段級的劃分方式可以任意定義,并不要求語義上的完整性,因此可以靈活地在 token 級與軌跡級之間自由調整粒度,并且可以適應不同的任務和應用場景。
SPO 框架主要包含三個核心部分:(1) 靈活的段級劃分策略;(2) 基于蒙特卡洛采樣的段級優勢值估計;(3) 利用段級優勢值進行策略優化。
這種模塊化的設計使框架具備高度的靈活性,不同的部分可以有不同的實現策略,以適用不同的應用場景。
該團隊進一步針對不同的推理場景提出 SPO 框架的兩個具體實例:對于短的思維鏈(chain-of-thought, CoT)場景,提出了 SPO-chain,該方法使用基于切分點(cutpoint-based)的段劃分和鏈式優勢值估計;對于長 CoT 場景,提出極大提升 MC 采樣效率的樹形結構優勢值估計方法。
此外,該團隊還提出了一種 token 概率掩碼(token probability-mask)策略優化方法,選擇性的對段內的低概率 token 計算損失而非段內的所有 token。作者認為這些 token 是模型推理軌跡可能發生分叉的地方,是段級優勢值產生的主要原因。這種方法可以用于 SPO-chain 和 SPO-tree,從而進一步強化信用分配。
框架及核心技術
SPO 框架主要圍繞以下三個具有挑戰性的問題進行設計:(1) 如何將生成的序列劃分為多個段?(2) 如何準確且高效地估計每個段對應的優勢值?(3) 如何利用段級優勢值來更新策略?SPO 的三個核心模塊分別解答上面三個問題,每個模塊包含多種可選策略,來適用于不同的場景:
1. 段劃分 (Segment Partition):
a) 基于切分點的段劃分 (Cutpoint-based Partition): 為短思維鏈場景設計,將段劃分點放置在狀態值(V 值)更有可能發生變化的地方。根據 token 概率動態確定段邊界,優先在模型 “猶豫” 或可能改變推理路徑的關鍵點(cutpoints)進行劃分,使信用分配更精確。比如,在下圖例子中,標記為紅色的 token 是關鍵點,而標記為藍色的豎杠是分段結果。
b) 固定 token 數量段劃分 (Fixed Token Count Partition): 將序列劃分為固定長度的段,便于樹形結構的組織和優勢值估計,為 SPO-tree 設計。
2. 段級優勢值估計(Segment Advantage Estimation):
a) 鏈式優勢值估計 (Chain-based) 方法:在短思維鏈場景下,MC 采樣的成本不高,該團隊采用一種直接的段級優勢值估計方式,獨立估計每個段邊界的狀態值(V 值),然后計算段級優勢值。以下公式展示了鏈式優勢值的估計方法。
b) 樹形優勢值估計 (Tree-based): 在長思維鏈場景下,MC 估計的代價很高,團隊提出了一種高效的樹形估計方法:將采樣軌跡組織成樹形結構,通過自底向上的獎勵聚合計算狀態價值(V 值),同一個父節點的子節點形成一個組,在組內計算每個段的優勢值。這種方式將用于 V 值估計的樣本同時用于策略優化,極大提高了樣本效率。以下公式展示了樹形優勢值估計方法。
3. 基于段級優勢值 token 概率掩碼策略優化(Policy Optimization Using Segment Advantages with Token Probability-mask):
在得到段級優勢值以后,為了進一步提高信用分配,團隊創新性地提出 token 概率掩碼策略優化方法,在策略更新僅將段級優勢值分配給該段內的低概率(關鍵)token,而非所有 token。這種方法能更精確地將獎勵 / 懲罰賦予關鍵的決策點,提升學習效率和效果。下面分別展示了 SPO-chain 和 SPO-tree 的優化目標。
a) SPO-chain 優化目標:
b) SPO-tree 優化目標:
對比基線方法
如下圖所示,在短思維鏈場景,使用 RhoMath1.1B 作為基座模型,使用 GSM8K 訓練集進行訓練,對比各種訓練算法,使用 SPO 訓練得到的模型測試集正確率更高。
對于長思維鏈場景,如下圖所示,使用 DeepSeek-R1-Distill-Qwen-1.5B 作為基座模型,使用 MATH 數據集進行訓練,在相同的訓練時間下,測試集正確率比 GRPO 更高。
下表展示了在長思維鏈場景下的更多對比結果:與同期基于相同基座模型(DeepSeek-R1-Distill-Qwen-1.5B)并使用 GRPO 方法訓練得到的模型(DeepScaleR、STILL-3)相比,盡管 SPO 僅使用 MATH 數據集且僅使用 4K 的最大上下文長度進行訓練,SPO-tree 在各個上下文長度評測下表現優秀。值得注意的是,盡管 DeepScaleR 在 32K 上下文長度評測下表現最佳,但它在較短上下文長度(2K 與 4K)下卻表現最差,甚至不及原始基座模型。這表明,GRPO 訓練方法可能未有效優化模型的 token 效率,導致輸出存在較多冗余,從而在上下文長度有限的情形下出現正確率下降的問題。
分段粒度的影響
通過實驗發現,很細的粒度 (int2,每個兩個切分點進行分段),相比于中等粒度 (int5),僅有微小提升,但是過粗的粒度 (int100),相比于中等粒度 (int5),正確率下降很大。證明了 SPO 采用中等粒度優勢值的有效性。
段劃分方式的影響
實驗表明,在短思維鏈場景下,采用提出的基于切分點的段劃分方式效果最好,優于采用換行符進行劃分(VinePPO)以及固定 token 數量劃分(Fixed-token-count)。
Token 概率掩碼消融
實驗表明,將 token 概率掩碼去除會導致 SPO-chain 正確率下降,更值得注意的是:將 token 概率掩碼應用到 GRPO 上,會讓其正確率有明顯上升。
不同樹結構的影響
實驗表明,更小的樹結構在早期正確率更高,可能因為更快掃過更多的數據樣本。然而隨著訓練的進行,更大的樹結構會有更好的正確率,因為更大的樹結構對于段級優勢值的估計更加準確。
總結
該工作提出了一種基于中間粒度段級優勢值的 RL 訓練框架 SPO,在 token 級和軌跡級之間更好的平衡,具有比軌跡級更好的信用分配,同時僅需要少量優勢值估計點,可以使用有效無偏的 MC 方式進行估計,不需要額外的 critic 模型。
文章同時提出了 SPO 的兩個實例,為短思維鏈場景設計的 SPO-chain 以及為長思維鏈場景設計的 SPO-tree,通過實驗證明了 SPO 框架和兩個實例的有效性。