如何最佳化測試時計算？解決「元強化學習」問題
机器之心發表於2025-02-10
最佳化大模型的測試時計算是提升模型部署效率和節省計算資源的關鍵一環。前段時間，黃仁勳在 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) 衡量的平均正確性給出。因此，我們可以將學習演算法表述為解決以下最佳化問題：接下來的問題是：我們如何解決由語言模型引數化的、在計算受限演算法類 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) 規定了一種在生成和生成驗證之間交替的策略，那麼獎勵將對應於生成和驗證的成功。然後我們可以最佳化：一般情況下，輸出 token 流可能無法清晰地分成生成和驗證片段。在這種情況下，可以考慮元強化學習問題的更抽象形式，直接使用資訊增益的某種估計作為獎勵。可以透過多輪強化學習方法來解決 (Obj-1) 和 (Obj-2)。實際上，只要能夠使用某種執行定期線上策略取樣的強化學習演算法來解決最佳化問題，強化學習方法的選擇 (基於價值還是基於策略) 可能並不重要。我們還可以考慮另一種制定元強化學習訓練目標的方法：只最佳化測試回合獲得的獎勵，而不是訓練回合的獎勵，從而避免量化資訊增益的需要。