訓練1000樣本就能超越o1，李飛飛等人畫出AI擴展新曲線
今年 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
測試時間擴展
本文將測試時間擴展方法分為兩類：
- 序列擴展，即后續計算依賴于先前的計算結果；
- 并行擴展，即計算獨立運行。
本文專注于序列擴展，因為直觀上其具有更好的擴展性，因為后續計算可以基于中間結果進行，從而實現更深層次的推理和迭代優化。
此外，本文還提出了新的序列擴展方法以及對其進行基準測試的方式。
預算強制（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 中蒸餾出來的，這表明本文的蒸餾程序可能是有效的。
最后，本文還進行了一系列消融實驗，感興趣的讀者，可以查看原論文，了解更多內容。
