3B模型長思考後擊敗70B！HuggingFace逆向出o1背後技術細節並開源

機器之心報道

機器之心編輯部

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

可以說，將目光放到提升較小模型的性能上來有其必然性。對於大語言模型而言，訓練時計算（train-time compute）的擴展主導了它們的發展。儘管這種模式已被證明非常有效，但越來越大模型的預訓練所需的資源卻變得異常昂貴，數十億美元的集羣已經出現。

因此，這一趨勢引發了人們對另外一種互補方法的極大興趣，即測試時計算擴展（test-time compute scaling）。測試時方法不依賴於越來越大的預訓練預算，而是使用動態推理策略，讓模型在更難的問題上「思考更長時間」。一個突出的例子是 OpenAI 的 o1 模型，隨着測試時計算量的增加，它在困難數學問題上表現出持續的進步。

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

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

在過去幾個月裡，HuggingFace 一直在深入研究，試圖對這些結果進行逆向工程並復現。他們在這篇博文將介紹：

那麼，計算最優擴展在實踐中效果如何呢？在下圖中，如果你給它們足夠的「思考時間」，規模很小的 1B 和 3B Llama Instruct 模型在具有挑戰性的 MATH-500 基準上，超越了比它們大得多的 8B、70B 模型。

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

各路網友看到這些結果也不淡定了，直呼不可思議，並認爲這是小模型的勝利。

接下來，HuggingFace 深入探討了產生上述結果背後的原因，並幫助讀者瞭解實現測試時計算擴展的實用策略。

擴展測試時計算策略

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

HuggingFace 專注於基於搜索的方法，它們是測試時計算優化的實用且可擴展的解決方案。下面是三種策略：

實驗設置

實驗設置包括以下步驟：

爲了比較各種搜索策略，本文使用了以下開源模型和數據集：

本文將從一個簡單的基線開始，然後逐步結合其他技術來提高性能。

多數投票

多數投票是聚合 LLM 輸出的最直接方法。對於給定的數學問題，會生成 N 個候選解，然後選擇出現頻率最高的答案。在所有的實驗中，本文采樣了多達 N=256 個候選解，溫度參數 T=0.8，併爲每個問題生成了最多 2048 個 token。

以下是多數投票應用於 Llama 3.2 1B Instruct 時的表現：

結果表明，多數投票比貪婪解碼基線有顯著的改進，但其收益在大約 N=64 generation 後開始趨於平穩。這種限制的出現是因爲多數投票難以解決需要細緻推理的問題。

基於多數投票的侷限性，讓我們看看如何結合獎勵模型來提高性能。

超越多數：Best-of-N

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

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

加權 Best-of-N：彙總所有相同響應的得分，並選擇總獎勵最高的回答。這種方法通過重複出現來提高分數，從而優先考慮高質量的回答。從數學上講，回答的權重 a_i：

其中，RM (p,s_i) 是對於問題 p 的第 i 個解決方案 s_i 的獎勵模型分數。

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

最常見的規約如下：

以下是應用 Best-of-N 的兩種變體得到的結果：

結果揭示了一個明顯的優勢：加權的 Best-of-N 始終優於普通的 Best-of-N，特別是在發電預算較大的情況下。它能夠彙總相同答案的分數，確保即使頻率較低但質量較高的答案也能得到有效的優先處理。

然而，儘管有這些改進，仍然達不到 Llama 8B 模型所達到的性能，並且在 N=256 時 Best-of-N 方法開始趨於穩定。

可以通過逐步監督搜索過程來進一步突破界限嗎？

使用 PRM 的集束搜索

作爲一種結構化搜索方法，集束搜索可以系統地探索解決方案空間，使其成爲在測試時改進模型輸出的強大工具。與 PRM 結合使用時，集束搜索可以優化問題解決中中間步驟的生成和評估。集束搜索的工作方式如下：

通過允許 PRM 評估中間步驟的正確性，集束搜索可以在流程早期識別並優先考慮有希望的路徑。這種逐步評估策略對於數學等複雜的推理任務特別有用，這是因爲驗證部分解決方案可以顯著改善最終結果。

實現細節

在實驗中，HuggingFace 遵循 DeepMind 的超參數選擇，並按照以下方式運行集束搜索：

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

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

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

爲了瞭解哪些問題最適合哪種策略，DeepMind 計算了估計問題難度的分佈，並將結果分成五等分。換句話說，每個問題被分配到 5 個級別之一，其中級別 1 表示較容易的問題，級別 5 表示最難的問題。爲了估計問題難度，DeepMind 爲每個問題生成了 2048 個候選解決方案，並進行了標準採樣，然後提出了以下啓發式方法：

下圖是根據 pass@1 分數和四個測試時計算預算 N=[4,16,64,256] 對各種方法的細分：

可以看到，每個條形表示測試時計算預算，並且在每個條形內顯示每種方法的相對準確度。例如在難度等級 2 的四個條形中：

多數投票是所有計算預算中表現最差的方法，除了 N=256（集束搜索表現最差）。

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

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

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

DVTS：通過多樣性提升性能

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

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

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

下圖是將 DVTS 應用於 Llama 1B 的結果：

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

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

計算 - 最優擴展（compute-optimal scaling）

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

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

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

擴展到更大的模型

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

接下來該怎麼辦？

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

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