〈Stanford CS329A Self-Improving AI Agents〉-EP-Part 2 | Test-Time Compute Scaling

https://www.youtube.com/watch?v=-Ggc37xLj_Y

模型參數不動、也不做任何微調,光是把 inference 階段的算力加上去,同一顆模型就能變強。但這條路能拿回多少,取決於你有沒有辦法在一堆生成出來的答案裡,認出哪一個是對的。Stanford CS329A 第二講講的就是這兩句話,前半段在證明模型其實早就知道答案,後半段在講為什麼你常常拿不回來。

訓練有三個階段,inference 是第三條可以加的軸

LLM 的開發可以拆成三段。pretraining 歷史上最貴:跑好幾個月、用掉一大堆 GPU、吃掉上兆個資料 token。fine-tuning 通常輕得多,用得上的資料量跟算力都少了幾個數量級。inference 是拿到訓練好的模型之後真正去用它,過去的用法很省,使用者丟一個 query,模型回一個答案,一次來回就結束,每次 inference 幾乎不花算力。

新的 paradigm 換了另一種花法。同一個 agent 可以離線工作,生出很多 token,一遍一遍修自己的輸出。講者把這形容成一個新的旋鈕:訓練結束之後,你還有一個地方可以把能力往上加。上一堂課講過 pretraining 的 scaling law,test loss 對資料量、算力、參數量有一條可以預測的曲線,換 setup 方程式會變,但方向是可預測的;這一堂要問的是 inference 這一端有沒有同樣的東西。

同一題問一千次,小模型可以反超大模型

先回頭講 large language monkeys。無限猴子定理的版本是這樣:同一顆 LLM、同一個問題,不要只問一次,問十次、一百次,然後讓一個 verifier 指出哪一份回應是對的,那一份就當系統的輸出。

理由是模型的生成不是決定性的。Llama 3 的 8B 或 7B 單次作答打不贏 GPT-4o,但同一題問很多次、再從候選裡挑出對的那一個,小模型可以顯著變好,在一些任務上反超這些更大、更封閉的模型。講者的說法是,這些小模型其實已經知道答案,只是第一次不告訴你;重複取樣是把已經存在的答案挖出來。

這個做法在硬數學、程式、問答,以及 SWE-bench 這類 agentic benchmark 上都成立。SWE-bench 模擬的是軟體工程師改程式、生出 patch 的工作。圖上量的是 coverage,也就是至少有一個樣本解對的題目比例,橫軸是樣本數,從一到一千。講者說圖上那顆 DeepSeek(她記得是 V3)在一千個樣本之後,解掉的程式與軟體工程題目比 Claude 3.5 或 o1 Preview 還多。這件事在程式題上特別有用:如果有 unit test 可以指出哪一個樣本是對的,就得到一條端到端的自動化路徑,把一顆開源模型變成更強的系統,完全不用重新訓練。

Coverage 有自己的 scaling law,背後是一條很長的難題尾巴

pretraining 的 scaling law 用來預測 test loss,test-time compute 也可以定義一條類似的關係:coverage 跟平行取樣的樣本數之間是一條 exponential power law,係數 a、b 從實際曲線擬合出來。實驗跨了 Llama 3 的 8B、7B,以及 Gemma、Pythia 這些模型,參數從 7,000 萬到 700 億,也跨了好幾個任務領域,預測曲線大致貼著實際量到的 coverage 走,連 7,000 萬參數的模型都看得到同樣的行為。實務上的意思很直接:你可以先算要多少樣本、配多少資源,才達到你要的 coverage。

為什麼一堆單題的指數曲線合起來會變成 power law,講者也解釋了。單一題目 pass@1 的機率是 p,那 k 個樣本全部答錯的機率是 1 減 p 連乘 k 次,至少一個答對的機率就是 1 減掉它,這是單題的算法。換到一整組題目,要出現他們觀察到的 power law,必要且充分條件是題目裡有一條很長的難題尾巴,也就是 pass@1 非常低的那些題。實際去看資料,形狀就是這樣:一大票簡單題在 pass@1 就被解掉,題目越難 pass@1 越低,最後拖出一條長尾。

對整個產業來說,這個發現把花錢的位置多開了一個。以前公司把幾億、幾十億美元壓在 pretraining,fine-tuning 少很多,每次 inference 幾乎像不用錢;現在 inference compute 本身可以是一筆大開銷,而且可以離線做:把 agent 放出去解一個問題,讓它一直生 token、一直把答案變好。講者說,test-time scaling 讓算力有第二個可以預測地花掉的地方。

驗證先行:有些領域天生驗得動

重複取樣要真的有用,前提是系統挑得出對的那一份,所以自動驗證是整件事的瓶頸。有些領域的驗證天生比別的領域簡單。

形式數學可以用 formal proof 的軟體跑,逐步檢查證明對不對。程式題可以寫 unit test,而寫一個 unit test 通常比寫整個程式簡單得多,這件事甚至交給人做就好。講者特別興奮的一個方向是 AI as a compiler:用 LLM 把 PyTorch 這種原始碼編譯成 CUDA 這類更貼近硬體、為 GPU 最佳化的低階語言,再比對兩邊在同樣輸入下的輸出是不是一樣。你不需要判斷 CUDA 寫得好不好,只要看行為等不等價。實驗室有一套叫 Kernel Bench 的 CUDA 生成 benchmark,樣本數往上加,coverage 就跟著往上,用的就是這個等價性當 perfect verifier。同樣的道理可以推到任何兩種語言之間的翻譯:Python、C++ 跟 Java 互轉,等價性比正確性容易量。

Generation-verification gap:生得出來,挑不出來

沒有 verifier 的領域就是另一回事。講者比較了幾種 best-of-N 的選法:majority voting,也就是選出現次數最多的那個答案;用模型排序;用 reward model 打分。圖上綠線是 majority voting,藍線是假設有一個 perfect verifier 的 coverage,兩條線之間差了一大截。majority voting 大概在 10 個或 50 個樣本之後就平掉了,coverage 卻還在往上走,而且這個差距在 MATH 這種比較難的資料集上比在 GSM8K 上更明顯。換成 reward model 打分數、取最高分那一份,也還是留下一大段沒用到。

這個落差就是講者口中的 generation-verification gap。生成端生得出很多好答案,驗證端決定你實際能拿回多少。原因出在難題上:在一千個、一萬個樣本裡,最難的那些題可能只對過一次、兩次、三次,majority voting 根本選不到這種罕見的正確答案。就算在比較簡單的 GSM8K 上,也有些難題的正確答案出現頻率低到讓驗證變得很難,你要怎麼找到它們?

verifier 的品質自己也是變數。課堂上有人問,量到的 coverage 到底可不可信。講者說數學題的答案他們人工檢查過一輪,準確率在九成多,她給的數字是 97、98% 上下;但 unit test 沒覆蓋到的地方就會漏,所以 verifier 有多好,直接決定這個 coverage 數字有多少是真的。那組每題一萬個樣本的資料集放在 Hugging Face 上,講者說想把 generation-verification gap 縮小的人,這可以當課程專案。

課堂討論裡學生提了幾個方向。有人提 hybrid 架構,用 knowledge graph 加上 advanced RAG 的文件,平常先靠知識庫回答,覺得不夠準才去查 RAG,因為 RAG 本身也吃延遲,目標是降低對 verifier 的依賴。有人提 self-study:先讓模型去探索一個領域、把學到的東西整理起來,再進入平行的解題取樣;講者說後面的課會講 self-study、search 跟 tool use 這些重複取樣以外的方法。也有人反過來想:如果確認一個答案是對的很貴,那剔除明顯錯的可能便宜得多,用模擬、外部工具或另一顆模型來過濾。至於「生一萬個 verifier 來多數決」的提議,講者說他們做過的版本是 weakly supervised 的 verifier 集成,十顆、二十顆,不是一萬顆,方向有趣但算力上很貴。

平行取樣之外還有序列修改,reward model 決定算力往哪走

接下來講者換到另一篇論文,標題直白:scaling LLM test time compute optimally can be more effective than scaling model parameters。她把 test-time scaling 拆成幾件事:平行取樣是一種,讓模型對同一個輸入獨立生好幾個答案;序列修改是另一種,模型先提出一個初步做法,從不同角度看它、一直改到可以回答為止。論文裡的序列修改是靠 prompting 誘導出來的,今天的 reasoning model 因為被訓練過,內部自己就會出現這種行為。

兩種可以混。最簡單的版本是生兩個平行的候選,各自做序列修改,最後用 PRM 或 ORM 挑比較好的那一個。講者說某種組合通常比只用一種好,但沒有通用的分配比例,這塊還是開放問題。

reward model 分成兩種。outcome reward model,簡稱 ORM,看最後的答案打分,best-of-N 就是挑分數最高的那一份。process reward model,簡稱 PRM,打的是解題過程裡每一步的分數。這裡的 step 不是 token,而是一個在智識上完整的段落,或一句話,看你怎麼定義。訓練資料可以找人在每一步標好或壞,用這些標註把 PRM 訓出來。PRM 在自己訓練領域附近的任務上表現最好,但因為它本質上是 language model,跨任務還是會有一點 generalization。實務上有現成的 PRM 可以拿來用在你的任務上,也可以用自己的資料自己訓一顆 PRM 跟 ORM。

PRM 可以拿來引導搜尋。課堂上示範的是 beam search:每一層給四個樣本的預算,PRM 挑分數最高的兩條分支,從這兩條繼續往下展,也可以用一個門檻決定要留哪些。這樣做會把沒希望的路徑剪掉,算力集中在有機會的分支上。

難度決定算力該怎麼分

講者引的論文用一個數學資料集,大約一萬兩千題訓練、500 題測試,難度用 pass@1 定義,再切成五個 bin。研究團隊自己訓了 PRM,也微調了一顆專門做修改的模型;不過講者補了一句,現在只要是 instruction-tuned 的模型,prompt 叫它修改,它就會改。

結果比的是 majority voting、ORM、PRM 引導,以及各種混合方式,目標是找出在每個準確度水準下,最少需要多少 generation 預算。容易的題目那一側,sequential 對 parallel 的比例越高(圖上顏色最深的那塊)表現越好。題目越難,最適比例就越難講,甚至從第四個 bin 到第五個 bin 就換了答案。課堂上有人把直覺講出來:容易的題目有很多條路都通,所以一條路慢慢改就好;難的題目可能只有少數幾條路走得通,那就得多開幾條平行路徑去撞。講者說這是個直覺上合理的講法。

還有一個好處是逐步驗證。有 PRM 之後,違反限制條件的分支可以在還沒把算力花下去之前就剪掉。

Test-time scaling 補位 pretraining,但補不到最難的那一格

這篇論文還有一個觀察,講者也說這到今天為止還是成立:對容易跟中等的題目來說,多加測試階段的算力,比多花 pretraining 的算力划算。圖上把 inference token 對 pretraining token 的比例畫出來,容易跟中等的題目這個比例是正的。但最難的題目換了答案:即使給一個相當大的 inference 預算,模型更大、pretraining 更多的那一邊還是贏。講者提醒,不要因此下結論說 pretraining 做完了、只要一直從手上的模型取樣就好。

課堂上有人問到帳要怎麼算,因為 pretraining 付一次,inference scaling 是每碰到一個新題目就再付一次。講者同意這個區分,但她補了另一面:不是每個組織都 pretrain 得起一顆大模型,那條路太貴;test-time scaling 對很多題目來說是可行、負擔得起的那一條,即使每次都要重付。開源模型跟小模型會隨著測試階段算力加得越多而越有用,這件事在她們實驗室很多研究裡都看到,可是最難的那些問題,還是 frontier model 做得比較好。

有人問,那能不能把一顆小模型變成難題的專家?講者把兩件事分開:針對一組特定題目微調,當然可以變好;但論文比的是通用的做法,資料是一般性的資料,訓練配方也是大家常用的那一套,比的是一般情況下的表現。

Archon:把 inference 的算力分配當成架構搜尋問題

課堂後半是重頭戲,Archon,一篇把 inference-time scaling 當成架構搜尋的論文,共同作者之一是 Yurii。命題是這樣:不同的 inference-time 技巧跟不同的模型可以混搭,那就把「正確性對成本」的那條前緣當成最佳化目標,去找出怎麼搭最好。

系統的輸入有四樣:要最佳化的目標 benchmark、inference call 的預算、手上有哪些 LLM、以及有哪些 inference-time 技巧可以用。中間的 optimizer 叫 ITAS,inference time architecture search,工作是在預算內把模型跟操作排成一個架構。輸出可能同時包含平行生成、序列處理、排序、批評、驗證,以及 fusion。

操作本身有幾種。generation 就是從模型取樣一個或好幾個回應。fusion 是把原始問題跟 k 個候選答案一起給 LLM,請它把所有候選裡的資訊合成一個最終答案。critic 描述一份回應的強項跟弱點。ranker 用 prompt 讓模型把候選排序。verifier 給一份回應打分或做選擇,順便講它為什麼這樣判斷。還有 unit test generation 跟 unit test evaluation:前者生出檢查用的 test;後者更妙,是把答案跟生成的 test 一起給模型,請模型判斷這份答案過不過,test 不一定要真的被執行。這些操作全是 prompting-based,底下的 LLM 沒有為了這些角色被特別訓練過。

舉例來說,一題「檢查括號有沒有配對」的程式題,模型可以生出「奇數個括號就輸出 No」這種 unit test,也可以生出「遇到右括號時,必須對上最近一個還沒配對的左括號」這種 test,甚至可以把實作這些 test 的程式碼也一起生出來。

Fusion 可以贏過 oracle selection

講者用一組實驗說明這些操作各自的份量。手上有 1 到 10 份候選回答時,比較幾種用法:隨機挑一份、請模型排序後挑第一名、用 oracle 挑出最好的那一份、把全部候選丟給模型 fusion 成一個新答案,以及先篩到前五名再 fusion。

排序贏過隨機,oracle 贏過模型排序,這兩個都不意外。有意思的是 fusion:它生出來的答案,比 oracle 從現有候選裡挑最好的那一個還要好。先篩選再 fusion 的分數又更高。講者的解釋是,fusion 不是從幾個答案裡選一個,它可以把散在好幾個不完美候選裡的資訊合起來。候選換成不同模型產出的(集成從最好那顆模型開始,一顆一顆加進來),同樣的趨勢還在,只是隨機選那條線變得更差,因為原本那顆模型是這一群裡比較好的。

加了層的 inference 架構,答案一層比一層準

Archon 搜的是有層次的架構:可能先讓幾個模型各自生成,接著上 critic,再排序,然後好幾個 fuser,中間再穿插一輪 critic 或 ranker,最後收在一個 final fusion。這種搜尋很貴,因為評估一個架構要跑很多次 inference call,所以他們先做了一輪預處理把搜尋空間縮小:第一層一定是 generation;一層只准放一種技巧;critic 一定排在 ranker 或 fuser 前面;unit test generator 後面一定要接 evaluator。離線實驗也看到 generation、critique、ranker、refusion 這串順序效果特別好,就把它保留在搜尋空間裡。

最佳化的目標是準確度對模型呼叫次數,另外可以加上手上有哪些模型、inference 時間、inference 成本預算這些條件。他們用 Bayesian optimization 搜,比 greedy 或隨機搜尋更省樣本,用更少的架構設定就走到最後的結果。

在程式題上,搜出來的一個架構是先生一大堆樣本、生 unit test,然後評估。架構的層數變多、排得仔細,準確度就往上。講者舉的例子是一個用不同模型的集成,三層 critic 跟 fuser,最後再一個 final fuser;它不只贏過「最好的模型只問一次」,也贏過「最好那顆模型問八次、接一層 fusion」。講者把這比成 deep learning 裡加層。她不主張任意加深都有用,但小心選過的 inference layer,確實讓答案在許多任務上一層比一層準。

14.1% 的 pass@1,跟一個還沒分配完的問題

Archon 最後只吐出一個回應,所以量的是 pass@1。它可以是針對單一任務最佳化的版本,也可以是跨多個任務訓練出來的通用版本;通用版本在最佳化時沒見過的任務上也轉移得過去。結果是只用開源模型,在多個任務上打平或超過當時的封閉源碼 frontier model,幅度不小。講者給的平均數字是:在 instruction following、推理、數學、程式這幾類題目上,Archon 的 pass@1 平均比 GPT-4o 或 Claude 3.5 Sonnet 高 14.1%。

課收在一個還沒解完的問題上。講者的說法是,inference token 該被當成一種可以分配的資源,真正難的是怎麼分:平行探索、序列修改、驗證、剪枝、批評、排序、fusion,每一種都在搶同一筆預算,而哪一種該分到多少,目前沒有通用答案。