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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

筆者讀后感：

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

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).