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