1/10訓練數據超越GPT-4o!清華等提出隱式過程獎勵模型PRIME,在線刷SOTA
Tell me and I forget, teach me and I remember, involve me and I learn.
告訴我,我會忘記,教我,我會記住,讓我參與,我就能學會。
——本杰明·富蘭克林
打破數據墻,我們還能做些什么?
近日,來自清華UIUC等機構的研究者提出了PRIME(Process Reinforcement through IMplicit REwards):通過隱式獎勵來進行過程強化。
GitHub地址:https://github.com/PRIME-RL/PRIME
這是一種帶有過程獎勵的在線RL開源解決方案,可以提高語言模型的推理能力,超越了SFT(監督微調)或者蒸餾等方法。
對比SFT,PRIME讓模型在重要基準測試上實現了巨大提升:平均提高了16.7%,在AMC和AIME中都提高了20%以上。
Eurus-2-7B-PRIME與Qwen2.5-Math-7B-Instruct,使用了相同的base model(Qwen-2.5-Math-7B),但在上表的6項測試中,5項都超越了instruct版本,同時也超越了GPT-4o。
而這個成績只用了Qwen Math 1/10的數據資源(230K SFT + 150K RL)!
作者發布了本研究中使用的所有模型和數據,感興趣的讀者請見文后鏈接。
過程獎勵模型
熱身階段(SFT)
如前所述,選擇Qwen2.5-Math-7B-Base作為起點,然后上點難度,采用競賽級別的數學和編程基準,包括AIME 2024、AMC、MATH-500、Minerva Math、OlympiadBench、LeetCode和LiveCodeBench(v2)。
首先對基礎模型進行監督微調,以獲得RL的入門模型(教模型學習某些推理模式)。
為此,研究人員設計了一個以動作為中心的鏈式推理框架,策略模型在每個步驟中選擇7個動作中的一個,并在執行每個動作后停止。
為了構建SFT數據集,研究者從幾個開源數據集中收集了推理指令。
值得注意的是,對于許多具有真實答案的數據集,作者選擇將其保留用于之后的RL訓練,目的是讓SFT和RL使用不同的數據集,以使RL中的探索多樣化,并且作者認為在PL中真實標簽更加重要。
作者用LLaMA-3.1-70B-Instruct來回答指令,并使用系統提示要求模型執行以動作為中心的思維鏈。
隱式PRM
下面接入過程獎勵模型(PRM),這里采用隱式PRM,只需要在響應級別標簽上訓練ORM。
過程獎勵模型簡單理解就是對每個推理步驟進行評分,舉個例子:
PRM是以這種粒度來評價響應的。
在本文的隱式PRM中,可以使用以下方式免費獲得過程獎勵:
通過簡單地收集響應水平數據和訓練ORM來獲得PRM,而無需注釋步驟標簽。
這與ORM訓練目標的具體選擇無關,比如使用交叉熵損失來實例化隱式PRM,就可以替換成:
強化學習
本文的目標是廣泛利用強化學習(RL)來提高推理能力。針對這種資源有限的情況,作者總結了一些最佳實踐:
從Ground Truth驗證器和高質量數據開始:作者進行了嚴格的數據收集和清理,以獲得可驗證的RL數據,并發現僅使用結果驗證器足以構建強大的基線。
作者比較了不同的RL算法得出結論,無價值模型的REINFORCE類方法足夠有效。
使用「mid-difficulty」問題進行穩定訓練:作者提出了一種名為在線提示過濾器的機制,通過過濾掉困難和簡單的問題,在很大程度上穩定了RL訓練。
使用PRM進行強化學習
將PRM集成到在線強化學習中并非易事,這里有幾個需要解決的關鍵挑戰。
如何為強化學習提供密集獎勵?
獎勵稀疏性一直是強化學習中長期存在的問題。到目前為止,我們仍然沒有特別好的解決方案來為LLM的在線強化學習構建密集獎勵。
以前的方法主要是為密集獎勵建立一個額外的價值模型,眾所周知,這樣的模型很難訓練,而且性能提升不大。
根據前文對隱式PRM的介紹,使用
可以免費從隱式PRM中獲得token級別的過程獎勵。
這種方式可以直接取代PPO中的價值模型,非常容易與任何優勢估計函數和結果獎勵相結合。在實踐中,作者將過程獎勵與REINFORCE、RLOO、GRPO、ReMax和PPO集成在一起,并進行了細微的修改。
如何設置一個好的PRM來啟動RL?
即使我們找到了在RL中使用過程獎勵的途徑,訓練好的PRM也并非易事:需要收集大規模(過程)獎勵數據(很貴),并且模型應該在泛化和分布偏移之間取得良好的平衡。
隱式PRM本質上是一種語言模型。因此從理論上講,可以使用任何語言模型作為PRM。在實踐中,作者發現最初的策略模型本身就是的一個很好的選擇。
如何在線更新PRM以防止獎勵黑客攻擊?
在線RL中,避免RM被過度優化或被黑客入侵至關重要,這需要RM與策略模型一起不斷更新。然而,鑒于步驟標簽的成本很高,在RL訓練期間很難更新PRM,——可擴展性和泛化問題。
但是,本文的隱式PRM僅要求更新結果標簽。也就是說,使用結果驗證器即可在訓練期間輕松更新PRM。
此外,還可以進行雙重轉發:首先使用策略部署更新PRM,然后使用更新的PRM重新計算過程獎勵,從而提供更準確的獎勵估算。
PRIME算法
下圖表示PRIME算法的整個循環:
策略模型和PRM都使用SFT模型進行初始化。對于每個RL迭代,策略模型首先生成輸出。然后,隱式PRM和結果驗證器對輸出進行評分,隱式PRM在輸出時通過結果獎勵進行更新。最后,將結果獎勵ro和過程獎勵rp組合在一起,用于更新策略模型。
以下是算法的偽代碼:
實驗
默認情況下,使用SFT模型初始化隱式PRM,并保留SFT模型作為參考對數探測器。超參數方面,策略模型的學習率固定為5e-7,PRM學習率為1e-6,使用AdamW優化器,mini batchsize大小為256,micro batchsize為8。
rollout階段收集256個提示,每個提示采樣4個響應。PRM訓練時β=0.05,所有實驗中將KL系數設置為0。
將PRIME與僅帶有結果驗證器(OV)的RLOO進行比較,與稀疏獎勵相比,PRIME將RL訓練加速了2.5倍,并將最終獎勵提高了6.9%,且方差更低。在下游任務上,PRIME的性能也始終優于OV。
下面展示PRM在線更新的重要性。比較兩種設置:在線PRM使用Eurus-2-7B-SFT初始化,離線PRM使用EurusPRM-Stage1初始化。
從下圖中可以看出,在線PRM在訓練集和測試集上的性能都大大優于離線PRM。