4500美元復刻DeepSeek神話,1.5B戰勝o1-preview只用RL!訓練細節全公開
強化學習迎來重大突破!
近日,來自UC伯克利的研究團隊基于Deepseek-R1-Distilled-Qwen-1.5B,通過簡單的強化學習(RL)微調,得到了全新的DeepScaleR-1.5B-Preview。
在AIME2024基準中,模型的Pass@1準確率達高達43.1% ——不僅比基礎模型提高了14.3%,而且在只有1.5B參數的情況下超越了OpenAI o1-preview!
目前,研究團隊已開源數據集、代碼和訓練日志。
只用不到5000美元的預算,團隊就復現了DeepSeek的成功。至此,開源又贏下一局。
網友們稱贊:當機器學習和數學相遇,就是超強組合的誕生!
訓練秘籍簡版:先短后長
1.5B模型,通過RL訓練,就能超越o1-preview,進行數學推理?
簡而言之,團隊這次的訓練策略就是四個字——先短后長。
第一步,研究人員會訓練模來型進行短思考。他們使用DeepSeek的GRPO方法,設定了8k的上下文長度來訓練模型,以鼓勵高效思考。
經過1000步訓練后,模型的token使用量減少了3倍,并比基礎模型提升了5%。
接下來,模型被訓練進行長思考。強化學習訓練擴展到16K和24K token,以解決更具挑戰性、以前未解決的問題。
隨著響應長度增加,平均獎勵也隨之提高,24K的魔力,就讓模型最終超越了o1-preview!
DeepScaleR-1.5B-Preview
最近,Deepseek-R1開源發布,對推理模型技術普及來說,是個重要突破。不過,它具體的訓練方法、超參數還有底層系統,都還沒公開。
在擴展強化學習的時候,最大的難題之一就是計算成本太高。
就拿DeepSeek-R1的實驗來說,要想完全復現,上下文長度得達到32K以上,訓練大概8000步,就算是只有1.5B參數的模型,起碼都得花70,000 GPU小時。
如何利用強化學習,把小型模型變成超厲害的推理模型呢?
為了解決這個問題,研究人員用了知識蒸餾模型,還創新性地引入了強化學習迭代延長方法。
團隊推出了DeepScaleR-1.5B-Preview模型,它經過4萬個高質量數學問題的訓練,訓練一共用了3800個A100 GPU小時。
最終,成本只需約4500美元,省了18.42倍!同時模型的性能還在幾個競賽級數學基準中,超過了o1-preview。
研究表明,用強化學習開發定制化的推理模型,既能大規模進行,還能控制成本,性價比超高!
AIME 2024測試集Pass@1準確率隨訓練進度而變:訓練至第1040步,上下文長度擴至16K token;到第1520步,上下文長度增至24K token
技術方案
數據集構建
在訓練數據集方面,研究人員收集了1984至2023年的美國國際數學邀請賽(AIME)、2023年之前的美國數學競賽(AMC),以及來自Omni-MATH和Still數據集的各國及國際數學競賽題目。
數據處理流程包含三個核心步驟:
- 答案提取:對于AMC和AIME等數據集,使用gemini-1.5-pro-002模型從AoPS官方解答中提取答案。
- 重復問題清理:基于RAG,并結合sentence-transformers/all-MiniLM-L6-v2的詞向量嵌入來消除重復問題。同時,對訓練集和測試集進行重疊檢測,以防止數據污染。
- 不可評分題目過濾:數據集(如Omni-MATH)中的部分問題,無法通過sympy數學符號計算庫評估(得靠LLM判斷)。這不僅會降低訓練速度,還會引入不穩定的獎勵信號,因此需要增加額外的過濾步驟,來剔除無法自動評分的問題。
在經過去重和過濾之后,就得到了約4萬個獨特的問題-答案對作為訓練數據集。
獎勵函數設計
按Deepseek-R1的經驗,用結果獎勵模型(ORM)而不是過程獎勵模型(PRM),來避免模型通過投機取巧得到獎勵。
獎勵函數返回值如下:
- 返回「1」:如果LLM的答案,既能通過LaTeX語法檢查,又能通過Sympy數學驗證,就給它獎勵。
- 返回「0」:要是LLM的答案是錯的,或者格式不對,比如少了<think>和</think>標記,那就不給獎勵。
迭代增加上下文長度:從短到長的思維擴展
推理任務由于會生成比標準任務更長的輸出,計算開銷較大,這會同時降低軌跡采樣(Trajectory Sampling)和策略梯度(Policy Gradient)更新的速度。
與此同時,上下文窗口大小翻倍,則會導致訓練計算量至少增加2倍。
這種情況產生了一個根本性的權衡取舍:較長的上下文能為模型提供更充足的思維空間,但會顯著降低訓練速度;而較短的上下文雖然可以加快訓練進度,但可能會限制模型解決那些需要長上下文的復雜問題的能力。
因此,在計算效率和準確性之間找到最佳平衡點至關重要。
基于Deepseek的廣義近端策略優化(GRPO)算法的訓練方案包含兩個主要步驟:
- 首先,使用8K token的最大上下文長度進行強化學習訓練,從而實現更有效的推理能力和訓練效率。
- 隨后,將上下文長度擴展到16K和24K token,使模型能夠解決更具挑戰性的、此前未能攻克的問題。
用8K上下文構建高效思維鏈推理
正式訓練之前,先用AIME2024測試集對Deepseek-R1-Distilled-Qwen-1.5B模型進行評估,并分析它的推理軌跡數據。結果發現,錯誤答案里平均包含的token數量,是正確答案的三倍。這說明回答越長,越容易出錯。
因此,直接采用長上下文窗口進行訓練效率可能不高,因為大部分token都沒有被有效利用。此外,冗長的回答還會表現出重復性模式,這表明它們并未對對思維鏈推理(CoT)產生實質性的貢獻。
基于這些發現,團隊決定先從8K token的上下文長度開始訓練。在AIME2024測試里,獲得了22.9%的初始準確率,只比原始模型低6%。
事實證明這個策略很有效:訓練的時候,平均訓練獎勵從46%提高到了58%,平均響應長度從5500 token減少到了3500 token。
把輸出限制在8K token以內,模型能更高效地利用上下文空間。如下表所示,不管是生成正確答案還是錯誤答案,token數量都大幅減少了。
在AIME準確率上,比原始基準模型還高了5%,用的token數量卻只有原來的1/3左右。
擴展至16K token上下文,關鍵轉折點出現
在大約1000步后,8K token運行中發生了一個有趣的變化:響應長度再次開始增加。然而,這卻沒有增加收益——輸出準確率達到了平臺期,并最終開始下降。
與此同時,響應截斷比例從4.2%上升到了6.5%,這表明更多的響應在上下文長度的限制下被截斷。
這些結果表明,模型試圖通過「延長思考時間」來提高訓練獎勵。然而,隨著更長的輸出,模型越來越頻繁地觸及到8K token上下文窗口的上限,從而限制了性能的進一步提升。
研究人員意識到這是一個自然的過渡點,于是決定「放開籠子,讓鳥兒飛翔」。
他們選擇了在第1040步的檢查點——即響應長度開始上升的地方——重新啟動訓練,并使用了16K上下文窗口。
這種兩階段的做法比從一開始就用16K token訓練效率高得多:8K的預熱階段讓平均響應長度保持在3K token而不是9K,這使得此階段的訓練速度至少提高了2倍。
在擴展上了下文窗口后,研究人員觀察到訓練獎勵、輸出長度和AIME準確率都呈現穩定提升趨勢。經過額外的500步訓練,平均輸出長度從3.5K增加至5.5K token,AIME2024的Pass@1準確率達到了38%。
24K魔法,超越o1-preview
在16K token上下文環境下額外訓練500步后,研究人員發現模型性能開始趨于平穩——平均訓練獎勵收斂在62.5%,AIME單次通過準確率徘徊在38%左右,輸出長度再次呈現下降趨勢。同時,最大輸出截斷比率逐漸升至2%。
為了最終推動模型性能達到o1級別,研究人員決定決定推出「24K魔法」——將上下文窗口擴大到24K token。
首先,將16K訓練時的檢查點設定在第480步,并重新啟動了一個24K上下文窗口的訓練。
隨著上下文窗口的擴展,模型終于突破了瓶頸。在大約50步后,模型的AIME準確率首次超過了40%,并在第200步時達到了43%。24K的魔力發揮得淋漓盡致!
總體來看,訓練歷時約1750步。最初的8K階段使用了8塊A100 GPU進行訓練,而16K和24K階段則擴展到32塊A100 GPU進行訓練。
整個訓練過程共耗時約3800個A100小時,相當于32塊A100 GPU上運行了大約5天,計算成本約為4500美元。
研究人員用多個競賽級別的數學評測基準來測試模型,像AIME 2024、AMC 2023、MATH-500、Minerva Math還有OlympiadBench。
這里報告的是Pass@1準確率,簡單說,就是模型第一次就答對的概率。每個問題的結果,都是16次測試取平均值得到的。
將DeepScaleR和DeepSeek模型,以及近期專注推理任務強化學習的成果對比。DeepScaleR在所有評測里,都比基礎模型強很多。
在AIME 2024測試中,成績更是大幅提升了14.4%,整體性能也提高了8.1%。
DeepScaleR比最新模型的表現還好,像從7B參數模型微調來的rSTAR、Prime和SimpleRL。DeepScaleR只用1.5B參數,就達到了o1-preview的性能水平——這是模型效率的重大突破!
AIME準確率與模型規模對比,DeepScaleR實現性能與規模最佳平衡(帕累托最優)。
關鍵發現
很多人認為強化學習只對大型模型有用,其實強化學習在小型模型上也能發揮顯著作用。
Deepseek-R1發現,直接在小型模型上用強化學習,效果不如知識蒸餾。在Qwen-32B模型上做對比實驗,強化學習只能讓AIME測試的準確率達到47%,但只用知識蒸餾就能達到72.6%。
不過,要是從更大的模型中,通過蒸餾得到高質量的SFT數據,再用強化學習,小模型的推理能力也能大幅提升。
研究證明了這一點:通過強化學習,小型模型在AIME測試中的準確率從28.9%提高到了43.1%。
不管是只用監督微調,還是只用強化學習,都沒辦法讓模型達到最佳效果。只有把高質量的監督微調蒸餾和強化學習結合起來,才能真正發揮LLM的推理潛力。
之前的研究發現,強化學習直接在16K token的上下文環境里訓練,和8K token比起來,效果并沒有明顯提升。這很可能是因為計算資源不夠,模型沒辦法充分利用擴大后的上下文。
最近的研究也指出,模型回復太長,里面就會有很多冗余的推理內容,這些內容容易導致錯誤結果。本文的實驗證實了這些發現。
團隊先在較短的8K token上下文里,優化模型的推理能力,這樣一來,后續在16K和24K token的環境里訓練時,就能取得更快、更明顯的進步。
這種一步一步增加長度的方法,能讓模型在擴展到更長的上下文之前,先建立起穩定的推理模式,從而提高強化學習擴展上下文長度的效率 。
核心貢獻者
項目主頁還展示了參與DeepScaleR設計的所有研究人員,其中有兩位核心貢獻者。
Michael Luo
Michael Luo目前是UC伯克利電氣工程與計算機科學系(EECS)的博士生,導師是Ion Stoica教授。
在此之前,他獲得了UC伯克利電氣工程與計算機科學碩士和工商管理雙學士學位。
他的研究興趣主要在人工智能和系統領域。目前,其研究主要是為機器學習從業者構建可擴展的系統,以實現Sky Computing的愿景。
Sijun Tan(譚嗣俊)
譚嗣俊目前是UC伯克利計算機科學專業的三年級博士生,導師是Raluca Ada Popa。
此前,他在弗吉尼亞大學獲得計算機科學和數學雙學士學位,導師是David Wu和Yuan Tian。
他曾在Facebook AI Research(FAIR)實習過一段時間,并在螞蟻集團擔任過高級算法工程師。
他的研究領域涵蓋機器學習、計算機安全和應用密碼學。目前,其研究重點是增強通用型AI智能體的能力和魯棒性。