SimKO：緩解RLVR訓(xùn)練中的概率過(guò)度集中，優(yōu)化pass@K性能

2025-11-10 08:46:00

隨著 Deepseek-R1，Kimi1.5 等模型展示了強(qiáng)化學(xué)習(xí)在提升大型語(yǔ)言模型復(fù)雜推理能力上的巨大潛力，使用可驗(yàn)證強(qiáng)化學(xué)習(xí)（RLVR）在數(shù)學(xué)、邏輯與編程等領(lǐng)域進(jìn)行訓(xùn)練提升模型性能受到了廣泛關(guān)注。

作者彭若天是西湖大學(xué)和浙江大學(xué)聯(lián)培博士生，任毅是不列顛哥倫比亞大學(xué)博士，郁晝亮是香港中文大學(xué)博士生，劉威楊是香港中文大學(xué)計(jì)算機(jī)系助理教授，溫研東是西湖大學(xué)人工智能系助理教授

然而，盡管現(xiàn)有 RLVR 方法在提升模型的 pass@1 性能（單次嘗試正確的概率）方面取得了顯著成果，但其在 pass@K（K 次嘗試中至少一次正確的概率，K>1）上的性能相比基礎(chǔ)模型卻下降了。

這一現(xiàn)象表明，雖然模型在「利用」（Exploitation）單一正確路徑的能力有所增強(qiáng)，但犧牲了對(duì)多樣化正確解的「探索」（Exploration）能力。

針對(duì)這一問(wèn)題，研究團(tuán)隊(duì)從「模型預(yù)測(cè)下一個(gè)詞的概率分布」這一新視角出發(fā)，深入研究了「探索」能力下降的內(nèi)在機(jī)制。大量實(shí)驗(yàn)發(fā)現(xiàn)，經(jīng)過(guò)現(xiàn)有 RLVR 算法訓(xùn)練后的模型，多數(shù)存在概率集中于單一推理路徑的問(wèn)題。

受該現(xiàn)象啟發(fā)，研究團(tuán)隊(duì)提出一種簡(jiǎn)潔且高效的算法 ——SimKO (Simple Pass@K Optimization)，顯著優(yōu)化了 pass@K（K=1 及 K>1）性能。

圖 1

論文地址：https://arxiv.org/abs/2510.14807
項(xiàng)目主頁(yè)：https://spherelab.ai/simko/
代碼倉(cāng)庫(kù)：https://github.com/CLR-Lab/SimKO

問(wèn)題分析：導(dǎo)致 RLVR 探索能力下降的原因是什么？

當(dāng)前主流的大語(yǔ)言模型強(qiáng)化學(xué)習(xí)算法（如 GRPO、DAPO 等）采用 RLVR 范式，其訓(xùn)練方法可以直觀理解為：模型對(duì)每個(gè)問(wèn)題生成多個(gè)答案，對(duì)正確答案給予獎(jiǎng)勵(lì)，對(duì)錯(cuò)誤答案施加懲罰。在理想的探索狀態(tài)下，模型能夠在多個(gè)潛在正確推理路徑之間分配相對(duì)均勻的概率質(zhì)量，不應(yīng)將概率過(guò)度集中于某一條正確路徑上，如圖 2 (a) 和 (b) 所示。

同時(shí)，團(tuán)隊(duì)認(rèn)為當(dāng)前的用熵（Entropy）作為指標(biāo)衡量多樣性存在局限：熵?zé)o法具體反映概率分布的形態(tài)。如圖 2（c）所示，兩個(gè)具有相同熵值的分布，一個(gè)可能包含多個(gè)峰值，而另一個(gè)則可能高度集中于一個(gè)峰值。

因此，熵?zé)o法精確描述模型在推理過(guò)程中對(duì)不同推理路徑的真實(shí)探索程度。

圖 2

因此，團(tuán)隊(duì)引入了一種新的分析指標(biāo)，用于更加精細(xì)地觀察訓(xùn)練過(guò)程中的學(xué)習(xí)動(dòng)態(tài)：該指標(biāo)通過(guò)測(cè)量模型在生成過(guò)程中，反映下一 token 后驗(yàn)概率分布的不同排名候選詞（rank-k candidate）的平均對(duì)數(shù)概率（average log-probability），從而實(shí)現(xiàn)對(duì)概率分布演化動(dòng)態(tài)的直接觀測(cè)。

通過(guò)這一新指標(biāo)，作者發(fā)現(xiàn) RLVR 訓(xùn)練機(jī)制存在一個(gè)系統(tǒng)性偏差（如圖 3 所示）：

它會(huì)持續(xù)強(qiáng)化 rank-1 候選詞的概率，同時(shí)顯著抑制其他較低排名（rank-k, k>1）的候選路徑，即使那些路徑同樣是正確的。

這種機(jī)制導(dǎo)致了模型輸出分布的「過(guò)度集中」。模型的概率質(zhì)量過(guò)度匯聚于單一的推理路徑，導(dǎo)致其喪失了生成多樣化正確答案的能力。

更重要的是，進(jìn)一步實(shí)驗(yàn)分析明確揭示了概率過(guò)度集中問(wèn)題與 pass@K 性能的下降之間存在強(qiáng)相關(guān)性：當(dāng)模型概率分布越集中于 rank-1 答案，而 rank-2 和 rank-3 的概率越低的時(shí)候，其 pass@K 指標(biāo)也隨之降低。

圖 3

SimKO：一種非對(duì)稱(chēng)的梯度調(diào)節(jié)策略

圖 4

為解決上述的概率過(guò)度集中問(wèn)題，研究團(tuán)隊(duì)提出了 SimKO (Simple Pass@K Optimization)。其核心機(jī)制在于對(duì)探索 token 施加非對(duì)稱(chēng)的更新策略（如圖 4 所示）：即在正確的推理路徑上實(shí)現(xiàn)概率平滑，而在錯(cuò)誤的推理路徑上施加精準(zhǔn)懲罰。

(A) 關(guān)鍵節(jié)點(diǎn)的識(shí)別

SimKO 并非對(duì)所有 token 進(jìn)行無(wú)差別調(diào)節(jié)。它首先識(shí)別推理路徑中具有高熵的 token，這些 token 代表了模型面臨多個(gè)高概率選項(xiàng)、可能產(chǎn)生不同推理方向的關(guān)鍵節(jié)點(diǎn)。因此 SimKO 更新策略只應(yīng)用于這些關(guān)鍵節(jié)點(diǎn)。

(B) 正確路徑：實(shí)施 top-K Label Smoothing

傳統(tǒng)方法 (GRPO): 僅提升被采樣到的單個(gè)正確 token 的概率。由于推理路徑是由模型自行采樣的，因此采樣到的的 token 是 rank-1 候選 token 的概率較高高，這是導(dǎo)致概率過(guò)度集中于 rank-1 的直接原因。這種機(jī)制容易使概率質(zhì)量集中在單路徑，從而抑制了對(duì)其他潛在正確路徑的探索。
SimKO ：對(duì)正確的 token 使用 top-K label smoothing 策略。當(dāng)一個(gè)采樣路徑被驗(yàn)證為正確時(shí)，SimKO 將獎(jiǎng)勵(lì)均勻分配給路徑上的關(guān)鍵節(jié)點(diǎn)處概率最高的 top-K 個(gè)候選 token。該策略防止概率分布集中在單一的正確路徑上，從而在強(qiáng)化正確路徑的同時(shí)，維持了模型的探索多樣性。

傳統(tǒng)方法 (GRPO): 對(duì)所有被采樣到的錯(cuò)誤 token 施加統(tǒng)一強(qiáng)度的懲罰。這種統(tǒng)一懲罰常導(dǎo)致「擠壓效應(yīng)」（Squeezing Effect）。即當(dāng)一個(gè)非 rank-1 的錯(cuò)誤候選 token 被懲罰時(shí)，這部分概率會(huì)被 rank-1 候選 token 吸收，從而提高 rank-1 token 的概率，這會(huì)讓模型的概率分布更加集中。
SimKO ：對(duì)錯(cuò)誤的 token 進(jìn)行非對(duì)稱(chēng)懲罰。當(dāng)采樣的錯(cuò)誤 token 恰好是 rank-1 候選時(shí)，施加顯著更強(qiáng)的懲罰；而對(duì)其他（rank-k, k>1）錯(cuò)誤候選則降低懲罰強(qiáng)度。這種策略緩解了懲罰非 rank-1 token 帶來(lái)的概率集中問(wèn)題，同時(shí)對(duì) rank-1 token 實(shí)施更強(qiáng)的懲罰也防止模型陷入單一的錯(cuò)誤推理模式。

實(shí)驗(yàn)分析：實(shí)現(xiàn)「探索」與「利用」的平衡

團(tuán)隊(duì)在多個(gè)數(shù)學(xué)推理基準(zhǔn)（MATH500、AIME 2024/25、Minerva Math、OlympiadBench、AMC23）上對(duì) Llama 和 Qwen 模型進(jìn)行了系統(tǒng)性評(píng)估。如表 1 所示，SimKO 策略在顯著提升 pass@K 性能的同時(shí)，成功保持（或略微提升）了 pass@1 準(zhǔn)確率，證明其有效平衡了「探索」與「利用」。

表 1

為了驗(yàn)證 SimKO 是否有效緩解了概率分布過(guò)度集中問(wèn)題，研究團(tuán)隊(duì)使用上述新提出的分析指標(biāo)，追蹤了模型使用不同 RLVR 算法訓(xùn)練過(guò)程中的學(xué)習(xí)動(dòng)態(tài)（如圖 5 所示）：

GRPO：隨著訓(xùn)練的進(jìn)行，模型表現(xiàn)出嚴(yán)重的概率過(guò)度集中現(xiàn)象。rank-1 的概率迅速收斂至接近 1，而 rank-2 與 rank-3 的概率則降低至 10?? 與 10?1?以下。這表明，幾乎所有的概率質(zhì)量都集中到了 rank-1 token 上，導(dǎo)致模型訓(xùn)練后喪失了探索能力。
SimKO：相比之下，SimKO 展示出了有效的緩解概率集中效果。其 rank-1 的概率顯著低于 GRPO，同時(shí)為 rank-2 與 rank-3 保留了更高的概率質(zhì)量。這一結(jié)果有力地表明了 SimKO 通過(guò)緩解了概率過(guò)度集中問(wèn)題，從而提升了模型的探索能力，進(jìn)而提升 pass@K 性能

圖 5

SimKO 不僅在數(shù)學(xué)推理任務(wù)上表現(xiàn)優(yōu)異，在邏輯推理任務(wù)中同樣具有出卓越的泛化效果（見(jiàn)表 2）：

Synlogic (分布內(nèi)任務(wù)): SimKO 相比 GRPO 表現(xiàn)出更好的性能，pass@1 提升 31.6%，pass@128 提升 26.3%。相比之下，NSR 和 W-REINFORCE 在此任務(wù)上未能有效收斂（pass@1 僅 1.1% 和 0.8%）。
BBH (分布外任務(wù)): SimKO 在提升 pass@1（+16%）的同時(shí)，pass@128 依然有良好的性能（92%）表現(xiàn)。而對(duì)照組 GRPO（88.2%）和 PSR （82.8%）則在 pass@128 相比基礎(chǔ)模型出現(xiàn)了顯著的性能衰減。