大模型推理時的尺度擴展定律

FesianXu at 20250212 at Wechat Search Team

前言

大模型的尺度擴展定律告訴我們：『LLM的性能會隨著模型的參數量、模型的訓練量、模型的訓練數據量的增加而增加』。訓練存在尺度擴展定律，測試也存在尺度擴展定律，實踐告訴我們在推理時候增大計算量，往往可以獲得模型性能收益。那么在給定了計算預算的前提下，如何安排預算才能達到最好的模型效果呢？本篇文章就在討論這一點。如有謬誤請見諒并聯系指出，本文遵守CC 4.0 BY-SA版權協議，轉載請聯系作者并注明出處，謝謝。

關鍵字：尺度擴展定律、推理時計算擴展
發表單位：Google DeepMind

$\nabla$ 聯系方式：

e-mail: FesianXu@gmail.com

github: https://github.com/FesianXu

知乎專欄: 計算機視覺/計算機圖形理論與應用

微信公眾號：機器學習雜貨鋪3號店

在大語言模型（Large Language Model, LLM）中，存在所謂的尺度擴展規律（Scaling Laws） [2]，如Fig 1所示，即是：

LLM的性能會隨著模型的參數量、模型的訓練量、模型的訓練數據量的增加而增加

fig1_llm_scaling_laws

Fig 1. 大模型中的尺度擴展規律，測試集損失隨著模型訓練量、訓練集數據量、模型參數量的增加而遞減（即是模型性能遞增）。

然而，也有一系列工作告訴我們，在保持LLM不變的情況下，在推理時增加計算量，如以下方法，可以有效地提高模型效果

增加并行搜索計算量 ：模型采樣多個答案（有多種采樣方法，Best-of-N、Beam Search、Lookahead Search等），通過獎勵模型從中挑選一個最佳答案。
增加串行修改計算量 ：模型給出一個初始答案，然后基于這個初始答案進行修正，以此類推直到得到一個最佳答案未知。

第一種方法也可以稱之為是通過并行搜索（Search）的方法，第二種則是串行修正模型的輸出Token概率分布的方法，文章中提到這兩種不同方法，適用于不同的問題類型。那么對于一個特定的問題，在給定了有限計算預算的前提下，應該如何分配預訓練計算量、推理時計算量呢？這兩種計算量是否可以相互兌換呢（Exchange）？這正是本文想要討論的，如何分配計算預算，特別是聚焦在如何分配推理時的計算預算。

如公式(1)所示，用 $\theta^{*}_{q, y^{*}(q)}(N)$ 表示在給定了輸入prompt為 $q$ ，推理計算預算 $N$ 的情況下的最佳推理策略， $\mathrm{Target}(\theta,N,q)$ 表示在給定推理策略 $\theta$ 、計算預算和prompt情況下的輸出分布。那么公式(1)就是表示，在給定計算預算 $N$ 的情況下，最大化模型效果（此處 $y=y^{*}(q)$ 表示一個精準預測）得到的推理策略。此處的推理策略 $\theta$ ，不妨簡單理解為是并行計算和串行計算的比例以及具體的一些超參數，比如beam width等。
$\theta^{*}_{q, y^{*}(q)}(N) = \arg \max_{\theta} (\mathbb{E}_{y \sim \mathrm{Target}(\theta,N,q)} \big [ \mathbb{I}_{y=y^{*}(q)} \big ]) \tag{1}$
顯然這個公式還是過于抽象，無法將問題細化。作者認為，給模型的問題難度（Difficulty）可以作為公式(1)的有效代理，即是假如問題難度分為離散的5個檔次，那么只要求出不同難度檔次下的最優推理策略，就是對公式(1)的有效擬合。那么問題就變成了如何定義一個問題的難度（注意到本文討論的問題都是數學問題，也就是有明確答案的），作者采用的是基于『后驗』的方法，也就是用一個基線LLM模型對問題進行2048次采樣，計算這些答案的一次通過率（pass@1），越高的一次通過率說明問題越簡單，越低的一次通過率則說明問題越難。

整個實驗是在數據數據集MATH [3] 基準上進行的，其中包含了不同難度的高中級別數學問題，所有的實驗中都采用分割出的12k個數據進行訓練，500個數據進行測試。本文采用的Verifier（用于給答案打分，判斷答案的好壞程度）是PRM（Process Reward Model），其能夠對結構化答案中每一個步驟進行打分，如Fig 2所示，比起ORM（Output Reward Model）只對整個答案粒度進行打分，是一種更為細粒度的建模方式，具體可見 [4]。本文采用的PRM的具體訓練方法請參考原論文 [1]，本博文不做詳細介紹。

fig2_prm_step_by_step_scoring

Fig 2. PRM可以對答案的每一個步驟進行檢驗，比起只能對最終答案進行檢驗的ORM，其打分粒度更細。其中每一行是一個解答步驟，用"\n"隔開，綠色背景越深，代表一個更高的PRM打分，紅色背景越深則表示一個低的PRM打分。圖來自 [4]。

基于并行搜索的方法

我們先來看到基于并行搜索的方法，本文探索了如Fig 3所示的幾種搜索方法：

最佳N選1（Best-of-N）：對問題進行多次采樣，得到 $N$ 個完整的回答，然后對所有采樣得到的答案過一遍PRM判斷答案的質量，然后將PRM打分最高的答案保留作為最終答案。考慮到我們的測試數據集是數學問題，可以認為是有明確的數字答案的，因此本文采用的是加權最佳N選1（Best-of-N Weighted）策略。不難看出，最佳N選1的計算預算即是 $N$ 。
束搜索（Beam Search）：一個問題的回答可以分為多個步驟（Step），在給定了固定數量的束數（The number of beams） $N$ 和束寬（beam width） $M$ 后，如Fig 3的中圖所示：
- 第一步： $t_0$ 時刻，首先采樣 $N$ 個答案步驟作為初始化。
- 第二步：對于這 $N$ 個步驟進行PRM打分，保留其打分排序前 $R = N / M$ 的步驟。
- 第三步： $t_1$ 時刻，對于這 $R$ 個步驟，每一個繼續采樣 $M$ 個步驟，這樣我們 $t_1$ 時刻就有了 $N/M \times M = N$ 個答案步驟，和初始化時候的數量一樣。
- 第四步： $t_2$ 時刻，重復第二步到第四步。
不難發現，束搜索的每一步采樣我們都固定采用 $N$ 個答案步驟，因此最終計算預算和最佳N選1是類似的，都是 $N$ 。
前瞻搜索（Lookahead Search）：是束搜索的『升級版』，每一輪采樣都采樣 $K$ 步前瞻（在束搜索中 $K = 0$ ），然后在通過RPM判斷保留哪些采樣軌跡，不難看出前瞻搜索的計算量通常都比較大，通常記作 $\times (K+1)$ 。

fig2_test_time_scaling_via_search

Fig 3. 對比不同的 PRM 搜索方法。左圖：最佳N選1。中間圖：束搜索。右圖：前瞻搜索。

讓我們來看到試驗結果，束搜索和前瞻搜索都有推理超參數需要選擇：

束搜索：其束寬 $M$ 需要設定，本文采用了兩種設定 $\sqrt{N}$ 和固定的 $M = 4$ 。
前瞻搜索：其前瞻 $K$ 和束寬 $M$ 需要設定，本文采用了三種設定：
- $\sqrt{N}$
- $\sqrt{N}$
- $K = 3, M = 4$

整個試驗結果如Fig 4所示，從左圖可以發現，在計算預算短缺的情況下，束搜索明顯表現最佳，然而當預算增長到 $N = 64$ 以后，最優N選1策略和前瞻搜索就開始趕上了束搜索，最終在高計算預算情況下，最優N選1策略效果拔得頭籌。從右圖來看，在不同的問題難度下，這三種搜索方法有著不同的優勢。在簡單問題上（問題難度1和2），當預算不足時候效果最好的是束搜索策略，隨著預算增加，束搜索出現了性能下降的情況，這個可以認為是過擬合（也即是過度優化），再簡單問題且有著充分預算的情況下，能看到最佳N選1策略的效果是最好的。

然而在困難問題上（問題難度3以上），我們發現幾乎在所有預算下，都是束搜索策略的表現明顯最好，在最難的問題（難度5）上，甚至只有束搜索還能產生有意義的結果。這個結論是很符合直觀印象的，對于簡單問題，模型一次性產生正確答案的幾率比較高，因此通過最優N選1或者多數投票的策略都能獲得很不錯的效果，但是遇到困難問題的時候，一個復雜問題通常會有著多個復雜的步驟，祈求模型一次性產生正確答案的幾率就很低了，束搜索能夠以每一步中最正確的步驟作為種子繼續搜索，因此解決復雜問題的能力就更強，當然這也需要更多的算力預算了。

fig3_diff_methods_for_searching_against_PRM_verifiers

Fig 4. 左圖：比較針對 PRM 驗證器進行搜索的不同方法。我們發現，在低生成預算情況下，束搜索（beam search）表現最佳，但隨著預算的進一步增加，改進效果逐漸減弱，并且低于最佳 N 選一（best-of-N）基線。前瞻搜索（lookahead-search）在相同生成預算下通常表現不如其他方法。右圖：按難度級別對束搜索和最佳 N 選一進行比較。每個難度組中的四根柱狀圖分別對應不斷增加的測試時計算預算（4、16、64 和 256 次生成）。在難度較低的問題（第 1 和第 2 組）上，隨著預算的增加，束搜索出現了過度優化的跡象，而最佳 N 選一則沒有這種情況。在中等難度的問題（第 3 和第 4 組）上，束搜索相較于最佳 N 選一展現出了一致的改進效果。

從上面的實驗中，我們發現不同難度的問題需要不同的搜索策略，并且對計算預算的需求也不同，那么從以上試驗得到的“計算最優測試時計算分配設置”（Compute-optimal test-time compute strategy）對比基線效果究竟有多大的優勢呢？從Fig 5中我們發現，幾乎在所有預算下，都是計算最優策略的效果顯著最佳，并且對比 $N = 16$ 和 $N = 64$ ，我們發現計算最優的配置幾乎可以用基線最優N選1策略的25%計算預算達到相同的效果。

fig5_compute_optimal_trends

Fig 5. 將計算最優的測試時計算分配與 PRM 搜索的基線進行對比。通過根據問題難度的概念來調整測試時計算，我們發現使用最多少 4 倍的測試時計算（例如 16 次生成對比 64 次生成）幾乎可以超越 PRM 的最佳 N 選一（best-of-N）。“計算最優 oracle”指的是使用從真實正確性信息中得出的 oracle 難度分組，“計算最優預測”指的是使用 PRM 的預測來生成難度分組。請注意，使用這兩種難度分組的曲線在很大程度上彼此重疊。

總的來說，我們有以下結論，直接拿走不謝。

我們發現，任何給定驗證器搜索方法的有效性在很大程度上取決于計算預算和具體問題。具體而言，束搜索在處理較難問題以及在較低計算預算下更為有效，而最佳 N 選一方法在處理較易問題以及在較高預算下則更為有效。此外，通過針對特定問題難度和測試時計算預算選擇最佳搜索設置，我們幾乎可以使用少至 4 倍的測試時計算量來超越最佳 N 選一方法。

基于串行修正的方法

以上是基于并行搜索的測試時提升方法，還有一系列的工作是基于串行修正的，做法也非常直觀，讓LLM去修改自己之前生成的答案，期望能將之前答案的錯誤部分修正，然后進行多次迭代直到獲得最終的正確答案，這個過程如Fig 6左圖的"Sequential Revisions"所示。基于現有的LLM，通過簡單的prompt去修正自己的錯誤答案在推理型任務中被認為是低效的，因此需要微調出一個修正模型（Revision model），這個過程在本文就不介紹了，有興趣的讀者請自行翻閱原論文。

fig6_sequential_revision_vs_paralle_search

Fig 6. 并行搜索（例如，最佳 N 選一）與串行修訂的對比。左圖：并行采樣并行獨立生成 N 個答案，而順序修訂則根據之前的嘗試依次生成每個答案。右圖：無論是順序采樣還是并行采樣，我們都可以使用驗證器來確定最佳 N 選一的答案（例如，通過應用最佳 N 選一加權）。我們還可以將部分預算分配給并行采樣，部分分配給順序采樣，從而有效地結合這兩種采樣策略。在這種情況下，我們使用驗證器首先在每個順序鏈中選擇最佳答案，然后在不同鏈之間選擇最佳答案。

雖然這個能夠修正自己錯誤的LLM是至多基于前4個答案進行修正訓練的，在測試時候卻可以生成長度更長的修改鏈（很簡單，只需要保留最近的4次修正答案作為模型的上文就行了，之前的修正答案可以截斷掉），如Fig 7左圖所示，嘗試更多次的修正能夠獲得更好的pass@1指標。在給定了 $N$ 的計算預算前提下，我們既可以進行并行采樣 $N$ 次選取最佳答案（并行搜索策略），也可以串行修改 $N$ 次選取其中的最佳答案，那么哪種策略的效果更好呢？如Fig 7右圖所示，作者發現不管是基于PRM模型校驗的結果（也就是最佳N選1的結果），亦或是基于多數服從的結果（也就是投票結果），在不同的計算預算下，基本都是串行修正方案下的效果更有優勢。

fig7_revison_model_vs_sequential

Fig 7. 左圖：我們修訂模型在每個修訂步驟的首次通過率（pass@1）。在每個修訂步驟之后，首次通過率逐漸提高，甚至超過了它訓練時的 4 個修訂步驟。我們通過在測試集中對每個問題的 4 個長度為 64 的修訂軌跡的性能進行平均，來估計每個步驟的首次通過率。右圖：修訂模型的順序采樣與并行采樣對比。比較從我們的修訂模型并行生成 N 個初始答案與使用模型順序生成 N 個修訂答案時的性能。當使用驗證器和多數投票來選擇答案時，我們發現使用修訂模型順序生成答案的表現略優于并行生成答案。

我們現在有串行修正和并行搜索兩大類方法可以使用計算預算 $N$ ，如Fig 6右圖的"Combining Sequential/Parallel"所示，如何在這兩大類方法中分配計算預算呢？從直觀上看，串行修正是對答案的局部錯誤進行修改，而并行搜索則是從答案整體正確性的角度進行搜索，是一種更為全局的方法。作者用 $\mathrm{N(串行修正)}/\mathrm{N(并行搜索)}$ 表示這兩種方法的預算分配，如Fig 8左圖所示，我們發現增加串行修正計算預算的比例能更容易帶來性能提升，然而也存在一個理想比例達到性能的峰值，大概是16左右。

我們再看到不同難度問題下，預算分配的規律。如Fig 8右圖所示，在簡單問題下，提高串行修正計算預算比例總是能帶來性能增長，然而在較難問題上，則需要一個合適的比例，意味著需要引入一定量的并行搜索才能得到更好的性能。這個現象也能解釋我們之前的結論，簡單問題就算LLM不能一次性給出完整的正確答案，但是能給出基本正確的答案，只需要少量的串行修正就能得到完全正確答案。但是較復雜的問題LLM很難一次性就給出基本正確的答案，而是需要多次搜索后才能找到基本正確的答案，因此需要引入串行搜索的計算。

fig8_vary_ratio_of_seq_parallel

Fig 8. 左圖：改變分配給串行修訂與并行搜索的生成預算比例。每條線代表在比例變化時一個固定的生成預算。我們使用驗證器進行答案選擇。我們發現，盡管增加串行修正往往比更多的并行搜索表現更佳，但在更高的生成預算下，存在一個理想比例，能夠在兩個極端之間取得平衡。右圖：對于 N = 128 的生成預算，在不同難度組中改變順序與并行的比例。使用基于驗證器的選擇，我們發現較簡單的問題在使用純串行修正時能夠達到最佳性能。而對于較難的問題，則存在一個串行與并行計算之間的理想比例。

作者驗證了在考慮了串行和并行的最優計算預算分配下，對于單獨并行計算的優勢，結果是能夠以純并行計算25%的預算達到相同的模型效果，具體結果就不展開了。總結以上的試驗，我們有結論如下，拿走不謝。

我們發現，在串行（例如修訂）和并行（例如標準的最佳 N 選一）測試時計算之間存在一種權衡，并且串行與并行測試時計算的理想比例在很大程度上取決于計算預算和具體問題。具體而言，較簡單的問題從純串行的測試時計算中受益，而較難的問題通常在串行與并行計算達到某個理想比例時表現最佳。此外，通過針對給定問題難度和測試時計算預算最優地選擇最佳設置，我們能夠使用少至 4 倍的測試時計算量來超越并行的最佳 N 選一基線。

以上的試驗讓我們看到了，在測試階段進行復雜推理策略帶來的模型性能提升，我們不可避免有一個疑問：增加測試階段的計算預算，是否能替代預訓練階段的計算預算呢？也就是是否能通過復雜的測試策略，從而在減少模型預訓練的情況下，提升模型性能。作者最后在本文也進行了試驗，不過我們就不繼續詳細討論了，作者的結論是：

測試時計算與預訓練計算并非一對一“可互換”。對于模型能力范圍內的簡單和中等難度問題，或者在推理（實時性）要求較低的情況下，測試時計算可以輕松彌補額外的預訓練。然而，對于超出基礎模型能力范圍的具有挑戰性的問題，或者在推理（實時性）要求較高的情況下，預訓練可能更有效于提升性能。

筆者讀后感：

無論是學術界還是工業界，測試時尺度擴展都是當前的研究熱點，這篇論文的信息量很大，作者做了很多有價值的試驗，去驗證擴展不同測試策略的計算量帶來的性能提升，筆者將其理解為LLM測試時的scaling law，同時也探索了預訓練階段和測試階段的scaling law，并且說明了預訓練在困難問題下是具有不可替代性的。不過本文的試驗都是在數學類的問題上進行試驗的，結論是否可以泛化到其他問題（比如問答類問題、代碼型問題等等），是一個值得繼續探索的開放研究問題。同時，作者本文的試驗沒有考慮交織訓練和測試，也就是復雜推理策略輸出的答案，可以回饋LLM進行進一步的訓練從而提升模型效果。這些都是可以進一步探索的工作。

Reference

[1]. Snell, Charlie, Jaehoon Lee, Kelvin Xu, and Aviral Kumar. “Scaling llm test-time compute optimally can be more effective than scaling model parameters.” arXiv preprint arXiv:2408.03314 (2024).

[2]. Kaplan, Jared, Sam McCandlish, Tom Henighan, Tom B. Brown, Benjamin Chess, Rewon Child, Scott Gray, Alec Radford, Jeffrey Wu, and Dario Amodei. “Scaling laws for neural language models.” arXiv preprint arXiv:2001.08361 (2020).

[3]. Hendrycks, Dan, Collin Burns, Saurav Kadavath, Akul Arora, Steven Basart, Eric Tang, Dawn Song, and Jacob Steinhardt. “Measuring mathematical problem solving with the math dataset.” arXiv preprint arXiv:2103.03874 (2021). aka MATH

[4]. Lightman, Hunter, Vineet Kosaraju, Yura Burda, Harri Edwards, Bowen Baker, Teddy Lee, Jan Leike, John Schulman, Ilya Sutskever, and Karl Cobbe. “Let’s verify step by step.” arXiv preprint arXiv:2305.20050 (2023).