GRPO能高效10倍嗎?Kwai AI的SRPO給出肯定答案
Kwai AI提出的SRPO框架將LLM強化學習後訓練步驟減少90%,同時在數學和代碼領域達到DeepSeek-R1水平。該兩階段強化學習方法結合歷史重採樣,克服了GRPO的侷限性。
OpenAI的o1系列和DeepSeek-R1的巨大成功明確展示了大規模強化學習(RL)在激發複雜推理行為和顯著提升大語言模型(LLM)能力方面的潛力。然而,這些突破性推理模型的核心訓練方法往往在其技術報告中保持神秘。最近的社區工作主要集中在數學推理上,跨域泛化的挑戰在很大程度上未被探索。此外,標準的GRPO算法存在性能瓶頸、樣本利用效率低以及混合領域數據集難以培養專門推理技能等問題,阻礙了RL方法在LLM上的有效擴展。
為解決這些限制,來自快手Kwaipilot團隊的研究人員引入了一種新穎的強化學習框架:兩階段歷史重採樣策略優化(SRPO)。該方法從多個維度系統性地應對上述訓練挑戰。團隊公開發布了技術報告,詳細闡述了其訓練方法的複雜性,並開源了SRPO-Qwen-32B模型。
值得注意的是,這項工作首次在數學和代碼領域同時實現了DeepSeek-R1-Zero級別的性能。通過使用與DeepSeek相同的基礎模型(Qwen2.5-32B),並採用純強化學習訓練方法,SRPO在AIME24(50)和LiveCodeBench(41.6)基準測試上取得了令人印象深刻的成績,超越了DeepSeek-R1-Zero-32B。更令人驚歎的是,SRPO僅用了R1-Zero所需訓練步驟的十分之一就達到了這一性能水平。
Vanilla GRPO面臨的挑戰
在初步探索中,Kwaipilot團隊嘗試了標準GRPO算法。然而,他們很快遇到了瓶頸,模型無法達到期望的R1-Zero性能水平。這些問題包括:
- 跨領域優化衝突(數學與代碼):數學問題傾向於產生更長、更詳細的推理軌跡(Long CoT),而代碼數據對此傾向較弱。直接混合這兩種數據類型導致了衝突,使得兩個領域的表現都不理想。
- 相似組獎勵導致訓練效率降低:GRPO算法依賴於採樣組內非零獎勵的方差來計算優勢。當組內的rollout產生幾乎相同的獎勵值時,計算出的優勢趨近於零。如果訓練批次中有很大一部分出現這種現象,有效的梯度貢獻就變得極小,大大降低了訓練效率。
- 過早的性能飽和:GRPO訓練在基準評估中遇到了早期性能平台和獎勵飽和。這個問題部分歸因於數據質量不足。當訓練數據缺乏足夠的複雜性或多樣性,尤其是包含過多簡單問題時,模型傾向於保守地維持其在簡單任務上的表現,阻礙了其開發複雜深入推理能力。
兩階段訓練
為了解決數學和代碼領域之間固有的響應長度衝突,Kwaipilot團隊實施了兩階段訓練範式:
- 第一階段:激發推理能力:這一初始訓練階段專注於高難度的數學數據。主要目標是充分激勵模型的測試時擴展,培養諸如反思性暫停、回溯和逐步分解等能力。
- 第二階段:技能整合:在這一階段,代碼數據被引入訓練過程。基於第一階段建立的推理基礎,此階段旨在進一步增強編碼能力,同時逐步強化程序思維、遞歸和工具調用能力。
不同訓練數據策略對響應長度的影響分析揭示了以下見解:混合訓練導致響應長度增長有限且基準性能不佳;純數學訓練帶來穩定的響應長度增長和優異的數學基準性能,並培養了強大且可泛化的推理能力;純代碼訓練在代碼基準上表現提升,但顯式推理行為發展不足;而兩階段訓練在數學和編程領域均取得了優異結果,模型能生成詳細的逐步推理,甚至自發利用代碼輔助數學推理。
歷史重採樣
Kwaipilot團隊觀察到,在訓練的中後期,批次中近50%的採樣組產生了相同的獎勵。這通常發生在模型在簡單問題上持續成功時,導致獎勵方差極小和梯度更新無效。為了解決這個問題並提高梯度信號質量,他們引入了歷史重採樣。在訓練過程中,他們記錄每個epoch內所有rollout的獎勵結果。在epoch結束時,他們根據以下標準重建下一個epoch的數據集:
- 過濾過於簡單的樣本:所有rollout都正確的樣本被排除,因為它們不能提供有信息量的信號用於策略改進。
- 保留信息量豐富的樣本:具有多樣化結果(既有正確也有錯誤)或全部錯誤的樣本被保留。這些樣本產生正的獎勵方差,確保非零優勢和有效的梯度信號。此外,當前epoch中所有rollout都錯誤的困難樣本也被保留。其理由是,這些最初具有挑戰性的問題對於更新後的策略可能變得相對容易,從而在後續訓練中產生有效的梯度。這種策略符合課程學習的原則,逐漸使模型平均面對更具挑戰性的樣本,以提高訓練效率。
與DAPO中提出的動態採樣方法相比,歷史重採樣顯著提高了計算效率,並帶來了更穩定的響應長度增長。
數據
Kwaipilot團隊對公開的Code&Math數據集進行了細緻的數據清洗和過濾。他們應用啓發式規則過濾掉無關URL、格式噪聲,並確保原始數據中核心字段(問題和答案真實值)的完整性。遵循PRIME對數學數據的數據清洗方法,他們移除了多部分問題、純證明題以及需要圖片或表格理解的問題。對於代碼數據,他們排除了依賴特定環境、文件I/O或網絡交互的問題,專注於算法邏輯。
在數據導入前,他們對數學和代碼問題進行了正確性驗證,確保答案的準確性和可解性,丟棄了不正確或模糊的解答。隨後,他們評估了每個問題的難度,根據通過率(Pass@k)將其分為簡單、中等和困難三個級別。
實驗結果
本節詳細介紹使用SRPO方法獲得的實驗結果。Kwaipilot團隊重點關注訓練過程中獎勵和響應長度等指標的變化。
在訓練過程中,初始獎勵增長進入平台期後,訓練進入第二階段。第二階段開始時,由於模型之前缺乏代碼訓練,整體獎勵下降,隨後在後續訓練中穩步上升。集成代碼數據並未顯著增加響應長度,這符合他們的預期。同時,基準測試結果表明,模型在數學和編碼能力上持續穩定提升,證明了新方法的有效性。
具體來説,歷史重採樣確保了每個訓練步驟的梯度更新保持有效,直接提高了信息量豐富的梯度比例。這種增強的採樣效率帶來了穩定的獎勵增長,清晰展示了重採樣策略帶來的訓練效率提升。
推理行為
Kwaipilot團隊識別了三種代表性的反思模式:重新檢查、猶豫和探索。他們對包含這些模式的響應進行了統計分析,並記錄了每種模式的平均響應長度。在RL訓練過程中,他們觀察到模型自我反思、糾正和回溯的頻率逐漸增加,表明“自我驗證”能力正在出現。他們認為,類似於人類認知過程的“反思”在RL訓練中出現在模型中,是策略優化過程產生的適應性行為。
訓練早期,模型幾乎沒有主動檢查和反思之前的推理步驟;但隨着訓練進行,模型表現出顯著的反思和回溯行為,形成了諸如逐步推理、數值代入、逐步驗證和自我優化等響應模式。有趣的是,他們還發現模型學會了在解決數學問題時自發使用程序代碼進行驗證。它會先通過數學推理提供解決過程,然後主動編寫代碼驗證答案的正確性。這些實例表明模型能夠利用程序思維進行自我糾正和多次嘗試,進一步表明在訓練後期,模型已經掌握了廣泛的思維和基於代碼的推理方法的綜合應用。
論文地址:arXiv
模型地址:HuggingFace