S1：簡單高效的測試時推理能力擴(kuò)展方法
引言：測試時計算擴(kuò)展的興起與挑戰(zhàn)
近年來，大型語言模型（LLMs）在自然語言處理領(lǐng)域取得了舉世矚目的成就，其性能的飛躍在很大程度上得益于大規(guī)模自監(jiān)督預(yù)訓(xùn)練和不斷增長的模型規(guī)模。然而，模型規(guī)模的不斷擴(kuò)大也帶來了訓(xùn)練成本和計算資源需求的急劇增加，這使得進(jìn)一步提升模型性能變得越來越困難。為了在不進(jìn)一步增加模型規(guī)模的前提下提升性能，一種新的范式——測試時計算擴(kuò)展（Test-Time Computation Extension）——應(yīng)運而生。測試時計算擴(kuò)展的核心思想是在測試階段（即模型的推理階段）投入額外的計算資源，從而增強(qiáng)模型的表現(xiàn)。這種方法避免了對模型進(jìn)行重新訓(xùn)練或大規(guī)模修改，提供了一種靈活且高效的性能提升途徑。
OpenAI 的 o1 模型（其具體實現(xiàn)細(xì)節(jié)并未公開）展示了測試時擴(kuò)展的巨大潛力，這引發(fā)了大量的復(fù)現(xiàn)和改進(jìn)工作。許多研究者試圖理解并復(fù)制 o1 的成功，但由于缺乏公開信息，這些努力面臨著巨大的挑戰(zhàn)。為了探索一種既簡單又有效的測試時擴(kuò)展方法，斯坦福研究人員的論文《s1: Simple test-time scaling》提出了 S1 方法。S1 通過在一個精心構(gòu)建的小規(guī)模推理數(shù)據(jù)集上進(jìn)行微調(diào)，并結(jié)合易于實現(xiàn)的預(yù)算強(qiáng)制技術(shù)，實現(xiàn)了強(qiáng)大的推理性能。S1 的目標(biāo)是提供一種可復(fù)現(xiàn)、易理解且高效的測試時擴(kuò)展方案，為 LLM 社區(qū)提供一個有價值的參考。
S1 方法的核心：簡約而不簡單
S1 的整體思路非常簡潔，可以概括為兩個主要步驟：首先，在一個名為 S1K 的、僅包含 1000 個高質(zhì)量推理樣本的數(shù)據(jù)集上，對預(yù)訓(xùn)練語言模型進(jìn)行監(jiān)督微調(diào)（Supervised Fine-Tuning, SFT）。這一步使得模型能夠更好地適應(yīng)推理任務(wù)的特定需求。然后，在測試階段，采用預(yù)算強(qiáng)制（Budget Forcing）技術(shù)，精確控制模型的推理過程，以此來提升模型的性能。預(yù)算強(qiáng)制通過限制模型的推理步數(shù)或推理時間，迫使模型在有限的資源內(nèi)進(jìn)行更高效的思考。
S1 的簡單性體現(xiàn)在多個方面。首先，S1K 數(shù)據(jù)集規(guī)模很小，但質(zhì)量非常高，其中的樣本都經(jīng)過了精心挑選，確保了訓(xùn)練數(shù)據(jù)的高效性。其次，模型的訓(xùn)練過程非常簡潔，只需要標(biāo)準(zhǔn)的監(jiān)督微調(diào)即可完成，無需復(fù)雜的訓(xùn)練策略或算法設(shè)計。最后，預(yù)算強(qiáng)制技術(shù)本身也易于理解和實現(xiàn)，不需要復(fù)雜的數(shù)學(xué)推導(dǎo)或編程技巧。這種簡單性使得 S1 方法易于復(fù)現(xiàn)和應(yīng)用，降低了研究門檻。
S1K 數(shù)據(jù)集：小而精的推理基石
S1K 數(shù)據(jù)集的構(gòu)建是 S1 方法成功的關(guān)鍵。該數(shù)據(jù)集遵循三個核心原則：質(zhì)量、難度和多樣性。首先，樣本必須具有高質(zhì)量，這意味著樣本必須避免格式錯誤、內(nèi)容不準(zhǔn)確或邏輯混亂等問題。只有高質(zhì)量的樣本才能為模型提供有效的學(xué)習(xí)信號。其次，樣本應(yīng)具有一定的挑戰(zhàn)性，能夠有效考察模型的推理能力，避免過于簡單的問題導(dǎo)致模型無法得到充分的訓(xùn)練。最后，樣本應(yīng)覆蓋多個領(lǐng)域和主題，以確保模型的泛化能力，避免模型在特定領(lǐng)域過擬合。
S1K 數(shù)據(jù)集的構(gòu)建過程包括以下幾個步驟：首先，從 16 個不同的數(shù)據(jù)源（涵蓋數(shù)學(xué)、物理、生物等多個領(lǐng)域）收集了 59,000 個初始樣本。這些樣本包含了推理問題，并使用 Google Gemini Flash Thinking API 生成了相應(yīng)的推理軌跡和答案。然后，進(jìn)行質(zhì)量過濾，移除包含 API 錯誤、格式問題、重復(fù)樣本等低質(zhì)量樣本，確保數(shù)據(jù)集的基本質(zhì)量。接下來，利用 Qwen2.5-7B-Instruct 和 Qwen2.5-32B-Instruct 模型的推理表現(xiàn)和推理軌跡長度作為指標(biāo)，篩選出具有一定難度的樣本，這保證了數(shù)據(jù)集能夠有效評估模型的推理能力。最后，根據(jù)數(shù)學(xué)主題分類系統(tǒng)（MSC）對樣本進(jìn)行分類，并從每個類別中挑選具有代表性的樣本，確保了數(shù)據(jù)集在不同領(lǐng)域和主題上的覆蓋度。最終形成的 S1K 數(shù)據(jù)集具有小規(guī)模、高質(zhì)量和多樣性的特點，為 S1 方法的成功奠定了堅實的基礎(chǔ)。
預(yù)算強(qiáng)制：S1 測試時擴(kuò)展的核心機(jī)制
測試時擴(kuò)展方法主要分為兩類：順序型和并行型。順序型方法指的是模型的推理過程是逐步進(jìn)行的，每一步的推理都依賴于前一步的結(jié)果。而并行型方法則允許模型同時進(jìn)行多個推理過程，例如同時探索多個推理路徑。S1 采用的是順序型擴(kuò)展方法，其核心技術(shù)是預(yù)算強(qiáng)制。預(yù)算強(qiáng)制通過兩種方式來控制模型的推理過程。一方面，當(dāng)模型生成的思考 Token 數(shù)量達(dá)到預(yù)設(shè)上限時，強(qiáng)制結(jié)束思考過程，并引導(dǎo)模型開始生成答案，這避免了模型無限制地進(jìn)行推理，導(dǎo)致資源浪費。另一方面，當(dāng)模型過早地嘗試結(jié)束思考過程時，抑制其生成結(jié)束思考的 Token，并通過添加 "Wait" 等提示詞，鼓勵模型進(jìn)行更深入的思考，這防止了模型在推理不充分的情況下給出答案。
預(yù)算強(qiáng)制的優(yōu)勢在于其簡單性和有效性。它易于理解和實現(xiàn)，無需復(fù)雜的算法設(shè)計。同時，它能夠有效控制模型的推理過程，顯著提升模型在多個任務(wù)上的性能。與其他測試時擴(kuò)展方法相比，例如條件長度控制、拒絕采樣和強(qiáng)化學(xué)習(xí)，預(yù)算強(qiáng)制在簡單性和有效性之間取得了較好的平衡。條件長度控制的局限性在于模型可能難以準(zhǔn)確理解和遵循 Prompt 中的長度限制。拒絕采樣的效率較低，并且可能導(dǎo)致模型過度依賴較短的推理路徑。強(qiáng)化學(xué)習(xí)方法雖然理論上能夠?qū)崿F(xiàn)更精細(xì)的資源控制，但訓(xùn)練過程較為復(fù)雜。
S1-32B 模型與實驗結(jié)果
S1-32B 模型是 S1 方法的具體實現(xiàn)，它基于 Qwen2.5-32B-Instruct 模型。S1-32B 的訓(xùn)練過程非常簡潔，僅僅是在 S1K 數(shù)據(jù)集上進(jìn)行監(jiān)督微調(diào)。訓(xùn)練過程使用 16 塊 H100 GPU，總訓(xùn)練時間僅需 26 分鐘。實驗結(jié)果表明，S1-32B 在多個推理任務(wù)（如 MATH500、AIME24 和 GPQA Diamond）上都取得了優(yōu)異的性能。而且，隨著預(yù)算強(qiáng)制計算量的增加，模型性能也隨之提升，這有力地驗證了預(yù)算強(qiáng)制的有效性。
為了深入理解 S1 的成功因素，研究團(tuán)隊進(jìn)行了一系列消融實驗。數(shù)據(jù)消融實驗旨在探究 S1K 數(shù)據(jù)集構(gòu)建的三個原則（質(zhì)量、難度和多樣性）對模型性能的影響。實驗結(jié)果表明，這三個原則的結(jié)合對于模型性能的提升至關(guān)重要，缺一不可。測試時擴(kuò)展方法消融實驗則旨在比較不同測試時擴(kuò)展方法對模型性能的影響。實驗結(jié)果表明，預(yù)算強(qiáng)制在控制性、擴(kuò)展性和最終性能方面均優(yōu)于其他方法。
S1 的優(yōu)勢、局限性與啟示
S1 方法的優(yōu)勢在于其簡單性、高效性和有效性。它易于理解、實現(xiàn)和復(fù)現(xiàn)，研究者和開發(fā)者可以快速上手，進(jìn)行實驗和應(yīng)用。同時，使用少量高質(zhì)量數(shù)據(jù)即可快速完成模型訓(xùn)練，顯著降低了訓(xùn)練成本。更重要的是，S1 在多個推理任務(wù)上取得了與 o1-preview 相當(dāng)?shù)男阅?，證明了其在提升 LLM 推理能力方面的潛力。 并且S1的開源特性肯定會進(jìn)一步促進(jìn)了LLM社區(qū)的協(xié)作發(fā)展。
然而，S1 方法也存在一些局限性。一是S1 的性能一般來說可能是存在上限的，低于更復(fù)雜的方法（如 DeepSeek-R1）。這可能是因為 S1 僅依賴于簡單的監(jiān)督微調(diào)和預(yù)算強(qiáng)制，而沒有引入更復(fù)雜的訓(xùn)練策略或模型架構(gòu)。二是預(yù)算強(qiáng)制本身也存在局限性。隨著預(yù)算強(qiáng)制計算量的增加，模型性能的提升會逐漸趨于平緩，甚至可能出現(xiàn)下降。這意味著預(yù)算強(qiáng)制存在一個最優(yōu)的計算量范圍，超過這個范圍可能導(dǎo)致資源浪費或性能下降。另外，預(yù)算強(qiáng)制的有效性受到模型上下文窗口長度的限制，對于需要更長推理路徑的任務(wù)，預(yù)算強(qiáng)制可能無法充分發(fā)揮作用。預(yù)算強(qiáng)制還缺乏自適應(yīng)性，它是一種靜態(tài)的控制方法，無法根據(jù)不同的輸入樣本動態(tài)調(diào)整計算資源。對于不同難度的樣本，使用相同的計算預(yù)算可能并不是最優(yōu)策略。 S1還對順序推理有依賴, 這限制了它利用并行計算加速推理的能力。 并且預(yù)算強(qiáng)制機(jī)制的內(nèi)部工作原理仍然不夠透明, 增加了理解模型決策的難度.
當(dāng)然，S1 的成功也帶來了一些重要的啟示。首先，它強(qiáng)調(diào)了數(shù)據(jù)質(zhì)量的重要性，即使是規(guī)模較小的數(shù)據(jù)集，只要質(zhì)量足夠高，也能顯著提升模型性能。其次，它表明簡單的策略在特定場景下也能取得與復(fù)雜方法相媲美的效果，這啟示我們不應(yīng)一味追求復(fù)雜性，而應(yīng)根據(jù)具體問題選擇合適的解決方案。最后，S1 的成功再次證明了測試時擴(kuò)展的巨大潛力，為 LLM 性能提升提供了一條新的途徑。
S1的總體評價
S1方法為測試時計算擴(kuò)展提供了一個新視角，證明了即使是簡單的方法，也可以在提升大型語言模型推理能力上取得顯著效果。 它的主要優(yōu)點是簡單, 高效, 并且開源。這使得研究者和開發(fā)者更容易復(fù)現(xiàn)和應(yīng)用S1, 促進(jìn)了整個LLM社區(qū)的合作與發(fā)展。
但是，S1 的局限性也是顯而易見的。 它的性能上限可能低于更復(fù)雜的方法。預(yù)算強(qiáng)制技術(shù)本身也存在局限, 比如性能提升會趨于平緩, 并且受限于模型的上下文窗口大小。 S1高度依賴于S1K數(shù)據(jù)集的質(zhì)量, 并且對順序推理的依賴限制了它利用并行計算的能力。 此外, 預(yù)算強(qiáng)制的內(nèi)部機(jī)制不夠透明, 使得模型的可解釋性較差。這些局限性也為未來的研究提供了方向。
什么是測試時擴(kuò)展？
這個術(shù)語太不友好，我花點篇幅解釋一下吧。
在機(jī)器學(xué)習(xí)領(lǐng)域，模型的生命周期通常劃分為兩個主要階段：訓(xùn)練階段和測試階段（也稱為推理階段）。訓(xùn)練階段是模型學(xué)習(xí)的過程，通過大量的數(shù)據(jù)和計算，模型不斷調(diào)整其內(nèi)部參數(shù)，以盡可能準(zhǔn)確地完成預(yù)定的任務(wù)，如圖像識別、文本生成或語言翻譯等。一旦模型訓(xùn)練完成，其內(nèi)部參數(shù)就被固定下來。測試階段則是對訓(xùn)練好的模型進(jìn)行性能評估的環(huán)節(jié)，在這個階段，模型會接收之前未見過的新數(shù)據(jù)，并根據(jù)這些新數(shù)據(jù)進(jìn)行預(yù)測或生成輸出。關(guān)鍵在于，在傳統(tǒng)的測試流程中，模型不會再更新其參數(shù)，而是直接利用訓(xùn)練階段學(xué)習(xí)到的知識來進(jìn)行推理。
測試時擴(kuò)展（Test-Time Computation Extension，簡稱 Test-Time Extension）是一種旨在提升模型在測試階段性能的技術(shù)范式。它的核心思想非常直觀：在測試階段，也就是模型進(jìn)行推理的時候，為模型提供額外的計算資源或者計算步驟，以此來提高模型輸出結(jié)果的準(zhǔn)確率或整體性能，而無需對模型本身進(jìn)行任何修改（即不改變模型在訓(xùn)練階段學(xué)習(xí)到的參數(shù)）。這種方法與傳統(tǒng)的“訓(xùn)練-測試”二階段流程有所不同，它在測試階段引入了額外的計算干預(yù)。
采用測試時擴(kuò)展的主要動機(jī)有以下幾點：首先，它允許在不增加模型規(guī)模的前提下提升性能。訓(xùn)練更大、更復(fù)雜的模型通常需要耗費更多的計算資源和時間，而測試時擴(kuò)展提供了一條在資源受限的情況下提高性能的途徑。其次，測試時擴(kuò)展具有很強(qiáng)的靈活性。由于它不涉及模型的重新訓(xùn)練，因此可以根據(jù)實際需要在測試時動態(tài)地調(diào)整計算量，從而在模型性能和計算成本之間取得平衡。此外，對于某些特定類型的任務(wù)，尤其是那些需要復(fù)雜推理的任務(wù)，標(biāo)準(zhǔn)的推理過程可能不足以讓模型充分發(fā)揮其潛力，測試時擴(kuò)展允許模型在這些任務(wù)上進(jìn)行更深入、更細(xì)致的“思考”。最后，測試時擴(kuò)展還可以在一定程度上彌補訓(xùn)練階段的不足，例如當(dāng)訓(xùn)練數(shù)據(jù)不夠充分或者訓(xùn)練方法存在局限時，通過在測試時增加計算，模型有可能更好地處理那些在訓(xùn)練時沒有充分學(xué)習(xí)到的復(fù)雜模式。
測試時擴(kuò)展與一些常見的機(jī)器學(xué)習(xí)技術(shù)有著明顯的區(qū)別。例如，它不同于訓(xùn)練時的數(shù)據(jù)增強(qiáng)（Data Augmentation）。數(shù)據(jù)增強(qiáng)是在訓(xùn)練階段對訓(xùn)練數(shù)據(jù)進(jìn)行各種變換，以增加數(shù)據(jù)的多樣性，從而提高模型的泛化能力。而測試時擴(kuò)展則完全作用于測試階段，與訓(xùn)練數(shù)據(jù)無關(guān)。它也不同于集成學(xué)習(xí)（Ensemble Learning）。集成學(xué)習(xí)通常需要訓(xùn)練多個不同的模型，然后在測試時將這些模型的預(yù)測結(jié)果進(jìn)行整合，例如通過投票或平均的方式。測試時擴(kuò)展通常只針對單個模型，通過增加該模型在測試時的計算量來提升其性能。測試時擴(kuò)展也不等同于直接增加模型的深度或復(fù)雜度，后者是在模型設(shè)計階段就確定下來的，并且會改變模型本身的結(jié)構(gòu)，而測試時擴(kuò)展不會改變模型的固有結(jié)構(gòu)。
常見的測試時擴(kuò)展方法包括多次推理（讓模型對同一輸入進(jìn)行多次獨立推理并選擇最一致的答案）、思考鏈（通過提示詞引導(dǎo)模型逐步推理）、自修正（讓模型對自己的輸出進(jìn)行迭代改進(jìn)）以及搜索算法（在測試時探索多個可能的輸出）等。每種方法都有其適用的場景和優(yōu)缺點。 S1方法，通過在小規(guī)模數(shù)據(jù)集上進(jìn)行微調(diào)，再結(jié)合推理時的預(yù)算強(qiáng)制，實際上也是一種測試時擴(kuò)展策略，它通過限制推理步驟來達(dá)到提升效果的目的。
總之，測試時擴(kuò)展是一種在不改變模型本身的情況下，通過在測試階段增加計算資源或計算步驟來提升模型性能的技術(shù)。它提供了一種靈活且有效的方法，尤其適用于那些需要復(fù)雜推理的任務(wù)，為大型語言模型在實際應(yīng)用中的性能提升開辟了新的可能性。
結(jié)論與未來展望
S1 的主要貢獻(xiàn)在于提出了一種簡單、高效且有效的提升 LLM 推理能力的方法，并構(gòu)建和開源了高質(zhì)量的 S1K 推理數(shù)據(jù)集。同時，S1 驗證了預(yù)算強(qiáng)制技術(shù)在測試時擴(kuò)展中的有效性，并強(qiáng)調(diào)了 S1 方法的開源特性。未來的研究方向包括探索更有效的測試時擴(kuò)展方法，例如改進(jìn)預(yù)算強(qiáng)制技術(shù)，或者將其與強(qiáng)化學(xué)習(xí)等更高級的技術(shù)結(jié)合起來。此外，可以將 S1 的思想應(yīng)用于其他領(lǐng)域和任務(wù)，如代碼生成、文本摘要、對話系統(tǒng)等，以驗證其通用性。另一個重要的研究方向是突破上下文窗口的限制，例如采用并行擴(kuò)展方法或分塊推理等，以處理更復(fù)雜的推理任務(wù)。增強(qiáng)模型的可解釋性, 深入理解預(yù)算強(qiáng)制機(jī)制的內(nèi)部工作原理也是一個重要的研究方向。繼續(xù)擴(kuò)大和完善 S1K 數(shù)據(jù)集，使其涵蓋更廣泛的領(lǐng)域和任務(wù), 以及探索將人類反饋融入S1, 都有助于提高模型的泛化能力和可靠性。
總而言之，S1 方法為 LLM 測試時計算擴(kuò)展提供了一個簡單而有效的解決方案，并為未來的研究提供了寶貴的啟示。隨著 LLM 技術(shù)的不斷發(fā)展，相信 S1 方法及其衍生技術(shù)將在更廣泛的領(lǐng)域得到應(yīng)用和推廣，推動 LLM 推理能力的持續(xù)提升。
參考論文： arXiv:2501.19393v2 [cs.CL] 3 Feb 2025
