如何優化測試時計算？解決「元強化學習」問題

選自ml.cmu.edu

作者：Amrith Setlur、Yuxiao Qu 等

機器之心編譯

優化大模型的測試時計算是提升模型部署效率和節省計算資源的關鍵一環。前段時間，黃仁勳在 CES 2025 的演講中把測試時 Scaling 形容為大模型發展的三條曲線之一。如何優化測試時計算成為業界關注的重要課題。

到目前為止，改進大型語言模型 (LLM) 的主要策略是使用越來越多的高質量數據進行監督微調 (SFT) 或強化學習 (RL)。

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

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

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

圖 1：訓練模型以優化測試時計算並學習「如何發現」正確答案，而不是學習輸出「什麼答案」的傳統學

習範式。

當前訓練模型的主要原則是監督它們為輸入產生特定的輸出。例如，監督微調嘗試匹配給定輸入的直接輸出 token，類似於模仿學習，而 RL 微調訓練響應以優化獎勵函數，該函數通常應該在 oracle 響應上取最高值。無論哪種情況，我們都在訓練模型以產生它可以表示的 y* 的最佳近似值。

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

我們如何訓練模型來滿足這些要求？

學習「如何回答」

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

實現能夠賦予模型系統性程序運行能力的元策略，應該能夠使其在測試時外推和泛化到不同複雜度的輸入查詢。例如，如果一個模型被教了柯西 – 舒華斯不等式的含義，它就應該能夠在簡單和困難的證明問題上在適當的時候運用它。換句話說，給定一個測試查詢，我們希望模型能夠執行包含多個基本推理步驟的策略，這可能需要消耗更多的 token。

圖 2 展示瞭解決一個給定問題的兩種不同策略的例子。我們如何訓練模型來做到這一點呢？我們將把這個目標形式化為一個學習問題，並通過元強化學習的思路來解決它。

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

將學習「如何做」形式化為一個目標

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

對於未知的測試問題分佈 P_test 和有限的測試時計算預算 C，我們可以從訓練問題數據集 D_train 中學習一個算法 A∈A_C (D_train), 這個算法屬於推理計算受限的測試時算法類 A_C。這個類中的每個算法都以問題 x∼P_test 作為輸入，並輸出一個 token 流。

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

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

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

接下來的問題是：我們如何解決由語言模型參數化的、在計算受限算法類 A_c 上的優化問題 (Op-How)？

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

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

一個非常簡短的元強化學習入門

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

此外，這種設定不是根據策略在測試任務上的零樣本表現來評估它，而是讓它通過在測試時執行幾個「訓練」回合來適應測試任務，在執行這些回合之後再評估策略。

回到我們的設定，你可能會想知道馬爾可夫決策過程（MDP）和多個任務 (用於元強化學習) 是如何體現的。每個問題 x∈X 都會引發一個新的強化學習任務，這個任務被形式化為馬爾可夫決策過程 (MDP) M_x，其中問題 x 中的 token 集合作為初始狀態，我們的大語言模型 A_θ(x) 產生的每個 token 作為一個動作，而簡單的確定性動態則通過將新的 token ∈T 與到目前為止的 token 序列連接來定義。注意，所有 MDP 共享動作集合和狀態集合

，這些對應於詞彙表中可能出現的變長 token 序列。然而，每個 MDP M_x 都有一個不同的未知獎勵函數，由比較器 r (x,⋅) 給出。

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

適應性策略如何適應測試問題？

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

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

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

圖 3：智能體 – 環境交互協議。圖源：https://arxiv.org/pdf/1611.02779

學習自適應策略的元強化學習

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

我們如何解決這樣一個元強化學習問題？也許解決元強化學習問題最明顯的方法是採用黑盒元強化學習方法。這將涉及最大化輸出軌跡 A_θ(x) 中想像的「episodes」的獎勵總和。例如，如果 A_θ(x) 對應於使用自我糾正策略，那麼每個 episode 的獎勵將對軌跡中出現的單個響應進行評分。如果 A_θ(x) 規定了一種在生成和生成驗證之間交替的策略，那麼獎勵將對應於生成和驗證的成功。然後我們可以優化：