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