機(jī)器之心編輯部

在技術(shù)如火如荼發(fā)展的當(dāng)下，業(yè)界常常在思考一個(gè)問(wèn)題：如何利用 AI 發(fā)現(xiàn)科學(xué)問(wèn)題的新最優(yōu)解？

一個(gè)普遍的解法是「測(cè)試時(shí)搜索」（Test-time search），即提示一個(gè)凍結(jié)的（不更新參數(shù)的）大語(yǔ)言模型（LLM）進(jìn)行多次嘗試，這一點(diǎn)類似人類在做編程作業(yè)時(shí)的「猜」解法，尤其是進(jìn)化搜索方法（如 AlphaEvolve），會(huì)將以往的嘗試存入緩沖區(qū)，并通過(guò)人工設(shè)計(jì)、與領(lǐng)域相關(guān)的啟發(fā)式規(guī)則生成新的提示。

可是，盡管這些提示能夠幫助 LLM 改進(jìn)以往的解法，但 LLM 本身并不會(huì)真正提升，就像一個(gè)學(xué)生始終無(wú)法內(nèi)化作業(yè)背后的新思想一樣。

實(shí)際上，能夠讓 LLM 真正進(jìn)步的最直接方式是學(xué)習(xí)。

盡管「學(xué)習(xí)」和「搜索」都能隨著算力擴(kuò)展而良好地增長(zhǎng)，但在 AI 的發(fā)展歷史中，對(duì)于圍棋、蛋白質(zhì)折疊等這類困難問(wèn)題，「學(xué)習(xí)」往往最終超越了「搜索」。因?yàn)椋?strong>科學(xué)發(fā)現(xiàn)本質(zhì)是：超出訓(xùn)練數(shù)據(jù)與人類現(xiàn)有知識(shí)的 out-of-distribution 問(wèn)題。

為此，斯坦福大學(xué)、英偉達(dá)等機(jī)構(gòu)聯(lián)合提出一種新方法：在測(cè)試時(shí)進(jìn)行強(qiáng)化學(xué)習(xí)（RL），即讓 LLM 在嘗試解決特定測(cè)試問(wèn)題的過(guò)程中持續(xù)訓(xùn)練自己。

• 論文鏈接：https://www.alphaxiv.org/abs/2601.16175
• 項(xiàng)目地址：https://github.com/test-time-training/discover

具體來(lái)看，團(tuán)隊(duì)只是把單個(gè)測(cè)試問(wèn)題定義為一個(gè)環(huán)境，并在其中執(zhí)行強(qiáng)化學(xué)習(xí)（RL），因此任何標(biāo)準(zhǔn) RL 技術(shù)原則上都可以應(yīng)用。然而，需要注意的是，這里的目標(biāo)與標(biāo)準(zhǔn) RL 存在關(guān)鍵差異，這里的目標(biāo)不是讓模型在各類問(wèn)題上平均表現(xiàn)更好，而是只為了解決眼前這一個(gè)問(wèn)題，并且只需要產(chǎn)出一個(gè)優(yōu)秀的解決方案，而不是平均產(chǎn)生多個(gè)良好的解決方案。

團(tuán)隊(duì)將該方法命名為「Test-Time Training to Discover」（TTT-Discover）。為了適應(yīng)上述目標(biāo)，其學(xué)習(xí)目標(biāo)函數(shù)和搜索子程序都旨在優(yōu)先考慮最有希望的解決方案 。

結(jié)果顯示，該方法在多種任務(wù)上取得了好成績(jī)，包括擊敗了 DeepMind 的 AlphaEvolve；數(shù)學(xué)領(lǐng)域 在 Erd?s 最小重疊問(wèn)題上取得了新突破；在 GPUMode 競(jìng)賽中，開發(fā)出了比人類最佳內(nèi)核快兩倍的全新 A100 GPU 內(nèi)核；在 AtCoder 測(cè)試中超越了最佳 AI 代碼和人類代碼；在單細(xì)胞分析的去噪任務(wù)中取得最好成績(jī)……

值得注意的是，該方法在使用開放模型 OpenAI gpt-oss-120b 基礎(chǔ)上，計(jì)算成本非常低，通過(guò)使用 Thinking Machines 的API Tinker ，每個(gè)問(wèn)題只需花費(fèi)幾百美元。

在業(yè)界看來(lái)，TTT-Discover 所提出的理念，或?yàn)槌掷m(xù)學(xué)習(xí)打開了新的想象空間。

TTT-Discover 方法創(chuàng)新

下圖展示了 TTT-Discover 的核心機(jī)制，展示 TTT-Discover 在測(cè)試階段針對(duì)單個(gè)問(wèn)題持續(xù)對(duì)大語(yǔ)言模型（LLM）進(jìn)行訓(xùn)練，記 πθi 為在測(cè)試時(shí)訓(xùn)練第 i 步更新權(quán)重后的策略。該圖繪制的是 TTT-Discover 在 GPUMode TriMul 競(jìng)賽中測(cè)試時(shí)，第 0 步、第 9 步、第 24 步以及第 49 步（最終階段）的獎(jiǎng)勵(lì)分布情況，每一步都會(huì)生成 512 個(gè)候選解。

可以看到，隨著訓(xùn)練過(guò)程的推進(jìn)，LLM 逐漸生成更優(yōu)的解，并最終超越了以往的最優(yōu)結(jié)果（即人類最佳方案）。

需要注意的是，TTT-Discover 沒(méi)有直接套用標(biāo)準(zhǔn)的 RL 算法（如 PPO/GRPO）。

因?yàn)閳F(tuán)隊(duì)認(rèn)為，標(biāo)準(zhǔn) RL 優(yōu)化的是期望獎(jiǎng)勵(lì)（平均分），而科學(xué)探索只在乎最大獎(jiǎng)勵(lì)（最高分），只要能找到一個(gè)突破性的解，策略在其他時(shí)候表現(xiàn)差也沒(méi)關(guān)系；這樣的策略容易讓發(fā)現(xiàn)探索僅僅止步于「安全但平庸」的高分區(qū)域，而不敢去嘗試可能帶來(lái)突破的高風(fēng)險(xiǎn)區(qū)域。另外，傳統(tǒng)算法每次都是從頭開始，無(wú)法逐步演化復(fù)雜解。

為此，團(tuán)隊(duì)引入兩個(gè)關(guān)鍵組件來(lái)解決上述問(wèn)題。

一是熵目標(biāo)函數(shù)，作用是通過(guò)指數(shù)加權(quán)來(lái)極端地偏向高獎(jiǎng)勵(lì)樣本。隨著 β → ∞，熵目標(biāo)函數(shù)趨近于最大值（max）。然而，團(tuán)隊(duì)發(fā)現(xiàn)，在訓(xùn)練早期若 β 過(guò)大，會(huì)導(dǎo)致訓(xùn)練不穩(wěn)定；在訓(xùn)練后期若 β 過(guò)小，則隨著改進(jìn)幅度越來(lái)越微小，優(yōu)勢(shì)函數(shù)會(huì)逐漸消失，這說(shuō)明為不同任務(wù)設(shè)定一個(gè)統(tǒng)一且固定的 β 常數(shù)是非常困難的。

為此，團(tuán)隊(duì)為每一個(gè)初始狀態(tài)自適應(yīng)地設(shè)置 β(s)，通過(guò)約束由該目標(biāo)函數(shù)誘導(dǎo)的策略的 KL 散度來(lái)實(shí)現(xiàn)。

二是受 PUCT 啟發(fā)的狀態(tài)復(fù)用策略，采用該規(guī)則來(lái)選擇初始狀態(tài)。每個(gè)狀態(tài) s 的評(píng)分為：

其中，Q (s) 表示當(dāng)初始狀態(tài)為 s 時(shí)所生成狀態(tài)中的最大回報(bào)（如果 s 尚未被選擇過(guò)，則取 R (s)）。不同于以往研究中采用「平均回報(bào)」的做法，團(tuán)隊(duì)在 Q (s) 中使用的是子狀態(tài)的最大回報(bào)，這也是關(guān)注的核心是從某個(gè)狀態(tài)出發(fā)所能達(dá)到的最佳結(jié)果，而不是平均結(jié)果。這種設(shè)計(jì)確保搜索集中在最有前景的解決路徑上，同時(shí)保持多樣性。

整體來(lái)看，熵目標(biāo)和 PUCT 復(fù)用策略的結(jié)合使 TTT-Discover 能夠優(yōu)先發(fā)現(xiàn)單一的最高獎(jiǎng)勵(lì)解決方案，而不是多個(gè)解決方案的平均表現(xiàn)。

結(jié)果評(píng)估

團(tuán)隊(duì)在四個(gè)截然不同的領(lǐng)域 —— 數(shù)學(xué)、GPU 內(nèi)核工程、算法設(shè)計(jì)和生物學(xué)問(wèn)題上評(píng)估了 TTT-Discover。

除了考慮潛在的影響力外，選擇領(lǐng)域的標(biāo)準(zhǔn)還考慮到兩個(gè)方面，首先，選擇能夠?qū)⒆陨肀憩F(xiàn)與人類專家進(jìn)行比較的領(lǐng)域，例如，可以通過(guò)與人類工程競(jìng)賽中的最佳提交方案或?qū)W術(shù)論文中報(bào)告的最佳結(jié)果進(jìn)行對(duì)比來(lái)實(shí)現(xiàn)，比如數(shù)學(xué)和算法設(shè)計(jì)，可以說(shuō)是近期相關(guān)工作取得非常大進(jìn)展的領(lǐng)域之一。

在每個(gè)應(yīng)用中，團(tuán)隊(duì)都報(bào)告了已知的人類最佳結(jié)果和 AI 最佳結(jié)果。

可以看到，在數(shù)學(xué)領(lǐng)域，關(guān)于構(gòu)造數(shù)學(xué)對(duì)象（如階躍函數(shù)）來(lái)證明不等式的更緊致邊界 ——Erd?s 最小重疊問(wèn)題任務(wù)上，之前人類最佳表現(xiàn)是 0.380927、AI 最佳表現(xiàn) (AlphaEvolve) 是 0.380924，而 TTT-Discover 刷新記錄，拿到了的成績(jī)。

在 GPU 內(nèi)核優(yōu)化任務(wù)中，首先需要說(shuō)明的是「新的最優(yōu)解」（state of the art）意味著實(shí)現(xiàn)了比現(xiàn)有方案更快的內(nèi)核實(shí)現(xiàn)。團(tuán)隊(duì)選擇 GPUMODE 作為評(píng)測(cè)平臺(tái)，因?yàn)槠渑判邪窠?jīng)過(guò)大量人類競(jìng)賽的充分驗(yàn)證，并配備了穩(wěn)健的評(píng)測(cè)框架，同時(shí)，其基準(zhǔn)測(cè)試避免了信噪比問(wèn)題，即避免因操作過(guò)于簡(jiǎn)單或輸入規(guī)模過(guò)小而使系統(tǒng)開銷主導(dǎo)運(yùn)行時(shí)間的情況。

結(jié)果是，團(tuán)隊(duì)的 TriMul 內(nèi)核在所有 GPU 類型上均達(dá)到了當(dāng)前最優(yōu)水平。在 A100 上，TTT-Discover 找到的最佳內(nèi)核比人類專家提交的最優(yōu)方案快 50%，盡管在訓(xùn)練階段團(tuán)隊(duì)的獎(jiǎng)勵(lì)函數(shù)并未在 A100 上直接計(jì)時(shí)。總體而言，在所有 GPU 類型上，該方法都相對(duì)于人類最佳結(jié)果實(shí)現(xiàn)了超過(guò) 15% 的性能提升。

而在另外兩項(xiàng)測(cè)試中，TTT-Discover 同樣取得了非凡的成績(jī)。

雖然當(dāng)前 TTT-Discover 方法取得了非常好的成績(jī)，但是團(tuán)隊(duì)也承認(rèn)，該方法目前的形式只能應(yīng)用于具有連續(xù)獎(jiǎng)勵(lì)的問(wèn)題中，而未來(lái)工作最重要的方向是針對(duì)具有稀疏獎(jiǎng)勵(lì)或二元獎(jiǎng)勵(lì)的問(wèn)題，比如數(shù)學(xué)證明、科學(xué)假說(shuō)，或者不可驗(yàn)證領(lǐng)域的問(wèn)題（物理、生物推理等）進(jìn)行測(cè)試時(shí)訓(xùn)練。