優(yōu)化大模型的測試時計算是提升模型部署效率和節(jié)省計算資源的關(guān)鍵一環(huán)。前段時間，黃仁勛在 CES 2025 的演講中把測試時 Scaling 形容為大模型發(fā)展的三條曲線之一。如何優(yōu)化測試時計算成為業(yè)界關(guān)注的重要課題。

到目前為止，改進大型語言模型 (LLM) 的主要策略是使用越來越多的高質(zhì)量數(shù)據(jù)進行監(jiān)督微調(diào) (SFT) 或強化學(xué)習(xí) (RL)。

不幸的是，這種擴展形式似乎很快就會遇到瓶頸，預(yù)訓(xùn)練的擴展定律會趨于穩(wěn)定，有報告稱，用于訓(xùn)練的高質(zhì)量文本數(shù)據(jù)可能在 2028 年耗盡，特別是對于更困難的任務(wù)，例如解決推理問題，這似乎需要將當(dāng)前數(shù)據(jù)擴展約 100 倍才能看到任何顯著的改進。LLM 在這些困難任務(wù)中的問題上的當(dāng)前表現(xiàn)仍然不盡如人意。因此，迫切需要數(shù)據(jù)高效的方法來訓(xùn)練 LLM，這些方法可以超越數(shù)據(jù)擴展并解決更復(fù)雜的挑戰(zhàn)。

在這篇文章中，我們將討論這樣一種方法：通過改變 LLM 訓(xùn)練目標，我們可以重用現(xiàn)有數(shù)據(jù)以及更多的測試時計算來訓(xùn)練模型以做得更好。

博客地址：https://blog.ml.cmu.edu/2025/01/08/optimizing-llm-test-time-compute-involves-solving-a-meta-rl-problem/

圖 1：訓(xùn)練模型以優(yōu)化測試時計算并學(xué)習(xí)「如何發(fā)現(xiàn)」正確答案，而不是學(xué)習(xí)輸出「什么答案」的傳統(tǒng)學(xué)習(xí)范式。

當(dāng)前訓(xùn)練模型的主要原則是監(jiān)督它們?yōu)檩斎氘a(chǎn)生特定的輸出。例如，監(jiān)督微調(diào)嘗試匹配給定輸入的直接輸出 token，類似于模仿學(xué)習(xí)，而 RL 微調(diào)訓(xùn)練響應(yīng)以優(yōu)化獎勵函數(shù)，該函數(shù)通常應(yīng)該在 oracle 響應(yīng)上取最高值。無論哪種情況，我們都在訓(xùn)練模型以產(chǎn)生它可以表示的 y* 的最佳近似值。

從抽象上講，這種范式訓(xùn)練模型以產(chǎn)生單個輸入輸出映射，當(dāng)目標是直接解決給定分布中的一組類似查詢時，這種方法很有效，但無法發(fā)現(xiàn)分布外查詢的解決方案。固定的、一刀切的方法無法有效適應(yīng)任務(wù)的異質(zhì)性。我們更希望有一個強大的模型，它能夠通過嘗試多種方法并在不同程度上尋求信息，或者在無法完全解決問題時表達不確定性，從而推廣到新的、未見過的問題。

我們?nèi)绾斡?xùn)練模型來滿足這些要求？

學(xué)習(xí)「如何回答」

為了解決上述問題，一個新想法是允許模型在測試時使用計算資源來尋找「元（meta）」策略或算法，這些策略或算法可以幫助它們理解「如何」得出一個好的答案。

實現(xiàn)能夠賦予模型系統(tǒng)性程序運行能力的元策略，應(yīng)該能夠使其在測試時外推和泛化到不同復(fù)雜度的輸入查詢。例如，如果一個模型被教了柯西 - 施瓦茨不等式的含義，它就應(yīng)該能夠在簡單和困難的證明問題上在適當(dāng)?shù)臅r候運用它。換句話說，給定一個測試查詢，我們希望模型能夠執(zhí)行包含多個基本推理步驟的策略，這可能需要消耗更多的 token。

圖 2 展示了解決一個給定問題的兩種不同策略的例子。我們?nèi)绾斡?xùn)練模型來做到這一點呢？我們將把這個目標形式化為一個學(xué)習(xí)問題，并通過元強化學(xué)習(xí)的思路來解決它。

圖 2:  展示了兩種算法及其各自生成的 token 流示例。

將學(xué)習(xí)「如何做」形式化為一個目標

對于每個問題 x∈X，假設(shè)我們有一個獎勵函數(shù) r (x,?):Y?{0,1}, 可以針對任何輸出 token 流 y 進行查詢。例如，對于數(shù)學(xué)推理問題 x，其 token 輸出流為 y，獎勵 r (x,y) 可以是檢查某個 token 子序列是否包含正確答案的函數(shù)。我們只獲得了訓(xùn)練問題數(shù)據(jù)集 D_train, 因此也只有獎勵函數(shù)集合 {r (x,?):x∈D_train}。我們的目標是在事先未知的測試問題分布 P_test 上獲得高獎勵。測試問題的難度可能與訓(xùn)練問題不同。

對于未知的測試問題分布 P_test 和有限的測試時計算預(yù)算 C，我們可以從訓(xùn)練問題數(shù)據(jù)集 D_train 中學(xué)習(xí)一個算法 A∈A_C (D_train), 這個算法屬于推理計算受限的測試時算法類 A_C。這個類中的每個算法都以問題 x～P_test 作為輸入，并輸出一個 token 流。

在圖 2 中，我們給出了一些例子來幫助理解這個 token 流可以是什么。例如，A_θ(x) 可能首先包含針對問題 x 的某些嘗試 token，然后是一些用于預(yù)測嘗試正確性的驗證 token，如果驗證為不正確，接著是對初始嘗試的一些改進，所有這些都以「線性」方式串聯(lián)在一起。另一個算法 A_θ(x) 可能是以線性方式模擬某種啟發(fā)式引導(dǎo)搜索算法。算法類 A_C (D_train) 將由上述所有可能的 A_θ(x) 產(chǎn)生的下一個 token 分布組成。注意，在這些例子中，我們希望使用更多的 token 來學(xué)習(xí)一個通用但可泛化的程序，而不是去猜測問題 x 的解決方案。

我們的學(xué)習(xí)目標是學(xué)習(xí)由自回歸大語言模型參數(shù)化的 A_θ(x)。我們將這整個流 (包括最終答案) 稱為響應(yīng) y～A_θ(x)。算法 A_θ(x) 的效用由獎勵 r (x,y) 衡量的平均正確性給出。因此，我們可以將學(xué)習(xí)算法表述為解決以下優(yōu)化問題：

將 (Op-How) 解釋為元強化學(xué)習(xí)問題

接下來的問題是：我們?nèi)绾谓鉀Q由語言模型參數(shù)化的、在計算受限算法類 A_c 上的優(yōu)化問題 (Op-How)？

顯然，我們既不知道測試問題的結(jié)果，也沒有任何監(jiān)督信息。因此，計算外部期望是徒勞的。對問題 x 猜測最佳可能響應(yīng)的「標準」大語言模型策略似乎也不是最優(yōu)的，因為如果充分利用計算預(yù)算 C，它可能會做得更好。

主要思路是優(yōu)化 (Op-How) 的算法 A_θ(x)∈A_c，類似于強化學(xué)習(xí)中的自適應(yīng)策略，它使用額外的 token 預(yù)算來實現(xiàn)某種算法策略來解決輸入問題 x。有了這個聯(lián)系，我們可以從類似問題通常的解決方式中獲得啟發(fā)：通過元學(xué)習(xí)的視角來看待 (Op-How)，具體來說是元強化學(xué)習(xí)：「元」是因為我們希望學(xué)習(xí)算法而不是直接回答給定問題，而「強化學(xué)習(xí)」是因為 (Op-How) 是一個獎勵最大化問題。

一個非常簡短的元強化學(xué)習(xí)入門

通常，強化學(xué)習(xí)訓(xùn)練一個策略來最大化馬爾可夫決策過程 (MDP) 中的給定獎勵函數(shù)。相比之下，元強化學(xué)習(xí)問題設(shè)定假設(shè)可以訪問任務(wù)分布 (每個任務(wù)都有不同的獎勵函數(shù)和動態(tài)特性)。在這種設(shè)定下，目標是在來自訓(xùn)練分布的任務(wù)上訓(xùn)練策略，使其能夠在從相同或不同測試分布抽取的測試任務(wù)上表現(xiàn)良好。

此外，這種設(shè)定不是根據(jù)策略在測試任務(wù)上的零樣本表現(xiàn)來評估它，而是讓它通過在測試時執(zhí)行幾個「訓(xùn)練」回合來適應(yīng)測試任務(wù)，在執(zhí)行這些回合之后再評估策略。

回到我們的設(shè)定，你可能會想知道馬爾可夫決策過程（MDP）和多個任務(wù) (用于元強化學(xué)習(xí)) 是如何體現(xiàn)的。每個問題 x∈X 都會引發(fā)一個新的強化學(xué)習(xí)任務(wù)，這個任務(wù)被形式化為馬爾可夫決策過程 (MDP) M_x，其中問題 x 中的 token 集合作為初始狀態(tài)，我們的大語言模型 A_θ(x) 產(chǎn)生的每個 token 作為一個動作，而簡單的確定性動態(tài)則通過將新的 token ∈T 與到目前為止的 token 序列連接來定義。注意，所有 MDP 共享動作集合和狀態(tài)集合  ，這些對應(yīng)于詞匯表中可能出現(xiàn)的變長 token 序列。然而，每個 MDP M_x 都有一個不同的未知獎勵函數(shù)，由比較器 r (x,?) 給出。

那么解決 (Op-How) 就對應(yīng)著找到一個策略，該策略能夠在計算預(yù)算 C 內(nèi)快速適應(yīng)測試問題 (或測試狀態(tài)) 的分布。從認知 POMDP 的視角來看這種測試時泛化的概念是另一種方式，這是一個將在 M_x 族上學(xué)習(xí)策略視為部分觀察強化學(xué)習(xí)問題的構(gòu)造。這個視角提供了另一種激發(fā)自適應(yīng)策略和元強化學(xué)習(xí)需求的方式：對于那些有強化學(xué)習(xí)背景的人來說，解決 POMDP 等同于運行元強化學(xué)習(xí)這一點應(yīng)該不足為奇。因此，通過解決元強化學(xué)習(xí)目標，我們正在尋求這個認知 POMDP 的最優(yōu)策略并實現(xiàn)泛化。

適應(yīng)性策略如何適應(yīng)測試問題？

在元強化學(xué)習(xí)中，對于每個測試 MDP M_x，策略 A_θ 在通過 A_θ 生成最終響應(yīng)進行評估之前，可以通過消耗測試時計算來獲取信息。在元強化學(xué)習(xí)術(shù)語中，獲得的關(guān)于測試 MDP M_x 的信息可以被視為在測試問題 x 引發(fā)的 MDP 上收集「訓(xùn)練」回合的獎勵，然后再在測試回合上進行評估。注意，所有這些回合都是在模型部署后執(zhí)行的。因此，為了解決 (Op-How)，我們可以將來自 A_θ(x) 的整個 token 流視為分成幾個訓(xùn)練回合的流。為了優(yōu)化測試時計算，我們需要確保每個回合都能提供一些信息增益，以便在測試 MDP M_x 的后續(xù)回合中表現(xiàn)更好。如果沒有信息增益，那么學(xué)習(xí) A_θ(x) 就退化為一個標準的強化學(xué)習(xí)問題 —— 只是計算預(yù)算更高 —— 這樣就不清楚學(xué)習(xí)「如何做」是否有用。

可以獲得什么樣的信息？當(dāng)然，如果 token 流中涉及外部接口，我們可以獲得更多信息。但是，如果不涉及外部工具，我們是否在利用免費午餐？我們指出不是這種情況，在 token 流進行過程中不需要涉及外部工具也能獲得信息。流中的每個回合都可以有意義地增加更多信息，也就是說，我們可以將消耗更多的測試時計算視為從模型對最優(yōu)解的后驗近似 P (?|x,θ) 中采樣的一種方式，其中每個回合 (或輸出流中的 token) 都在改進這個近似。因此，顯式地調(diào)整先前生成的 token 可以提供一種計算可行的方式，用固定大小的大語言模型來表征這個后驗。

綜上所述，當(dāng)被視為元強化學(xué)習(xí)問題時，A (?|?) 成為一個歷史條件 (「自適應(yīng)」) 策略，通過在給定測試問題上花費最多 C 的計算來優(yōu)化獎勵 r。

圖 3：智能體 - 環(huán)境交互協(xié)議。圖源：https://arxiv.org/pdf/1611.02779

學(xué)習(xí)自適應(yīng)策略的元強化學(xué)習(xí)

圖 4：A_θ(x) 的響應(yīng)包括一串 token。

我們?nèi)绾谓鉀Q這樣一個元強化學(xué)習(xí)問題？也許解決元強化學(xué)習(xí)問題最明顯的方法是采用黑盒元強化學(xué)習(xí)方法。這將涉及最大化輸出軌跡 A_θ(x) 中想象的「episodes」的獎勵總和。例如，如果 A_θ(x) 對應(yīng)于使用自我糾正策略，那么每個 episode 的獎勵將對軌跡中出現(xiàn)的單個響應(yīng)進行評分。如果 A_θ(x) 規(guī)定了一種在生成和生成驗證之間交替的策略，那么獎勵將對應(yīng)于生成和驗證的成功。然后我們可以優(yōu)化：

一般情況下，輸出 token 流可能無法清晰地分成生成和驗證片段。在這種情況下，可以考慮元強化學(xué)習(xí)問題的更抽象形式，直接使用信息增益的某種估計作為獎勵。

可以通過多輪強化學(xué)習(xí)方法來解決 (Obj-1) 和 (Obj-2)。實際上，只要能夠使用某種執(zhí)行定期在線策略采樣的強化學(xué)習(xí)算法來解決優(yōu)化問題，強化學(xué)習(xí)方法的選擇 (基于價值還是基于策略) 可能并不重要。

我們還可以考慮另一種制定元強化學(xué)習(xí)訓(xùn)練目標的方法：只優(yōu)化測試回合獲得的獎勵，而不是訓(xùn)練回合的獎勵，從而避免量化信息增益的需要。