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

LLM的性能會隨著模型的參數(shù)量、模型的訓練量、模型的訓練數(shù)據(jù)量的增加而增加

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

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

增加并行搜索計算量 ：模型采樣多個答案（有多種采樣方法，Best-of-N、Beam Search、Lookahead Search等），通過獎勵模型從中挑選一個最佳答案。增加串行修改計算量 ：模型給出一個初始答案，然后基于這個初始答案進行修正，以此類推直到得到一個最佳答案未知。第一種方法也可以稱之為是通過并行搜索（Search）的方法，第二種則是串行修正模型的輸出Token概率分布的方法，文章中提到這兩種不同方法，適用于不同的問題類型。那么對于一個特定的問題，在給定了有限計算預(yù)算的前提下，應(yīng)該如何分配預(yù)訓練計算量、推理時計算量呢？這兩種計算量是否可以相互兌換呢（Exchange）？這正是本文想要討論的，如何分配計算預(yù)算，特別是聚焦在如何分配推理時的計算預(yù)算。

如公式(1)所示，用表示在給定了輸入prompt為q，推理計算預(yù)算N NN的情況下的最佳推理策略，表示在給定推理策略、計算預(yù)算和prompt情況下的輸出分布。那么公式(1)就是表示，在給定計算預(yù)算N的情況下，最大化模型效果（此處表示一個精準預(yù)測）得到的推理策略。此處的推理策略，不妨簡單理解為是并行計算和串行計算的比例以及具體的一些超參數(shù)，比如beam width等。

顯然這個公式還是過于抽象，無法將問題細化。作者認為，給模型的問題難度（Difficulty）可以作為公式(1)的有效代理，即是假如問題難度分為離散的5個檔次，那么只要求出不同難度檔次下的最優(yōu)推理策略，就是對公式(1)的有效擬合。那么問題就變成了如何定義一個問題的難度（注意到本文討論的問題都是數(shù)學問題，也就是有明確答案的），作者采用的是基于『后驗』的方法，也就是用一個基線LLM模型對問題進行2048次采樣，計算這些答案的一次通過率（pass@1），越高的一次通過率說明問題越簡單，越低的一次通過率則說明問題越難。

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

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

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

最佳N選1（Best-of-N）：對問題進行多次采樣，得到N個完整的回答，然后對所有采樣得到的答案過一遍PRM判斷答案的質(zhì)量，然后將PRM打分最高的答案保留作為最終答案?？紤]到我們的測試數(shù)據(jù)集是數(shù)學問題，可以認為是有明確的數(shù)字答案的，因此本文采用的是加權(quán)最佳N選1（Best-of-N Weighted）策略。不難看出，最佳N選1的計算預(yù)算即是N。

束搜索（Beam Search）：一個問題的回答可以分為多個步驟（Step），在給定了固定數(shù)量的束數(shù)（The number of beams） N和束寬（beam width） M后，如Fig 3的中圖所示：

第一步：時刻，首先采樣N個答案步驟作為初始化。第二步：對于這N個步驟進行PRM打分，保留其打分排序前的步驟。第三步：時刻，對于這R個步驟，每一個繼續(xù)采樣M個步驟，這樣我們時刻就有了個答案步驟，和初始化時候的數(shù)量一樣。第四步：時刻，重復(fù)第二步到第四步。不難發(fā)現(xiàn)，束搜索的每一步采樣我們都固定采用N NN個答案步驟，因此最終計算預(yù)算和最佳N選1是類似的，都是N。

前瞻搜索（Lookahead Search）：是束搜索的『升級版』，每一輪采樣都采樣K步前瞻（在束搜索中K=0），然后在通過RPM判斷保留哪些采樣軌跡，不難看出前瞻搜索的計算量通常都比較大，通常記作。

Fig 3. 對比不同的 PRM 搜索方法。左圖：最佳N選1。中間圖：束搜索。右圖：前瞻搜索。讓我們來看到試驗結(jié)果，束搜索和前瞻搜索都有推理超參數(shù)需要選擇：

束搜索：其束寬M需要設(shè)定，本文采用了兩種設(shè)定和固定的M = 4。

前瞻搜索：其前瞻K和束寬M需要設(shè)定，本文采用了三種設(shè)定

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

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

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

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

Fig 5. 將計算最優(yōu)的測試時計算分配與 PRM 搜索的基線進行對比。通過根據(jù)問題難度的概念來調(diào)整測試時計算，我們發(fā)現(xiàn)使用最多少 4 倍的測試時計算（例如 16 次生成對比 64 次生成）幾乎可以超越 PRM 的最佳 N 選一（best-of-N）。“計算最優(yōu) oracle”指的是使用從真實正確性信息中得出的 oracle 難度分組，“計算最優(yōu)預(yù)測”指的是使用 PRM 的預(yù)測來生成難度分組。請注意，使用這兩種難度分組的曲線在很大程度上彼此重疊?？偟膩碚f，我們有以下結(jié)論，直接拿走不謝。

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

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

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

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

Fig 7. 左圖：我們修訂模型在每個修訂步驟的首次通過率（pass@1）。在每個修訂步驟之后，首次通過率逐漸提高，甚至超過了它訓練時的 4 個修訂步驟。我們通過在測試集中對每個問題的 4 個長度為 64 的修訂軌跡的性能進行平均，來估計每個步驟的首次通過率。 右圖：修訂模型的順序采樣與并行采樣對比。比較從我們的修訂模型并行生成 N 個初始答案與使用模型順序生成 N 個修訂答案時的性能。當使用驗證器和多數(shù)投票來選擇答案時，我們發(fā)現(xiàn)使用修訂模型順序生成答案的表現(xiàn)略優(yōu)于并行生成答案。我們現(xiàn)在有串行修正和并行搜索兩大類方法可以使用計算預(yù)算N，如Fig 6右圖的"Combining Sequential/Parallel"所示， 如何在這兩大類方法中分配計算預(yù)算呢？從直觀上看，串行修正是對答案的局部錯誤進行修改，而并行搜索則是從答案整體正確性的角度進行搜索，是一種更為全局的方法。作者用 表示這兩種方法的預(yù)算分配，如Fig 8左圖所示，我們發(fā)現(xiàn)增加串行修正計算預(yù)算的比例能更容易帶來性能提升，然而也存在一個理想比例達到性能的峰值，大概是16左右。

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

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

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

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

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

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

筆者讀后感：

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