16張H100訓26分鐘,超越o1-preview,李飛飛等用1K樣本,揭秘測試時Scaling
大模型推理性能的提升,真的只能靠堆數據、加算力嗎?李飛飛等用僅1000個樣本微調模型,並提出預算強製(budget forcing)技術,成功讓推理能力隨測試計算量增加而提升。他們的s1-32B模型在多個基準測試中超越閉源模型OpenAI o1-preview,成為目前最具樣本效率的推理模型。
OpenAI o系列模型為何性能如此強大?
OpenAI將他們的方法描述為使用大規模強化學習(RL),暗示使用了大量的數據。
最近大火的DeepSeek-R1模型也通過使用數百萬個樣本和多個訓練階段使用強化學習的方式,成功地達到了o1級別的性能。
然而,至今為止沒有人公開成功複現清晰的測試時擴展行為。
那麼問題來了,實現測試時擴展和強推理性能的最簡單方法是什麼?
近日,來自史丹福大學、華盛頓大學、Ai2等機構的研究人員發表了一篇題為「s1: Simple test-time scaling」的論文,回答了上述問題。

團隊證明,僅使用1000個樣本進行下一個token的預測訓練,並在測試時通過一種簡單的預算強製(budget forcing)技術來控制思維持續時間,就能獲得一個強大的推理模型,其性能隨著測試計算量的增加而提升。
預算強製(budget forcing)可以簡單理解為通過強製提前結束模型的思考過程,或通過重覆添加「Wait」來延長思考時間,從而影響模型的推理深度和最終答案。
這種方法可以引導模型進行自我檢查,並修正推理過程中的錯誤,從而提高推理性能。
具體來說,他們構建了一個叫做「s1K」的數據集,由1000個精心篩選的問題組成,每個問題都配有推理軌跡(reasoning traces)和從Gemini Thinking Experimental蒸餾而來的答案。
接著團隊在一個預訓練模型上進行監督微調(SFT),僅使用16張H100 GPU訓練26分鐘。
訓練完成後,使用預算強製(budget forcing)方法來控制模型在測試時的計算量:
若模型生成的推理token超過設定的上限,則強製結束推理過程,並附加思維結束(end-of-thinking)token,促使模型進入答案生成階段。
若希望模型在問題上投入更多測試時計算資源,則抑制思維結束token的生成,並在推理軌跡中追加 「Wait」,鼓勵模型進行更深入的推理探索。
基於這個簡單的方法,並在1000個樣本上進行SFT訓練 + 測試時的預算強製(budget forcing)後,團隊提出的s1-32B展現出了測試時擴展(test-time scaling)的能力。

此外,s1-32B也是目前最具樣本效率(sample-efficient)的推理模型,在推理能力上超越了OpenAI的o1-preview等閉源模型。

如何創建s1K數據集
s1K數據集是一個包含1000個高質量推理問題的精選數據集。
團隊創建過程主要分為兩個階段。
初始階段,研究人員從16個不同的來源收集了59029個問題,並遵循三個指導原則:質量、難度和多樣性。
這些來源包括現有的數學問題數據集(如 NuminaMATH、AIME、OmniMath 和 AGIEval),以及研究人員自己創建的概率問題集 (s1-prob) 和腦筋急轉彎問題集 (s1-teasers)。
為了確保質量,研究人員檢查了所有樣本,並忽略了格式不佳的數據集。為了增加難度,他們選擇需要大量推理努力的問題。為了確保多樣性,他們涵蓋了不同的領域和推理任務.
第二階段,最終篩選1K樣本。研究人員通過三個階段的過濾,從59K樣本中篩選出1000個樣本,並繼續依賴質量、難度和多樣性這三個原則。
通過這些過程,研究人員創建了s1K數據集,該數據集包含50個不同領域的1000個高質量、多樣化和高難度的問題,並附帶推理過程。
這個數據集對於訓練s1-32B模型至關重要。
測試時擴展方法
其核心思想是通過在測試時增加計算量來提高語言模型的性能。
論文將測試時擴展方法分為兩類:順序(Sequential)和並行(Parallel)。
順序擴展是指後面的計算依賴於前面的計算,例如長的推理過程;並行擴展是指計算是獨立運行的,例如多數投票。
論文主要關注順序擴展,因為作者認為它可以更好地利用中間結果進行更深入的推理和迭代改進。
預算強製(Budget Forcing):通過限制模型在測試時使用的最大和/或最小思考token數量來控制計算量。
論文通過實驗證明了,這種簡單的方法能夠引導模型修正答案。
下圖這個例子中,模型最初在回答一個關於「raspberry」中 「r」的數量的問題時給出了錯誤的答案 「2」。
然而,通過抑制結束思考的token生成,並追加「Wait」來強製模型繼續推理,模型最終意識到自己快速閱讀導致了錯誤,並最終給出了正確的答案「3」。

測試時擴展方法
如下圖所示,s1-32B模型在使用預算強製技術後,其性能會隨著測試時計算量的增加而提高。

具體來說,通過增加模型思考的token數量(例如,通過追加「Wait」),模型在 AIME24 基準測試上的表現得到了提升。
然而,這種提升最終會趨於平緩,過度抑制結束思考的token會導致模型進入重覆循環。
結果表明,s1-32B模型是目前樣本效率最高的開源推理模型。儘管只使用了1000個樣本進行微調,s1-32B的性能仍明顯優於其基礎模型Qwen2.5-32B-Instruct。

同時,雖然DeepSeek r1-32B模型性能更強,但其使用了800倍的訓練樣本。
此外,s1-32B模型在AIME24上的表現幾乎與Gemini 2.0 Thinking API持平,表明其蒸餾過程是有效的。

總之,實驗結果證明了s1-32B模型在測試時擴展、樣本效率和推理能力方面的優勢,並驗證了預算強製技術的有效性。
消融實驗
數據消融實驗:研究人員通過以下對比實驗,驗證了高質量、多樣性、和難度這三個數據選擇標準的重要性:
僅質量 (1K-random):隨機選取1000個高質量樣本,性能明顯低於s1K,表明難度和多樣性過濾的重要性。
僅多樣性 (1K-diverse):均勻選取各個領域樣本,性能也遠不如s1K,表明只關注多樣性是不夠的。
僅難度 (1K-longest):選擇推理軌跡最長的1000個樣本,在GPQA上有提升,但整體不如s1K,表明難度只是一個方面。
最大化數據量 (59K-full):使用所有59K樣本訓練,雖然性能略有提升,但訓練資源消耗巨大,且提升幅度有限,說明精心挑選的少量數據比大量數據更高效。
結果表明,將質量、難度和多樣性相結合是實現樣本高效推理訓練的關鍵。

測試時擴展方法消融實驗:研究人員通過比較不同的測試時擴展方法,驗證了預算強製的優越性:
Token/步驟/類別 條件控制 (TCC/SCC/CCC):這些方法都無法有效控制計算量或獲得良好的擴展效果,表明僅在提示中告知模型計算量或步驟是不足的。
拒絕采樣 (RS):使用拒絕采樣會導致性能隨著計算量的增加而下降,因為更短的生成往往是模型一開始就走在正確軌道上的結果。
預算強製 (BF):實驗表明,預算強製在控制性、擴展性和性能方面都優於其他方法。
在預算強製中,追加「Wait」能夠鼓勵模型進行額外的思考,從而提高性能。

下圖展示了在AIME24數據集上使用s1-32B模型進行拒絕采樣(rejection sampling)的實驗結果,結果表明,隨著平均思考時間(以token數量衡量)的增加,模型的準確率反而下降,呈現出反向擴展趨勢。

更長的推理過程並不一定意味著更好的性能。 這一結果也反襯了論文提出的預算強製方法(budget forcing)的優越性,因為預算強製可以更有效地控制測試時的計算量,並促使模型進行更有目的性的思考。
實驗結果表明,預算強製是測試時擴展的最佳方法。
總結
儘管很多模型,例如DeepSeek-r1和k1.5,通過強化學習或使用數萬個蒸餾樣本來構建強大的推理模型。
但該研究表明,僅需在1000個樣本上進行監督微調 (SFT) 就足以構建一個具有競爭力的推理模型,該模型可以與 OpenAI的o1-preview相媲美。
李飛飛團隊研究人員認為,預訓練階段模型已經接觸了大量的推理數據,因此微調階段只需少量樣本就能激活模型的推理能力,這與LIMA論文中提出的表面「對齊假說」類似。
論文提出的預算強製 (budget forcing) 是一種簡單有效的順序擴展方法,它通過控制模型思考的token數量來提高性能,並首次複現了OpenAI的測試時擴展曲線。
儘管預算強製有其局限性,例如最終會趨於平緩和受上下文窗口的限制,但它證明了測試時擴展的潛力,並為未來的研究提供了明確的指標:可控性、擴展性和性能。
為了克服順序擴展的局限性,論文還探討了並行擴展方法,如多數投票和基於REBASE的樹搜索(下圖)。實驗結果表明,將順序擴展與並行擴展方法相結合,還可以進一步擴展測試時的計算量。

未來方向
論文強調,未來的研究可以探索如何進一步改進預算強製,例如輪換使用不同的字符串或結合頻率懲罰。
一個有前景的方向是將預算強製應用於通過強化學習訓練的推理模型,並研究新的測試時擴展方法。
此外,可以研究如何進一步擴展測試時的計算量,以克服現有語言模型上下文窗口的限制。
參考資料:
https://arxiv.org/pdf/2501.19393
本文來自微信公眾號「新智元」,作者:新智元,36氪經授權發佈。