訓練1000樣本就能超越o1,李飛飛等人畫出AI擴充套件新曲線

机器之心發表於2025-02-05
跟大模型說:要多想。

今年 1 月,DeepSeek R1 引爆了全球科技界,它創新的方法,大幅簡化的算力需求撼動了英偉達萬億市值,更引發了全行業的反思。在通往 AGI(通用人工智慧)的路上,我們現在不必一味擴大算力規模,更高效的新方法帶來了更多的創新可能。

最近一段時間,全世界的科技公司、研究團隊都在嘗試復現 DeepSeek,但如果這個時候有人說「我還能大幅改進 AI 的推理效率」,你會怎麼想?
圖片
s1 論文作者,史丹佛大學在讀博士 Niklas Muennighoff 表示,DeepSeek r1 令人興奮,但其缺少 OpenAI 的測試時間擴充套件圖並且需要大量資料。我們推出的 s1 僅使用 1K 樣本和簡單的測試時間干預即可重現 o1 的預覽擴充套件和效能。

這個新方法叫 s1。本週,史丹佛大學、華盛頓大學等研究機構嘗試了最簡化實現測試時間擴充套件(test-time scaling)的方法,僅讓模型訓練 1000 個問題就獲得了超越 o1 的強推理效能。

測試時間擴充套件是一種有前途的語言建模新方法,它使用額外的測試時間計算來提高模型效能。此前,OpenAI 的 o1 模型展示了這種能力,但並未公開分享其方法。很多工作都在嘗試復現 o1,這些嘗試包含蒙特卡洛樹搜尋、多智慧體等等。今年 1 月開源的 DeepSeek R1 成功實現了 o1 級別的效能,它是在數百萬個樣本上透過多訓練階段強化學習實現的。

在 s1 的新工作中,研究人員尋求最簡單的方法來實現測試時間擴充套件。它們構建了一個小型資料集 s1K,其中包含 1000 個問題,並根據三個標準(難度、多樣性和質量)與推理軌跡進行配對。

在此基礎上,研究人員開發了「預算強制」來控制測試時間計算,方法是強制終止模型的思考過程,或者在模型試圖結束時多次將「等待」附加到模型的生成中以延長思考。這有可能會導致模型仔細檢查其答案,修復其不正確的推理步驟。

在 s1K 上對 Qwen2.5-32B-Instruct 語言模型進行監督微調(16 塊 H100 GPU,26 分鐘)併為其設定預算強制後,新模型 s1-32B 在競賽數學問題上的表現比 o1-preview 高出 27%(MATH 和 AIME24)。
圖片
s1 效能與其他大模型的對比。
圖片
  • 論文:《s1: Simple test-time scaling》

  • 論文連結:https://arxiv.org/abs/2501.19393

  • 專案連結:https://github.com/simplescaling/s1


測試時間擴充套件

本文將測試時間擴充套件方法分為兩類:
  1. 序列擴充套件,即後續計算依賴於先前的計算結果;

  2. 並行擴充套件,即計算獨立執行。

本文專注於序列擴充套件,因為直觀上其具有更好的擴充套件性,因為後續計算可以基於中間結果進行,從而實現更深層次的推理和迭代最佳化。

此外,本文還提出了新的序列擴充套件方法以及對其進行基準測試的方式。

預算強制(Budget forcing)。本文提出了一種簡單的解碼時間(decoding-time )干預方法,透過在測試時強制設定最大或最小思考 token 數量來實現。圖 3 為該方法的一個示例展示,說明了這種簡單的方法可以引導模型得出更好的答案。
圖片
具體來說,本文透過簡單地追加思考結束(end-of-thinking)token 分隔符和「Final Answer:」來強制設定最大 token 數量,從而提前退出思考階段,使模型提供其當前的最佳答案。為了強制設定最小 token 數量,本文抑制思考結束 token 分隔符的生成,並選擇性地在模型的當前推理軌跡後追加字串「Wait」,以鼓勵模型反思其當前生成的內容。

基線。本文用以下方法對預算強制進行基準測試:

(I)條件長度控制方法,該方法依賴於在提示中告訴模型它應該生成多長時間。本文按粒度將它們分組為(a)token 條件控制,在提示中指定思考 token 的上限;(b)步驟條件控制,指定思考步驟的上限;(c)類條件控制,編寫兩個通用提示,告訴模型思考一小段時間或很長一段時間。

(II)拒絕取樣,即取樣直到生成符合預定的計算預算。

實驗

在訓練階段。本文使用 s1K 資料集對 Qwen2.5-32B-Instruct 進行監督微調,以獲得本文的模型 s1-32B。微調是在 16 臺 NVIDIA H100 GPU 上使用 PyTorch FSDP 進行的,耗時 26 分鐘。

評估。本文采用了三個推理基準進行評估。
  • AIME24 包含 30 個問題,這些問題來自 2024 年 1 月 31 日至 2 月 1 日舉行的美國 AIME 數學競賽。AIME 用來測試模型在算術、代數、計數、幾何、數論、機率等領域的能力;

  • MATH500 是一個包含不同難度競賽數學問題的基準;

  • GPQA Diamond 包含 198 個來自生物學、化學和物理學的博士級科學問題。

其他模型。本文將 s1-32B 與以下模型進行基準測試對比:OpenAI o1 閉源系列模型;DeepSeek r1 開源模型;Qwen 的 QwQ-32B-preview 等模型。

值得一提的是,s1-32B 是完全開源的,包括權重、推理資料和程式碼。

效能

測試時間擴充套件。圖 1 展示了 s1-32B 在使用預算強制技術後,隨著測試時間計算資源的增加,效能的變化情況。
圖片
圖 4(左)擴充套件了圖 1(中)的圖表,結果顯示雖然本文可以透過預算強制技術和更多的測試時計算資源提升 AIME24 的效能,但最終在六倍計算量時趨於平緩。可以得出過於頻繁地抑制思考結束 token 分隔符可能會導致模型陷入迴圈重複,而不是持續推理。

圖 4(右)展示了在對 Qwen2.5-32B-Instruct 進行 1,000 個樣本的訓練,從而生成 s1-32B,併為其配備簡單的預算強制技術後,它進入了一種不同的擴充套件正規化。透過多數投票在基礎模型上擴充套件測試時間計算資源無法趕上 s1-32B 的效能,這驗證了這一直覺,即序列擴充套件比並行擴充套件更有效。
圖片
圖 5 提供了 s1-32B 的生成示例。
圖片
樣本效率。圖 2(右)和表 1 將 s1-32B 與其他模型進行了比較。

結果顯示, s1-32B 是樣本效率最高的開放資料推理模型。儘管只在額外的 1000 個樣本上進行訓練,但它的表現明顯優於基礎模型(Qwen2.5-32B-Instruct)。

r1-32B 在僅使用 SFT 的情況下表現出比 s1-32B 更好的效能,但前者是在 800 倍以上的推理樣本上進行訓練的。僅用 1000 個樣本是否能達到這個效能還是一個懸而未決的問題。

s1-32B 在 AIME24 上幾乎與 Gemini 2.0 Thinking 相匹配,因為 s1-32B 是從 Gemini 2.0 中蒸餾出來的,這表明本文的蒸餾程式可能是有效的。
圖片
圖片
最後,本文還進行了一系列消融實驗,感興趣的讀者,可以檢視原論文,瞭解更多內容。

相關文章