TTT-Discover:斯坦福&英伟达提出测试时强化学习,多项任务超越人类专家

斯坦福、英伟达提出TTT-Discover,一种测试时强化学习方法,部分任务超越人类专家2倍,或为持续学习带来新思路。

原文标题:比人类专家快2倍,斯坦福联合英伟达发布TTT-Discover:用「测试时强化学习」攻克科学难题

原文作者:机器之心

冷月清谈:

斯坦福大学、英伟达等机构联合提出了 TTT-Discover,这是一种在测试时进行强化学习(RL)的新方法,旨在解决科学问题的新最优解。该方法将单个测试问题定义为一个环境,并让 LLM 在解决问题的过程中持续训练自己。与标准 RL 不同,TTT-Discover 的目标是只解决当前问题,并产出一个优秀的解决方案。通过引入熵目标函数和状态复用策略,TTT-Discover 能够优先考虑最有希望的解决方案。实验结果表明,TTT-Discover 在多个领域都取得了显著成果,包括在 Erdős 最小重叠问题上取得新突破,开发出比人类最佳内核快两倍的全新 A100 GPU 内核等。该方法或为持续学习打开了新的想象空间,但目前仅适用于具有连续奖励的问题。

怜星夜思:

1、TTT-Discover的“测试时训练”思路,对比传统AI训练方式,有哪些优势和局限?这种思路未来可能在哪些领域有更广泛的应用?
2、TTT-Discover在GPU内核优化任务中,奖励函数似乎没有直接在A100上进行计时,但最终结果却优于人类专家。这说明了什么?
3、文章提到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 刷新记录,拿到了 0.380876 的成绩。



在 GPU 内核优化任务中,首先需要说明的是「新的最优解」(state of the art)意味着实现了比现有方案更快的内核实现。团队选择 GPUMODE 作为评测平台,因为其排行榜经过大量人类竞赛的充分验证,并配备了稳健的评测框架,同时,其基准测试避免了信噪比问题,即避免因操作过于简单或输入规模过小而使系统开销主导运行时间的情况。


结果是,团队的 TriMul 内核在所有 GPU 类型上均达到了当前最优水平。在 A100 上,TTT-Discover 找到的最佳内核比人类专家提交的最优方案快 50%,尽管在训练阶段团队的奖励函数并未在 A100 上直接计时。总体而言,在所有 GPU 类型上,该方法都相对于人类最佳结果实现了超过 15% 的性能提升。



而在另外两项测试中,TTT-Discover 同样取得了非凡的成绩。




虽然当前 TTT-Discover 方法取得了非常好的成绩,但是团队也承认,该方法目前的形式只能应用于具有连续奖励的问题中,而未来工作最重要的方向是针对具有稀疏奖励或二元奖励的问题,比如数学证明、科学假说,或者不可验证领域的问题(物理、生物推理等)进行测试时训练。


更多详细信息,可阅读原文获取!


© THE END 

转载请联系本公众号获得授权

投稿或寻求报道:liyazhou@jiqizhixin.com

我认为TTT-Discover的核心在于其目标导向性,它不像传统RL那样追求平均性能的提升,而是All in on 解决当前问题。这在科学研究中尤为重要,因为科学发现往往是偶然的、突破性的。缺点也很明显,就是泛化能力可能较弱,换个问题可能就得重新学习。至于应用前景,我认为在需要定制化解决方案的场景,比如药物发现、材料设计等领域,会有很大的潜力。

稀疏奖励问题确实是强化学习的一大挑战。我觉得可以借鉴一些解决稀疏奖励问题的经典方法,比如:
1. 奖励塑造(Reward Shaping):设计更密集的中间奖励,引导模型逐步接近目标。
2. 好奇心驱动(Curiosity-driven Exploration):鼓励模型探索未知的状态,提高探索的效率。
3. 分层强化学习(Hierarchical Reinforcement Learning):将问题分解为多个子问题,逐层解决。
针对数学证明这种二元奖励问题,可以考虑引入博弈的思想,让模型自己与自己对弈,通过对抗学习来提高证明能力。

我猜测,解决稀疏奖励问题的一个关键在于如何有效地进行探索。在数学证明或科学假说验证中,每一步的尝试可能都是错的,只有最终证明成功才有奖励。 这时候,传统的随机探索可能效率非常低。 或许可以考虑将人类的先验知识融入到探索策略中,比如利用已有的数学定理或科学理论来指导模型的探索方向。 另外,模仿学习也是一个可能的方向,可以学习人类专家是如何进行证明或提出假说的。

我理解这体现了一种迁移学习的思想。虽然没有直接在A100上训练,但模型在其他GPU上的学习经验可以泛化到A100上。这提示我们在设计AI系统时,可以考虑利用不同环境的数据进行训练,以提高模型的泛化能力。当然,这可能也说明GPUMODE这个评测平台本身存在一些可以被“绕过”的漏洞,毕竟AI的目标是优化指标,而不是真正理解问题本身。

测试时训练的最大优势在于它的适应性!传统AI训练好后就基本定型了,遇到没见过的新问题可能就抓瞎。但TTT-Discover就像一个能边考试边学习的学霸,针对特定问题优化,效果当然更好啦。不过,它也依赖于奖励函数的设置,如果奖励函数设计不好,那可能就南辕北辙了。我觉得除了文章里提到的领域,像金融风控、智能制造这种需要快速适应变化的环境,TTT-Discover说不定也能大显身手。

这其实给我一种“无心插柳柳成荫”的感觉。有可能奖励函数的设计,出发点并不是为了A100,而是为了解决更通用的问题。结果AI在学习过程中,意外地发现了A100上的更优解。 这也提醒我们,在科研探索中,有时候坚持初心,反而能有意想不到的收获。

其实我有一个脑洞大开的想法:能不能把人类的反馈也引入到训练过程中? 比如,让数学家或科学家来评价模型生成的中间步骤,并给予一定的奖励或惩罚。 这样就可以利用人类的智慧来弥补稀疏奖励的不足。 但是,这种方法也面临着一些挑战,比如如何保证人类反馈的质量和一致性,以及如何处理人类反馈中的偏见等等。

这说明奖励函数的设计非常关键!即使没有直接针对A100进行优化,但设计的奖励函数能够捕捉到内核优化的本质,引导模型找到更高效的解决方案。 就像高考,知识点掌握牢固了,即使没做过完全一样的题,也能举一反三。也可能说明咱们人类在内核优化方面,还有很大的提升空间啊(手动狗头)。

谢邀,利益相关,我是一个科研工作者。我觉得这个思路很棒!传统的AI训练模式,就像是填鸭式教育,学了很多知识,但是不一定能解决实际问题。TTT-Discover这种测试时训练,更像是项目制学习,在解决问题的过程中学习,更有针对性。但是,这种方法也需要大量的计算资源和时间,成本较高。另外,如何设计合适的奖励函数也是一个挑战。未来,我希望可以在我的研究领域尝试应用这种方法。