斯坦福聯合英偉達TTT-Discover:用測試時強化學習攻克科學難題

機器之心編輯部

在技術如火如荼發展的當下，業界常常在思考一個問題：如何利用 AI 發現科學問題的新最優解？

一個普遍的解法是「測試時搜索」（Test-time search），即提示一個凍結的（不更新參數的）大語言模型（LLM）進行多次嘗試，這一點類似人類在做編程作業時的「猜」解法，尤其是進化搜索方法（如 AlphaEvolve），會將以往的嘗試存入緩沖區，并通過人工設計、與領域相關的啟發式規則生成新的提示。

可是，盡管這些提示能夠幫助 LLM 改進以往的解法，但 LLM 本身并不會真正提升，就像一個學生始終無法內化作業背后的新思想一樣。

實際上，能夠讓 LLM 真正進步的最直接方式是學習。

盡管「學習」和「搜索」都能隨著算力擴展而良好地增長，但在 AI 的發展歷史中，對于圍棋、蛋白質折疊等這類困難問題，「學習」往往最終超越了「搜索」。因為，科學發現本質是：超出訓練數據與人類現有知識的 out-of-distribution 問題。

為此，斯坦福大學、英偉達等機構聯合提出一種新方法：在測試時進行強化學習（RL），即讓 LLM 在嘗試解決特定測試問題的過程中持續訓練自己。

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

具體來看，團隊只是把單個測試問題定義為一個環境，并在其中執行強化學習（RL），因此任何標準 RL 技術原則上都可以應用。然而，需要注意的是，這里的目標與標準 RL 存在關鍵差異，這里的目標不是讓模型在各類問題上平均表現更好，而是只為了解決眼前這一個問題，并且只需要產出一個優秀的解決方案，而不是平均產生多個良好的解決方案。

團隊將該方法命名為「Test-Time Training to Discover」（TTT-Discover）。為了適應上述目標，其學習目標函數和搜索子程序都旨在優先考慮最有希望的解決方案 。

結果顯示，該方法在多種任務上取得了好成績，包括擊敗了 DeepMind 的 AlphaEvolve；數學領域 在 Erd?s 最小重疊問題上取得了新突破；在 GPUMode 競賽中，開發出了比人類最佳內核快兩倍的全新 A100 GPU 內核；在 AtCoder 測試中超越了最佳 AI 代碼和人類代碼；在單細胞分析的去噪任務中取得最好成績……

值得注意的是，該方法在使用開放模型 OpenAI gpt-oss-120b 基礎上，計算成本非常低，通過使用 Thinking Machines 的API Tinker ，每個問題只需花費幾百美元。

在業界看來，TTT-Discover 所提出的理念，或為持續學習打開了新的想象空間。

TTT-Discover 方法創新

下圖展示了 TTT-Discover 的核心機制，展示 TTT-Discover 在測試階段針對單個問題持續對大語言模型（LLM）進行訓練，記 πθi 為在測試時訓練第 i 步更新權重后的策略。該圖繪制的是 TTT-Discover 在 GPUMode TriMul 競賽中測試時，第 0 步、第 9 步、第 24 步以及第 49 步（最終階段）的獎勵分布情況，每一步都會生成 512 個候選解。

可以看到，隨著訓練過程的推進，LLM 逐漸生成更優的解，并最終超越了以往的最優結果（即人類最佳方案）。

需要注意的是，TTT-Discover 沒有直接套用標準的 RL 算法（如 PPO/GRPO）。

因為團隊認為，標準 RL 優化的是期望獎勵（平均分），而科學探索只在乎最大獎勵（最高分），只要能找到一個突破性的解，策略在其他時候表現差也沒關系；這樣的策略容易讓發現探索僅僅止步于「安全但平庸」的高分區域，而不敢去嘗試可能帶來突破的高風險區域。另外，傳統算法每次都是從頭開始，無法逐步演化復雜解。

為此，團隊引入兩個關鍵組件來解決上述問題。

一是熵目標函數，作用是通過指數加權來極端地偏向高獎勵樣本。隨著 β → ∞，熵目標函數趨近于最大值（max）。然而，團隊發現，在訓練早期若 β 過大，會導致訓練不穩定；在訓練后期若 β 過小，則隨著改進幅度越來越微小，優勢函數會逐漸消失，這說明為不同任務設定一個統一且固定的 β 常數是非常困難的。

為此，團隊為每一個初始狀態自適應地設置 β(s)，通過約束由該目標函數誘導的策略的 KL 散度來實現。

二是受 PUCT 啟發的狀態復用策略，采用該規則來選擇初始狀態。每個狀態 s 的評分為：

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

整體來看，熵目標和 PUCT 復用策略的結合使 TTT-Discover 能夠優先發現單一的最高獎勵解決方案，而不是多個解決方案的平均表現。

結果評估

團隊在四個截然不同的領域 —— 數學、GPU 內核工程、算法設計和生物學問題上評估了 TTT-Discover。

除了考慮潛在的影響力外，選擇領域的標準還考慮到兩個方面，首先，選擇能夠將自身表現與人類專家進行比較的領域，例如，可以通過與人類工程競賽中的最佳提交方案或學術論文中報告的最佳結果進行對比來實現，比如數學和算法設計，可以說是近期相關工作取得非常大進展的領域之一。

在每個應用中，團隊都報告了已知的人類最佳結果和 AI 最佳結果。

可以看到，在數學領域，關于構造數學對象（如階躍函數）來證明不等式的更緊致邊界 ——Erd?s 最小重疊問題任務上，之前人類最佳表現是 0.380927、AI 最佳表現 (AlphaEvolve) 是 0.380924，而 TTT-Discover 刷新記錄，拿到了的成績。

在 GPU 內核優化任務中，首先需要說明的是「新的最優解」（state of the art）意味著實現了比現有方案更快的內核實現。團隊選擇 GPUMODE 作為評測平臺，因為其排行榜經過大量人類競賽的充分驗證，并配備了穩健的評測框架，同時，其基準測試避免了信噪比問題，即避免因操作過于簡單或輸入規模過小而使系統開銷主導運行時間的情況。

結果是，團隊的 TriMul 內核在所有 GPU 類型上均達到了當前最優水平。在 A100 上，TTT-Discover 找到的最佳內核比人類專家提交的最優方案快 50%，盡管在訓練階段團隊的獎勵函數并未在 A100 上直接計時。總體而言，在所有 GPU 類型上，該方法都相對于人類最佳結果實現了超過 15% 的性能提升。

而在另外兩項測試中，TTT-Discover 同樣取得了非凡的成績。

雖然當前 TTT-Discover 方法取得了非常好的成績，但是團隊也承認，該方法目前的形式只能應用于具有連續獎勵的問題中，而未來工作最重要的方向是針對具有稀疏獎勵或二元獎勵的問題，比如數學證明、科學假說，或者不可驗證領域的問題（物理、生物推理等）進行測試時訓練。