如果給小模型更長的思考時間，它們性能可以超越更大規(guī)模的模型。

最近一段時間，業(yè)內(nèi)對小模型的研究熱情空前地高漲，通過一些「實(shí)用技巧」讓它們在性能上超越更大規(guī)模的模型。

可以說，將目光放到提升較小模型的性能上來有其必然性。對于大語言模型而言，訓(xùn)練時計算（train-time compute）的擴(kuò)展主導(dǎo)了它們的發(fā)展。盡管這種模式已被證明非常有效，但越來越大模型的預(yù)訓(xùn)練所需的資源卻變得異常昂貴，數(shù)十億美元的集群已經(jīng)出現(xiàn)。

因此，這一趨勢引發(fā)了人們對另外一種互補(bǔ)方法的極大興趣，即測試時計算擴(kuò)展（test-time compute scaling）。測試時方法不依賴于越來越大的預(yù)訓(xùn)練預(yù)算，而是使用動態(tài)推理策略，讓模型在更難的問題上「思考更長時間」。一個突出的例子是 OpenAI 的 o1 模型，隨著測試時計算量的增加，它在困難數(shù)學(xué)問題上表現(xiàn)出持續(xù)的進(jìn)步。

雖然我們不清楚 o1 是如何訓(xùn)練的，但 DeepMind 最近的研究表明，可以通過迭代自我改進(jìn)或使用獎勵模型在解決方案空間上進(jìn)行搜索等策略來實(shí)現(xiàn)測試時計算的最佳擴(kuò)展。通過自適應(yīng)地按 prompt 分配測試時計算，較小的模型可以與較大、資源密集型模型相媲美，有時甚至超越它們。當(dāng)內(nèi)存受限且可用硬件不足以運(yùn)行較大模型時，擴(kuò)展時間時計算尤其有利。然而這種有前途的方法是用閉源模型演示的，沒有發(fā)布任何實(shí)現(xiàn)細(xì)節(jié)或代碼。

DeepMind 論文：https://arxiv.org/pdf/2408.03314

在過去幾個月里，HuggingFace 一直在深入研究，試圖對這些結(jié)果進(jìn)行逆向工程并復(fù)現(xiàn)。他們在這篇博文將介紹：

• 計算最優(yōu)擴(kuò)展（compute-optimal scaling）：通過實(shí)現(xiàn) DeepMind 的技巧來提升測試時開放模型的數(shù)學(xué)能力。
• 多樣性驗(yàn)證器樹搜索 (DVTS)：它是為驗(yàn)證器引導(dǎo)樹搜索技術(shù)開發(fā)的擴(kuò)展。這種簡單高效的方法提高了多樣性并提供了更好的性能，特別是在測試時計算預(yù)算較大的情況下。
• 搜索和學(xué)習(xí)：一個輕量級工具包，用于使用 LLM 實(shí)現(xiàn)搜索策略，并使用 vLLM 實(shí)現(xiàn)速度提升。

那么，計算最優(yōu)擴(kuò)展在實(shí)踐中效果如何呢？在下圖中，如果你給它們足夠的「思考時間」，規(guī)模很小的 1B 和 3B Llama Instruct 模型在具有挑戰(zhàn)性的 MATH-500 基準(zhǔn)上，超越了比它們大得多的 8B、70B 模型。

HuggingFace 聯(lián)合創(chuàng)始人兼 CEO Clem Delangue 表示，在 OpenAI o1 公開亮相僅 10 天后，我們很高興地揭曉了其成功背后的突破性技術(shù)的開源版本：擴(kuò)展測試時計算。通過給模型更長的「思考時間」，1B 模型可以擊敗 8B、3B 模型可以擊敗 70B。當(dāng)然，完整的技術(shù)配方是開源的。

各路網(wǎng)友看到這些結(jié)果也不淡定了，直呼不可思議，并認(rèn)為這是小模型的勝利。

接下來，HuggingFace 深入探討了產(chǎn)生上述結(jié)果背后的原因，并幫助讀者了解實(shí)現(xiàn)測試時計算擴(kuò)展的實(shí)用策略。

擴(kuò)展測試時計算策略

擴(kuò)展測試時計算主要有以下兩種主要策略：

• 自我改進(jìn)：模型通過在后續(xù)迭代中識別和糾錯來迭代改進(jìn)自己的輸出或「想法」。雖然這種策略在某些任務(wù)上有效，但通常要求模型具有內(nèi)置的自我改進(jìn)機(jī)制，這可能會限制其適用性。
• 針對驗(yàn)證器進(jìn)行搜索：這種方法側(cè)重于生成多個候選答案并使用驗(yàn)證器選擇最佳答案。驗(yàn)證器可以是基于硬編碼的啟發(fā)式方法，也可以是學(xué)得的獎勵模型。本文將重點(diǎn)介紹學(xué)得的驗(yàn)證器，它包括了 Best-of-N 采樣和樹搜索等技術(shù)。這種搜索策略更靈活，可以適應(yīng)問題的難度，不過它們的性能受到驗(yàn)證器質(zhì)量的限制。

HuggingFace 專注于基于搜索的方法，它們是測試時計算優(yōu)化的實(shí)用且可擴(kuò)展的解決方案。下面是三種策略：

• Best-of-N：通常使用獎勵模型為每個問題生成多個響應(yīng)并為每個候選答案分配分?jǐn)?shù)，然后選擇獎勵最高的答案（或稍后討論的加權(quán)變體）。這種方法強(qiáng)調(diào)答案質(zhì)量而非頻率。
• 集束搜索：一種探索解決方案空間的系統(tǒng)搜索方法，通常與過程獎勵模型 (PRM) 結(jié)合使用，以優(yōu)化問題解決中間步驟的采樣和評估。與對最終答案產(chǎn)生單一分?jǐn)?shù)的傳統(tǒng)獎勵模型不同，PRM 會提供一系列分?jǐn)?shù)，其中推理過程的每個步驟都有一個分?jǐn)?shù)。這種細(xì)粒度反饋能力使得 PRM 成為 LLM 搜索方法的自然選擇。
• 多樣性驗(yàn)證器樹搜索 (DVTS)：HuggingFace 開發(fā)的集束搜索擴(kuò)展，將初始集束拆分為獨(dú)立的子樹，然后使用 PRM 貪婪地擴(kuò)展這些子樹。這種方法提高了解決方案的多樣性和整體性能，尤其是在測試時計算預(yù)算較大的情況下。

實(shí)驗(yàn)設(shè)置

實(shí)驗(yàn)設(shè)置包括以下步驟：

• 首先給 LLM 提供一個數(shù)學(xué)問題，讓其生成 N 個部分解，例如，推導(dǎo)過程中的中間步驟。
• 每個 step 都由 PRM 評分，PRM 估計每個步驟最終達(dá)到正確答案的概率。
• 一旦搜索策略結(jié)束，最終候選解決方案將由 PRM 排序以產(chǎn)生最終答案。

為了比較各種搜索策略，本文使用了以下開源模型和數(shù)據(jù)集：

• 模型：使用 meta-llama/Llama-3.2-1B-Instruct 為主要模型，用于擴(kuò)展測試時計算；
• 過程獎勵模型 PRM：為了指導(dǎo)搜索策略，本文使用了 RLHFlow/Llama3.1-8B-PRM-Deepseek-Data，這是一個經(jīng)過過程監(jiān)督訓(xùn)練的 80 億獎勵模型。過程監(jiān)督是一種訓(xùn)練方法，模型在推理過程的每一步都會收到反饋，而不僅僅是最終結(jié)果；
• 數(shù)據(jù)集：本文在 MATH-500 子集上進(jìn)行了評估，這是 OpenAI 作為過程監(jiān)督研究的一部分發(fā)布的 MATH 基準(zhǔn)數(shù)據(jù)集。這些數(shù)學(xué)問題涵蓋了七個科目，對人類和大多數(shù)大語言模型來說都具有挑戰(zhàn)性。

本文將從一個簡單的基線開始，然后逐步結(jié)合其他技術(shù)來提高性能。

多數(shù)投票

多數(shù)投票是聚合 LLM 輸出的最直接方法。對于給定的數(shù)學(xué)問題，會生成 N 個候選解，然后選擇出現(xiàn)頻率最高的答案。在所有的實(shí)驗(yàn)中，本文采樣了多達(dá) N=256 個候選解，溫度參數(shù) T=0.8，并為每個問題生成了最多 2048 個 token。

以下是多數(shù)投票應(yīng)用于 Llama 3.2 1B Instruct 時的表現(xiàn)：

結(jié)果表明，多數(shù)投票比貪婪解碼基線有顯著的改進(jìn)，但其收益在大約 N=64 generation 后開始趨于平穩(wěn)。這種限制的出現(xiàn)是因?yàn)槎鄶?shù)投票難以解決需要細(xì)致推理的問題。

基于多數(shù)投票的局限性，讓我們看看如何結(jié)合獎勵模型來提高性能。

超越多數(shù)：Best-of-N

Best-of-N 是多數(shù)投票算法的簡單且有效的擴(kuò)展，它使用獎勵模型來確定最合理的答案。該方法有兩種主要變體：

普通的 Best-of-N：生成 N 個獨(dú)立響應(yīng)，選擇 RM 獎勵最高的一個作為最終回答。這確保了選擇置信度最高的響應(yīng)，但它并沒有考慮到回答之間的一致性。

加權(quán) Best-of-N：匯總所有相同響應(yīng)的得分，并選擇總獎勵最高的回答。這種方法通過重復(fù)出現(xiàn)來提高分?jǐn)?shù)，從而優(yōu)先考慮高質(zhì)量的回答。從數(shù)學(xué)上講，回答的權(quán)重 a_i：

其中，RM (p,s_i) 是對于問題 p 的第 i 個解決方案 s_i 的獎勵模型分?jǐn)?shù)。

通常，人們使用結(jié)果獎勵模型 (ORM) 來獲得單個解決方案級別的分?jǐn)?shù)。但為了與其他搜索策略進(jìn)行公平比較，使用相同的 PRM 對 Best-of-N 的解決方案進(jìn)行評分。如下圖所示，PRM 為每個解決方案生成一個累積的步驟級分?jǐn)?shù)序列，因此需要對步驟進(jìn)行規(guī)約（reduction）以獲得單個解決方案級分?jǐn)?shù)：

最常見的規(guī)約如下：

• Min：使用所有步驟中的最低分?jǐn)?shù)。
• Prod：使用階梯分?jǐn)?shù)的乘積。
• Last：使用步驟中的最終分?jǐn)?shù)。該分?jǐn)?shù)包含所有先前步驟的累積信息，因此將 PRM 有效地視為能夠?qū)Σ糠纸鉀Q方案進(jìn)行評分的 ORM。

以下是應(yīng)用 Best-of-N 的兩種變體得到的結(jié)果：

結(jié)果揭示了一個明顯的優(yōu)勢：加權(quán)的 Best-of-N 始終優(yōu)于普通的 Best-of-N，特別是在發(fā)電預(yù)算較大的情況下。它能夠匯總相同答案的分?jǐn)?shù)，確保即使頻率較低但質(zhì)量較高的答案也能得到有效的優(yōu)先處理。

然而，盡管有這些改進(jìn)，仍然達(dá)不到 Llama 8B 模型所達(dá)到的性能，并且在 N=256 時 Best-of-N 方法開始趨于穩(wěn)定。

可以通過逐步監(jiān)督搜索過程來進(jìn)一步突破界限嗎？

使用 PRM 的集束搜索

作為一種結(jié)構(gòu)化搜索方法，集束搜索可以系統(tǒng)地探索解決方案空間，使其成為在測試時改進(jìn)模型輸出的強(qiáng)大工具。與 PRM 結(jié)合使用時，集束搜索可以優(yōu)化問題解決中中間步驟的生成和評估。集束搜索的工作方式如下：

• 通過保持固定數(shù)量的「集束」或活動路徑 N ，迭代生成多個候選解決方案。
• 在第一次迭代中，從溫度為 T 的 LLM 中抽取 N 個獨(dú)立步驟，以引入響應(yīng)的多樣性。這些步驟通常由停止標(biāo)準(zhǔn)定義，例如終止于新行 \n 或雙新行 \n\n。
• 使用 PRM 對每個步驟進(jìn)行評分，并選擇前 N/M 個步驟作為下一輪生成的候選。這里 M 表示給定活動路徑的「集束寬度」。與 Best-of-N 一樣，使用「最后」的規(guī)約來對每次迭代的部分解決方案進(jìn)行評分。
• 通過在解決方案中采樣 M 個后續(xù)步驟來擴(kuò)展在步驟 (3) 中選擇的步驟。
• 重復(fù)步驟 (3) 和 (4)，直到達(dá)到 EOS token 或超過最大搜索深度。

通過允許 PRM 評估中間步驟的正確性，集束搜索可以在流程早期識別并優(yōu)先考慮有希望的路徑。這種逐步評估策略對于數(shù)學(xué)等復(fù)雜的推理任務(wù)特別有用，這是因?yàn)轵?yàn)證部分解決方案可以顯著改善最終結(jié)果。

實(shí)現(xiàn)細(xì)節(jié)

在實(shí)驗(yàn)中，HuggingFace 遵循 DeepMind 的超參數(shù)選擇，并按照以下方式運(yùn)行集束搜索：

• 計算擴(kuò)展為 4、16、64、256 時的 N 個集束
• 固定集束寬度 M=4
• 在溫度 T=0.8 時采樣
• 最多 40 次迭代，即最大深度為 40 步的樹

如下圖所示，結(jié)果令人震驚：在 N=4 的測試時預(yù)算下，集束搜索實(shí)現(xiàn)了與 N=16 時 Best-of-N 相同的準(zhǔn)確率，即計算效率提高了 4 倍！此外，集束搜索的性能與 Llama 3.1 8B 相當(dāng)，每個問題僅需 N=32 解決方案。計算機(jī)科學(xué)博士生在數(shù)學(xué)方面的平均表現(xiàn)約為 40%，因此對于 1B 模型來說，接近 55% 已經(jīng)很不錯了！

哪些問題集束搜索解決得最好

雖然總體上很明顯，集束搜索是一種比 Best-of-N 或多數(shù)投票更好的搜索策略，但 DeepMind 的論文表明，每種策略都有權(quán)衡，這取決于問題的難度和測試時計算預(yù)算。

為了了解哪些問題最適合哪種策略，DeepMind 計算了估計問題難度的分布，并將結(jié)果分成五等分。換句話說，每個問題被分配到 5 個級別之一，其中級別 1 表示較容易的問題，級別 5 表示最難的問題。為了估計問題難度，DeepMind 為每個問題生成了 2048 個候選解決方案，并進(jìn)行了標(biāo)準(zhǔn)采樣，然后提出了以下啟發(fā)式方法：

• Oracle：使用基本事實(shí)標(biāo)簽估計每個問題的 pass@1 分?jǐn)?shù)，對 pass@1 分?jǐn)?shù)的分布進(jìn)行分類以確定五分位數(shù)。
• 模型：使用每個問題的平均 PRM 分?jǐn)?shù)分布來確定五分位數(shù)。這里的直覺是：更難的問題分?jǐn)?shù)會更低。

下圖是根據(jù) pass@1 分?jǐn)?shù)和四個測試時計算預(yù)算 N=[4,16,64,256] 對各種方法的細(xì)分：

可以看到，每個條形表示測試時計算預(yù)算，并且在每個條形內(nèi)顯示每種方法的相對準(zhǔn)確度。例如在難度等級 2 的四個條形中：

多數(shù)投票是所有計算預(yù)算中表現(xiàn)最差的方法，除了 N=256（集束搜索表現(xiàn)最差）。

集束搜索最適合 N=[4,16,64]，但 Best-of-N 最適合 N=256。

應(yīng)該說，集束搜索在中等難度和困難難度問題（3-5 級）中取得了持續(xù)的進(jìn)展，但在較簡單問題上，尤其是在計算預(yù)算較大的情況下，它的表現(xiàn)往往比 Best-of-N（甚至多數(shù)投票）更差。

通過觀察集束搜索生成的結(jié)果樹，HuggingFace 意識到，如果單個步驟被分配了高獎勵，那么整棵樹就在該軌跡上崩潰，從而影響多樣性。這促使他們探索一種最大化多樣性的集束搜索擴(kuò)展。

DVTS：通過多樣性提升性能

正如上面所看到的，集束搜索比 Best-of-N 具有更好的性能，但在處理簡單問題和測試時計算預(yù)算較大時往往表現(xiàn)不佳。

為了解決這個問題，HuggingFace 開發(fā)了一個擴(kuò)展，稱之為「多樣性驗(yàn)證器樹搜索」（DVTS），旨在最大限度地提高 N 較大時的多樣性。

DVTS 的工作方式與集束搜索類似，但有以下修改：

• 對于給定的 N 和 M，將初始集束擴(kuò)展為 N/M 個獨(dú)立子樹。
• 對于每個子樹，選擇具有最高 PRM 分?jǐn)?shù)的步驟。
• 從步驟 (2) 中選擇的節(jié)點(diǎn)生成 M 個新步驟，并選擇具有最高 PRM 分?jǐn)?shù)的步驟。
• 重復(fù)步驟 (3)，直到達(dá)到 EOS token 或最大樹深度。

下圖是將 DVTS 應(yīng)用于 Llama 1B 的結(jié)果：

可以看到，DVTS 為集束搜索提供了一種補(bǔ)充策略：在 N 較小時，集束搜索更有效地找到正確的解決方案；但在 N 較大時，DVTS 候選的多樣性開始發(fā)揮作用，可以獲得更好的性能。

此外在問題難度細(xì)分中，DVTS 在 N 較大時提高了簡單 / 中等問題的性能，而集束搜索在 N 較小時表現(xiàn)最佳。

計算 - 最優(yōu)擴(kuò)展（compute-optimal scaling）

有了各種各樣的搜索策略，一個自然的問題是哪一個是最好的？在 DeepMind 的論文中（可參考《Scaling LLM Test-Time Compute Optimally can be More Effective than Scaling Model Parameters 》），他們提出了一種計算 - 最優(yōu)擴(kuò)展策略，該策略可以選擇搜索方法和超參數(shù) θ，以便在給定的計算預(yù)算 N 下達(dá)到最佳性能：

其中是問題 q 的正確答案。表示計算 - 最優(yōu)的擴(kuò)展策略。由于直接計算有些棘手，DeepMind 提出了一種基于問題難度的近似方法，即根據(jù)哪種搜索策略在給定難度級別上達(dá)到最佳性能來分配測試時的計算資源。

例如，對于較簡單的問題和較低的計算預(yù)算，最好使用 Best-of-N 等策略，而對于較難的問題，集 shu 搜索是更好的選擇。下圖為計算 - 最優(yōu)曲線！

擴(kuò)展到更大的模型

本文還探索了將計算 - 最優(yōu)（compute-optimal）的方法擴(kuò)展到 Llama 3.2 3B Instruct 模型，以觀察 PRM 在與策略自身容量相比時在哪個點(diǎn)開始減弱。結(jié)果顯示，計算 - 最優(yōu)的擴(kuò)展效果非常好，3B 模型的性能超過了 Llama 3.1 70B Instruct（后者是前者大小的 22 倍！）。

接下來該怎么辦？

對測試時計算擴(kuò)展的探索揭示了利用基于搜索的方法的潛力和挑戰(zhàn)。展望未來，本文提出了幾個令人興奮的方向：

• 強(qiáng)驗(yàn)證器：強(qiáng)驗(yàn)證器在提高性能方面發(fā)揮著關(guān)鍵作用，提高驗(yàn)證器的穩(wěn)健性和通用性對于推進(jìn)這些方法至關(guān)重要；
• 自我驗(yàn)證：最終目標(biāo)是實(shí)現(xiàn)自我驗(yàn)證，即模型可以自主驗(yàn)證自己的輸出。這種方法似乎是 o1 等模型正在做的事情，但在實(shí)踐中仍然難以實(shí)現(xiàn)。與標(biāo)準(zhǔn)監(jiān)督微調(diào) (SFT) 不同，自我驗(yàn)證需要更細(xì)致的策略；
• 將思維融入過程：在生成過程中融入明確的中間步驟或思維可以進(jìn)一步增強(qiáng)推理和決策能力。通過將結(jié)構(gòu)化推理融入搜索過程，可以在復(fù)雜任務(wù)上實(shí)現(xiàn)更好的表現(xiàn)；
• 搜索作為數(shù)據(jù)生成工具：該方法還可以充當(dāng)強(qiáng)大的數(shù)據(jù)生成過程，創(chuàng)建高質(zhì)量的訓(xùn)練數(shù)據(jù)集。例如，根據(jù)搜索產(chǎn)生的正確軌跡對 Llama 1B 等模型進(jìn)行微調(diào)可以帶來顯著的收益。這種基于策略的方法類似于 ReST 或 V-StaR 等技術(shù)，但具有搜索的額外優(yōu)勢，為迭代改進(jìn)提供了一個有希望的方向；
• 調(diào)用更多的 PRM：PRM 相對較少，限制了其更廣泛的應(yīng)用。為不同領(lǐng)域開發(fā)和共享更多 PRM 是社區(qū)可以做出重大貢獻(xiàn)的關(guān)鍵領(lǐng)域。

原文鏈接：https://huggingface.co/spaces/HuggingFaceH4/blogpost-scaling-test-time-compute