28年AGI撞上數據墻,以后全靠測試時計算?CMU詳解優化原理
2025年主導AI的將是第3代scaling law:測試時計算。
正如Michael Dell轉述所言:
第一代scaling:預訓練像讀大學
第二代scaling:微調像讀博士
第三代scaling:測試時計算就像深度思考和推理
近日,計算機強校CMU機器學習系,發表博客文章解釋LLM測試時計算優化問題,特別是涉及到的元強化學習(meta-RL)問題。
文章亮點如下:
- 監督學習只是在訓練模型「答案是什么」,只要學習「如何解答」,模型泛化性會更好。
- 學習「如何解答」類似于強化學習中的自適應策略,相當于找到一個策略,使其能夠在計算預算C內適應測試問題。
- 測試時訓練等價于原強化學習,這個視角提供了多方面的有價值見解:(1)優化測試時計算資源時,與信息增益相關的中間過程獎勵的作用;(2)模型崩潰和預訓練初始化在學習meta策略中的作用;以及(3)缺乏外部反饋的情況下,不對稱性如何成為測試時改進的驅動力。
數據Scaling,窮途末路
目前為止,改進大語言模型(LLM)的主要策略,是使用越來越多的高質量數據進行監督微調(SFT)或強化學習(RL)。
不幸的是,這種擴展方式似乎很快會遇到瓶頸,預訓練的擴展法則趨于平穩。
并且有報告稱,到2028年,用于訓練的高質量文本數據可能會耗盡。
因此,迫切需要數據高效的方法來訓練LLM,這些方法超越了數據擴展(data scaling),并且能解決更加復雜的問題。
當前的LLM是訓練「答案是什么」
目前訓練模型的主導原則是監督它們為輸入生成特定的輸出。
例如,給個輸入,監督微調試圖匹配直接輸出的token,類似于模仿學習;而RL微調則訓練響應以優化獎勵函數,該函數通常假設在oracle響應上取最大值。
在這兩種情況下,都是在訓練模型生成它可以表示的最佳近似值y*。
抽象地說,這種范式訓練模型以生成單一的輸入輸出映射。
當目標是直接解決一組來自給定分布的相似查詢時,這種方法效果很好,但無法發現超出分布的查詢的解決方案。
固定的、一刀切的方法無法有效適應任務的異質性。
相反,需要的是一種穩健的模型,它能夠嘗試不種方法,在不同程度上尋求信息,或在完全無法完全解決問題時表達不確定性,從而概括出新的、未見過的問題。
該如何訓練模型來滿足這些要求呢?
學習「如何解答」,泛化性會更好
為了解決上述問題,需要新的理念:在測試時,允許模型通過計算來尋找「元」(meta)策略或算法,幫助其理解如何得出更好的答案。
實施這些元策略,模型可以系統化地推理,在面對不同復雜度的輸入時,也可以做到外推和泛化(extrapolation and generalization)。
請參見下圖2,了解兩種不同的策略如何解決特定問題。
圖2:兩種算法的示例及每種算法生成的token流。包括從模型權重中,獲取相關信息、規劃證明大綱、驗證中間結果以及必要時修正的token。
第一種算法(左)生成初始答案并驗證其正確性,如有必要,修正錯誤步驟。
第二種算法(右)一次性生成多個解決策略,并按線性順序逐個執行這些策略,最后選擇最有效的策略。
如何訓練模型達到這一目標呢?
這一目標可以形式化為一個學習問題,并通過元強化學習(meta RL)中的概念來解決。
將「如何學習」作為目標
對于問題,強化學習需要一個獎勵函數
,而且獎勵函數可以查詢任何輸出token流y。
比如,在一個數學推理問題中,帶有標記輸出流y,獎勵可以用于檢查某些子序列的標記是否包含正確答案。
只要有訓練問題的數據集和相應的獎勵函數集
,就能利用更多的token,學習通用但具有泛化能力的過程,而不僅僅猜測問題的答案。
學習的目標是實現測試問題分布上的高額獎勵,而且這些測試問題服從一個未知的先驗分布(apriori)。
最終目的是對于未知的測試問題分布,在有限的測試時計算預算C條件下,從推理計算受限的算法類中,利用訓練問題數據集學習一個算法
。
算法類包含了所有可能的
所引起的下一個token分布。
學習目標是學習,也就是說一個自回歸LLM(參見圖1中的示例)。
將整個流包括最終答案(圖1中的綠色和藍灰色部分)稱為響應y,完全由學習到的LLM模型生成。
通過獎勵的平均正確性,來衡量算法
的效果。
因此,學習算法可被視為求解以下優化問題:
將優化問題解釋為元強化學習
接下來的問題是:如何在計算受限的算法類中,利用語言模型來求解優化問題?
顯然,對于測試問題,既不知道結果,也沒有任何監督信號。
因此,沒辦法計算(Op-How)問題中的外層的期望。
標準的LLM策略,隨便猜測一下可能最好的答案,也不是最佳策略,因為如果能充分利用計算預算C,可能會表現得更好。
主要思路是,優化(Op-How)的算法類似于強化學習中的自適應策略。
它使用額外的token預算來執行某種算法策略,從而解決輸入問題$$x$$(類似「上下文搜索」或「上下文探索」)。
通過這種聯系,可以借鑒解決類似問題的方法,也就是將(Op-How)視為元學習,尤其是元強化學習(meta RL)來處理:「元」(meta)表示目的是學習算法而非直接給出問題的答案;「強化學習」(RL)則表明(Op-How)是一個獎勵最大化問題。
元強化學習
通常,強化學習訓練一個策略,以最大化馬爾可夫決策過程(MDP)中的給定獎勵函數。
與此不同,元強化學習問題則假設能夠利用任務分布(這些任務擁有不同的獎勵函數和動態)。
在這種設定下,目標是通過訓練任務分布中的任務來學習策略,從而使得策略能夠在測試任務上表現良好,無論該測試任務是否來自原來的測試任務分布。
此外,這種設定不以策略在測試任務上的零樣本表現作為評估標準,而是允許策略在測試時通過執行幾個「訓練」回合來適應測試任務,并在這些回合結束后對其進行評估。
大多數元強化學習方法的差異在于適應過程的設計。例如,RL2通過上下文強化學習對適應過程進行參數化;MAML在測試時執行顯式的梯度更新;PEARL通過適應潛在變量來識別任務。
元強化學習的作用
你可能會想,馬爾可夫決策過程(MDP)和元強化學習需要的多個任務,從何而來?
每個問題x都會引發一個新的強化學習任務,形式化為一個馬爾可夫決策過程(MDP):初始狀態是問題x中的token集合,LLM生成的token表示
的動作,并且將新token與當前token序列串聯(concatenating)來定義平凡的確定性動態(trivial deterministic dynamics)。
需要注意的是,所有的MDP共享相同的動作集A,同時也共享狀態集S,這代表了詞匯表中可能的變長token序列。而且,每個馬爾可夫決策過程都有由比較器給出的不同的未知獎勵函數。
然后,求解(Op-How)就等同于找到一個策略,使其能夠在計算預算C內迅速適應測試問題(或測試狀態)的分布。
另一種看待測試時泛化的方式,是所謂的認識性POMDP(Epistemic POMDP)。它將從馬爾可夫決策過程Mx算法族中學習策略,被認為是部分可觀測強化學習問題。
從這個角度來看,可以進一步理解為何需要自適應策略和元強化學習:對于那些來自強化學習背景的人來說,解決POMDP等同于進行元強化學習。
因此,解決元強化學習,就是在尋找認識性POMDP的最優策略,從而實現泛化能力。
真的有用嗎?
既然元強化學習本身就非常困難,這種元強化學習視角有什么用?
作者認為,盡管元強化學習完全從頭學習策略很難,但對那些已經通過預訓練獲得豐富先驗知識的模型,用元強化學習對它們微調時,非常有效。
此外,上述的元強化學習問題可能呈現出特殊的結構(比如,已知且確定的動態,不同的初始狀態),從而可以開發出非通用但有用的元強化學習算法。
如何使自適應策略適應測試問題?
在元強化學習中,對于每個測試MDP Mx,策略通過在測試時利用計算資源來獲取信息,然后根據
生成的最終響應進行評估。
因此,為了解決(Op-How)問題,可以將策略的整個token流視為拆分成多個訓練階段。
為了優化測試時的計算資源,需要確保每個訓練階段提供某些信息增益,以便在測試MDP的后續階段表現得更好。
如果沒有信息增益,那么就會退化為一個標準的強化學習問題——只不過計算預算更高——并且也搞不清楚「如何學習」是否有用。
可以獲得什么信息?
當然,如果在token流中涉及外部接口,可能會獲得更多的信息。
然而,如果沒有涉及外部工具,是否可以享受「免費午餐」?
作者指出,不需要外部工具參與,信息仍然可以隨著token流的進展而獲得。
在流中的每個階段,都可能通過提升模型對真實獎勵函數r(x,?)的后驗信念,從而獲得更多有意義的信息(例如,通過單獨訓練的驗證器或策略本身進行自我驗證),并且因此獲得最優響應y?。
換句話說,更多的測試時計算,可以視為從模型逼近的后驗分布P(?∣x,θ)中采樣的方式,其中每個階段(或輸出流中的token)都在改進對后驗分布的逼近。
因此,明確地對先生成的token條件化,是用固定大小的LLM表示后驗的可計算方法。
這也意味著,即使沒有外部輸入,隨著生成更多的tokens,也期望下列互信息會增加:
在這種情況下,由于所有的監督來自本身,需要在生成和驗證之間存在不對稱性,才能讓驗證引發信息增益。
另一個想法是,當模型在訓練數據上欠擬合時,僅僅增加生成token的長度,也可能提供顯著的信息增益,因為計算資源的增加會提升模型的容量(參見下列文章的第2節)。
顯然還需要更多的工作來形式化這些論點,但已經有一些文章,表明自我改進可以隱式或顯式地利用這種不對稱性。
總結起來,當將優化問題(Op-how)視為一個元強化學習問題時,A(?|?)變成了一個歷史條件化的(「自適應的」)策略,通過在給定的測試問題上花費最多的計算量來優化獎勵r。
學習一個基于過去階段條件化的自適應策略,正是黑箱元強化學習方法的目標。
元強化學習也與學習如何探索緊密相關,事實上,可以將這些額外的token視為在探索特定問題的策略。
圖3:RL2中智能體與環境交互的過程
通過元強化學習學習自適應策略
解決元強化學習問題的最明顯的方法,可能是使用黑箱元強化學習方法,例如RL2。
這就需要最大化輸出軌跡(trace)中想象的「情節」的獎勵總和。
例如,如果對應于使用自我糾正策略,則每個階段的獎勵將根據軌跡中出現的個體響應進行評分。
如果指定一種交替生成和生成性驗證的策略,則獎勵將對應于生成和驗證成功的程度。可以進行下列優化:
其中,對應于響應的索引,這些響應標記了階段的結束,獎勵
則表示該階段的標量獎勵信號(例如,驗證段的驗證正確性,生成段的生成正確性,等等)。
此外,作者還優化了答案的最終正確性獎勵。請注意,這一公式規定了一個密集的、基于過程的獎勵(這不同于使用逐步過程獎勵模型(PRM),而是采用密集的額外獎勵(reward bonus);這種密集的額外獎勵與探索之間的關系可以在下列論文中找到)。
還可以通過顯式地添加損失項或者隱式地(例如,剪裁掉違反計算預算的模型生成的結果),選擇限制使用的計算上限C。
上述方法只針對生成和驗證。
然而,輸出的token流,通常無法清晰地劃分為生成和驗證。
在這種情況下,可以考慮一種更抽象的meta RL形式,使用某種信息增益的估計作為獎勵。
來自QuietSTaR論文中的度量標準可能就是這樣一個估計,盡管目前尚不清楚如何準確地定義該度量標準。
可以通過多輪RL方法解決(Obj-1)和(Obj-2),例如基于策略梯度的中間密集獎勵方法,或者基于演員-評論家(actor-critic)架構的方法(例如,先前的ArCHer工作)。
也許,只要能使用某種周期性的on-policy rollouts RL算法來解決優化問題,那RL方法(基于值還是基于策略)的選擇甚至都可能無關緊要。
還可以考慮另一種設計meta RL訓練目標的不同方法:只優化測試回合中獲得的獎勵(例如,最后一次嘗試的最終答案正確性),而不優化訓練回合,從而避免量化信息增益的需求。
作者認為,這將面臨優化非常稀疏監督信號的難題,尤其是在長軌跡的末端(在meta RL術語中,軌跡由多個推理段或多個「回合」組成);而密集獎勵應該能夠更好地解決這一問題。
文中也討論了面臨的其他問題。
作者介紹
值得一提的是,博文6位作者中有3位華人。
Yuxiao Qu,卡內基梅隆大學計算機科學學院機器學習系的一年級博士。在CMU之前,他在威斯康星大學麥迪遜分校計算機科學系獲得了學士學位。更早之前,他還在香港中文大學工作過一段時間。
Matthew Yang,是CMU機器學習系的碩士生。此前,他在滑鐵盧大學學習計算機科學和統計學。
Lunjun Zhang,是多倫多大學機器學習小組的一名計算機科學博士生。2024年,他在谷歌DeepMind實習,研究LLM。2021年至2024年,他在自動駕駛初創公司擔任研究員。更早之前,他在多倫多大學攻讀工程科學專業。