s1-32B 模型：超越 o1-preview，一起探索其原因
原文鏈接：????https://arxiv.org/abs/2501.19393??
??代碼鏈接：??https://github.com/simplescaling/s1???
Test-time scaling是一種語言建模方法，它利用額外的測試時計算資源來提升性能，OpenAI的o1模型近期也證實了這一方法的有效性。本文旨在探尋實現Test-time scaling以及強大推理性能的最簡途徑。
主要貢獻
- 精心策劃小型數據集s1K：該數據集包含1000個問題，并配有推理過程，其構建基于難度、多樣性和質量這三個標準，且通過消融實驗進行了驗證。
- 開發預算強制策略：通過強行終止模型的思考過程，或在模型試圖結束時多次添加“Wait”來延長生成過程，以此控制測試時間計算。這一策略能促使模型復查答案，糾正推理過程中的錯誤步驟。
在s1K數據集上對Qwen2.532B-Instruct語言模型進行監督微調，并應用預算強制策略后，我們的模型s1-32B在競賽數學問題（MATH和AIME24）上的表現，比o1-preview高出27%。
推理數據管理以創建s1K
初始收集59K樣本
依據三個指導原則，從16個不同來源收集了最初的59,029個問題：
- 質量：數據集需保證高質量，通過檢查樣本，剔除格式不佳等有問題的數據集。
- 難點：數據集應具備挑戰性，解答問題需要大量的推理工作。
- 多元化：數據集應涵蓋不同領域，以覆蓋各類推理任務。
數據來源包括NuminaMATH、AIME問題、OlympicArena、AGIEval等現有數據集，同時還創建了兩個用于定量推理的新數據集：
- s1-prob：包含斯坦福大學統計系博士資格考試概率部分的182道題（??https://statistics.stanford.edu??），并附有涵蓋復雜證明過程的手寫答案。
- s1-teasers：由23個常用于量化交易職位面試的腦筋急轉彎組成。每個樣本包含問題及答案，均取自PuzzledQuant（??https://www.puzzledquant.com/??）。
針對每個問題，利用Google Gemini Flash Thinking API生成推理過程和解決方案，并提取其推理過程和響應，最終得到59K個由問題、生成的推理過程和生成的解決方案組成的三元組。
最終選擇1K樣本
- 質量：首先剔除在API調用過程中出現錯誤的問題，數據集樣本數量減少至54,116個。接著，通過檢查是否存在格式問題的字符串模式（如ASCII藝術圖、不存在的圖像引用或不一致的問題編號等），過濾掉低質量的示例，數據集進一步縮減至51,581個。從這部分數據中，挑選出384個被認為高質量且無需進一步篩選的樣本作為最終1000個樣本的一部分。
- 難點：在每個問題上評估Qwen2.5-7B-Instruct和Qwen2.5-32B-Instruct兩個模型，由Claude 3.5 Sonnet將模型的每次嘗試與參考解決方案對比，評估正確性。基于更難的問題需要更多思考標記的假設，使用Qwen2.5分詞器測量每個推理軌跡的標記長度來衡量問題難度。剔除Qwen2.5-7B-Instruct或Qwen2.5-32B-Instruct能正確解答的問題，此時樣本數量降至24,496個。
- 多元化：運用Claude 3.5 Sonnet，依據美國數學學會的數學學科分類（MSC）系統（如幾何、動態系統、實分析等），將每個問題分類到特定領域。從24,496個問題中選擇最終樣本時，先隨機均勻選擇一個領域，然后根據傾向于更長推理軌跡的分布從該領域抽取一個問題。重復此過程，直至收集到1000個樣本。經過這三個階段的篩選，最終得到的數據集涵蓋了50個不同領域。
測試時縮放
測試時縮放方法可分為：
- 順序計算：后續計算依賴于前面的計算結果（例如較長的推理軌跡）。
- 并行計算：計算過程相互獨立（例如多數投票）。
本文聚焦于順序擴展，因為作者直觀認為這種方式擴展性更佳，后續計算可基于中間結果進行，有助于更深入的推理和迭代優化。
預算強制
提出一種簡單的解碼時間干預方法，在測試時強制設定最大和/或最小數量的思維標記：
- 為強制設定最大令牌數，只需附加思維結束標記分隔符和“Final Answer:”，使模型提前退出思考階段，給出當前最佳答案。
- 為強制設定最小令牌數，禁止生成思維結束標記分隔符，并可選擇在模型當前推理軌跡后附加字符串“Wait”，鼓勵模型對當前生成結果進行反思。
以下圖為例，展示了該方法如何引導模型得出更好的答案。模型原本在“...is 2.”處試圖停止，但我們禁止了思維結束標記分隔符，改為附加“Wait”，促使s1-32B自行糾正答案。
基線
預算強制的基準對比方法包括：
- 條件長度控制方法：依靠在提示中告知模型應生成的長度，再按粒度分為：
Token條件控制：在提示中指定思維標記的上限。
步進條件控制：指定思考步驟的上限，每個步驟約100個標記。
類條件控制：編寫兩個通用提示，讓模型進行短期或長期思考。
- 抑制采樣：不斷采樣，直至生成結果符合預先設定的計算預算。該方法根據生成結果的長度獲取后驗響應。
指標
測量了三個指標：
結果
設置
使用s1K數據集對Qwen2.5-32B-Instruct進行監督微調，得到模型s1-32B。
性能
- 測試時縮放：下圖展示了s1-32B在預算強制策略下，隨著測試時間計算量增加的性能變化。從圖中可以看出，使用預算強制技術和更多測試時間計算，能提升AIME24的性能，但在六倍計算量時性能趨于平緩。通過多數投票在基礎模型上擴展測試時間計算，其性能無法趕上s1-32B，這驗證了順序擴展比并行擴展更有效的觀點。
- 樣品效率：下圖和表格將s1-32B與其他模型進行了對比。結果顯示，s1-32B是樣本效率最高的開放數據推理模型。盡管s1-32B僅在額外的1000個樣本上進行訓練，但其性能明顯優于基礎模型Qwen2.5-32B-Instruct。同時發布的r1-32B性能比s1-32B更強，不過它是在多800倍推理樣本上訓練得到的。
討論
進一步Test-time scaling的限制
結果表明，預算強制策略可通過外推測試時間計算提升性能，如將AIME24的性能從50%提高到57%。然而，在進一步擴展時存在兩個關鍵限制：一是性能最終會趨于平緩；二是底層語言模型的上下文窗口會對其形成約束。盡管存在這些限制，但測試時間擴展在廣泛的精度范圍內仍有效，部分原因是縮減測試時計算的行為具有可預測性，不受這些限制的影響。若要繼續進行Test-time scaling，需要能進一步推斷測試時計算的新方法。
并行擴展作為解決方案
并行擴展為順序擴展的限制提供了一種解決思路。我們通過兩種方法增強順序擴展模型：
- 多數表決：生成個解決方案后，選擇出現頻率最高的作為最終答案。
- 通過REBASE進行樹搜索：使用REBASE流程獎勵模型（從LLaMA-34B初始化，并在合成流程獎勵建模數據集上進一步微調），然后通過多數投票匯總REBASE生成的解決方案。
如下圖所示，在這種場景下，使用REBASE增強模型的擴展性優于多數投票，甚至比順序擴展效果更好。不過，REBASE在每個步驟都需要為獎勵模型進行額外的前向傳遞，會增加一定的計算開銷。
