← 返回 2026-01-23

测试时训练发现(TTT-Discover):在推理时持续训练LLM以发现科学问题的新SOTA Learning to Discover at Test Time

Mert Yuksekgonul, Daniel Koceja, Xinhao Li, Federico Bianchi, Jed McCaleb, Xiaolong Wang, Jan Kautz, Yejin Choi, James Zou, Carlos Guestrin, Yu Sun 📅 2026-01-22 👍 45 2026-07-13 08:35
GPU核优化 大语言模型 强化学习 数学优化 测试时训练 科学发现

测试时强化学习训练LLM,在单个科学问题上发现新SOTA

前置知识

Test-Time Training (TTT, 测试时训练)

测试时训练是一种在推理阶段对模型进行适应性训练的技术范式。其核心思想是:与其让模型在部署前对所有可能的输入做好准备,不如在遇到具体测试实例时,利用该实例生成相关数据并进行针对性训练。这一概念最早可追溯到1970年代的局部加权回归(locally weighted regression),近年来在NLP的动态评估、计算机视觉的任务适应等场景中广泛应用。在本文中,TTT被扩展为对单个科学问题进行强化学习训练,使LLM能够从自身的尝试中学习并持续改进。

本文的核心方法TTT-Discover正是建立在TTT的范式之上,理解TTT的概念是理解本文创新点的前提。

Reinforcement Learning (RL, 强化学习)

强化学习是机器学习的一个分支,智能体通过与环境交互来学习最优策略。标准RL的目标是找到一个策略 π_θ,使其在所有可能状态下的期望累积奖励最大化。常见的算法包括PPO(近端策略优化)和GRPO(广义近端策略优化)。本文的关键洞察是:发现类问题的RL目标与标准RL不同——我们只需要找到一个最优解,而非最大化平均表现。

TTT-Discover在单个测试问题上执行RL,但其目标函数和搜索策略与标准RL有本质区别,理解标准RL才能理解本文的创新。

LoRA (Low-Rank Adaptation)

LoRA是一种参数高效微调方法,通过在预训练权重矩阵旁添加低秩分解矩阵来实现微调。对于权重矩阵 W,LoRA引入可训练的低秩矩阵 A 和 B,使得前向传播变为 h = (W + BA)x。本文使用秩为32的LoRA对gpt-oss-120b进行测试时训练,使得在数百美元的预算内完成50步训练成为可能。

LoRA使得在推理时对大型模型进行训练在计算和经济上可行,是本文方法能以低成本运行的关键技术支撑。

PUCT (Predictor + UCB applied to Trees)

PUCT是一种蒙特卡洛树搜索(MCTS)中的节点选择启发式,最初在AlphaGo中使用。它通过结合节点的价值估计 Q(s)、先验概率 P(s) 和探索奖励来平衡利用与探索。其评分为 Q(s) + c * P(s) * sqrt(1+T) / (1+n(s)),其中 T 是总展开次数,n(s) 是该节点被选中的次数,c 是探索系数。本文创新性地将PUCT用于选择初始状态(而非树搜索),并使用最大奖励而非平均奖励作为 Q(s)。

PUCT是本文复用策略(reuse routine)的核心组件,决定了如何从缓冲区中选择最有潜力的初始状态进行下一步探索。

Entropy-Regularized Objective (熵正则化目标)

传统RL优化期望奖励,而熵正则化目标(或称entropic objective)通过引入指数权重来偏向高奖励的动作。其目标函数为 J_β(θ) = E_{s~reuse(H)} log E_{a~π_θ(·|s)} [e^{β(s)R(s,a)}]。当 β 趋向无穷大时,该目标趋向于最大化奖励(而非平均奖励),这更符合发现类问题的本质——我们只需要一个最优解。本文还通过KL散度约束自适应地设置 β(s),避免训练不稳定。

这是本文区别于标准RL的核心技术贡献,直接解决了发现类问题中'只关心最佳结果而非平均表现'的目标差异。

Evolutionary Search (进化搜索)

进化搜索方法(如AlphaEvolve)通过维护一个解决方案缓冲区,并使用变异(mutation)和交叉(crossover)等启发式操作生成新方案。这些方法通常需要精心设计的领域特定启发式来定义如何修改和组合现有方案。进化搜索的核心限制是:LLM本身保持冻结状态,只能通过改变输入提示(prompt)来间接改进输出,而无法从经验中真正学习和内化新知识。

本文的TTT-Discover正是对进化搜索方法的替代——通过在测试时训练LLM权重而非仅改变提示,实现了更深层次的学习。

研究动机

当前AI在科学发现问题上主要依赖测试时搜索(test-time search),即通过提示一个冻结的LLM来进行大量尝试。具体来说,AlphaEvolve等进化搜索方法将过去的尝试存储在缓冲区中,通过手工设计的、领域特定的启发式来生成新的提示。虽然这些提示可以帮助LLM改进之前的解决方案,但LLM本身无法改进——就像一个学生永远无法内化作业背后的新思想。在标准RL中,策略 π_θ 的权重 θ 在测试时保持不变,所有改进都来自于改变输入 (d, s_i, c_i)。这导致了三个根本性问题:(1)朴素RL优化平均表现,对SOTA无感——例如在核工程问题中,从2000µs优化到1900µs需要突破性进展,但两者获得的奖励几乎相同;(2)有效视野短——从零开始尝试限制了策略能达到的深度;(3)探索不足——优化期望奖励会导致策略坍缩到安全的高奖励动作,而不敢尝试可能带来突破的冒险动作。

本文的目标是本文旨在提出一种新的测试时训练范式TTT-Discover,使LLM能够在测试时通过强化学习持续训练自身,从而在单个科学问题上发现超越人类专家和现有AI系统的新SOTA解决方案。具体目标包括:在数学(Erdős最小重叠问题、自相关不等式)、GPU核工程(TriMul竞赛)、算法设计(AtCoder竞赛)和生物学(单细胞分析去噪)四个领域均达到或超越SOTA,并且仅使用开源模型gpt-oss-120b而非闭源前沿模型,且每问题成本控制在几百美元以内。

与已有工作不同的是,本文的独特切入角度在于识别出发现类问题与标准RL问题的两个关键区别:(1)策略只需要解决单个问题而非泛化到其他问题;(2)我们只需要一个最优解,策略只是达到这一目标的手段。基于这两个区别,本文从两个维度进行创新:在学习目标上,用熵正则化目标 J_β(θ) 替代期望奖励目标,使训练偏向最大奖励动作;在复用策略上,用PUCT启发式替代进化搜索的复杂启发式,用最大奖励(而非平均奖励)评估状态价值。与同时期的工作(ThetaEvolve、MiGrATe、EvoTune)相比,TTT-Discover在相同模型和计算预算下仍能产生显著改进,这归功于其专门针对发现目标设计的学习目标和搜索子程序。

核心方法

TTT-Discover的高层思路可以用一个类比来理解:当你面对一道超难编程作业时,你不是反复猜答案(搜索),而是从每次失败的尝试中学习新思想,最终理解问题的本质。具体来说,TTT-Discover将每个科学问题定义为一个马尔可夫决策过程(MDP),其中状态 s 是候选解,动作 a 是包含思考token和代码的LLM输出,奖励 R(s) 由问题本身定义(如核运行时间的倒数、数学界限的改善程度等)。然后,TTT-Discover在这个由单个测试问题定义的环境中执行强化学习,使LLM的权重 θ 在50个训练步中持续更新。每一步生成512个rollout,使用LoRA(秩32)进行参数高效更新,整个训练运行约500美元。

TTT-Discover的两个核心创新点直接针对发现类问题的本质区别。第一,熵正则化目标 J_β(θ) = E_{s~reuse(H)} log E_{a~π_θ(·|s)} [e^{β(s)R(s,a)}]。与标准RL优化期望奖励不同,当 β 趋向无穷大时该目标趋向最大化,即'只关心最好的那个解'。更重要的是,β(s) 通过KL散度约束自适应设置,避免了固定 β 带来的训练不稳定——太大的 β 在训练早期导致不稳定,太小的 β 在后期使梯度消失。第二,PUCT复用策略。不同于进化搜索使用复杂的变异和交叉启发式,PUCT用简单的评分规则选择初始状态:Q(s) + c * P(s) * sqrt(1+T)/(1+n(s)),其中 Q(s) 使用最大奖励(而非平均奖励),因为我们关心的是从某状态出发能达到的最佳结果。

方法步骤详情

TTT-Discover的完整算法流程如下:输入为问题描述 d 和初始策略 π_{θ_0}。首先,问题 d 诱导出环境的奖励函数 R 和转移函数 T,并用空解初始化缓冲区 H_0。然后进入主循环,共 N=50 步:(步骤5)使用PUCT启发式从缓冲区 H_i 中采样初始状态 s_i 和上下文 c_i;(步骤6)从当前策略采样动作 a_i,每步生成512个rollout,分为8组每组64个,同组使用相同上下文;(步骤7)通过转移函数获得新状态 s'_i;(步骤8)评估奖励 r_i;(步骤9)将当前尝试加入缓冲区;(步骤10)使用熵正则化目标更新权重,应用重要性采样比校正。最终返回所有步中奖励最高的状态。实现细节包括:使用LoRA(秩32)、上下文窗口限制32768 token、推理努力设置为high、通过token强制机制确保模型在上下文耗尽前生成最终响应。

技术新颖性

TTT-Discover的技术新颖性体现在多个层面。首先,从范式层面,这是首次将强化学习(而非仅搜索)应用于测试时发现,使LLM能够从单个问题的尝试中真正内化新知识,而非仅通过改变提示来间接改进。其次,在目标函数设计上,熵正则化目标与自适应 β 的结合是独创的:自适应 β 通过约束诱导策略的KL散度来设置,使得 β(s) 根据每个初始状态的特性动态调整。第三,在复用策略上,将PUCT从蒙特卡洛树搜索的节点选择改造为缓冲区中的状态选择,并创新性地使用最大奖励替代平均奖励作为 Q(s)——这一设计直觉是:快速的核更可能产生更快的核。第四,从工程实现上,使用LoRA在有限上下文窗口内实现了高效的测试时训练。

TTT-Discover在单个问题上持续训练LLM的奖励分布演化
Figure 1: TTT-Discover在单个问题上持续训练LLM的奖励分布演化

实验结果

TTT-Discover在几乎所有尝试的问题上都设置了新SOTA,且仅使用开源模型gpt-oss-120b。在Erdős最小重叠问题上,TTT-Discover将上界改进至0.380876,超越AlphaEvolve的0.380924,改善幅度是AlphaEvolve相对于前SOTA改善幅度的16倍,发现了一个600片的非对称阶梯函数。在第一自相关不等式上,证明了 C_1 ≤ 1.50286(30000片阶梯函数),超越ThetaEvolve的1.50314,且从零开始发现而非在现有构造上精炼。在TriMul核竞赛中,TTT-Discover的核在所有GPU类型上均超越人类最佳:A100上2198µs vs 人类4531µs(快51.5%),H100上1161µs vs 人类1371µs(快15.3%),B200上914µs vs 人类1039µs。在AtCoder竞赛ahc039和ahc058上,若在竞赛期间提交将获得第一名。在单细胞去噪基准上,PBMC得分0.71(MSE 0.15),Tabula得分0.73(MSE 0.14),均超越MAGIC的0.64。消融实验表明完整TTT-Discover的核运行时间为1203µs,远优于无训练(2061µs)、无复用(5274µs)、朴素RL(5329µs)和Best-of-N(5352µs)。特别值得注意的是,使用Qwen3-8B的TTT-Discover在数学不等式问题上也超越了使用相同模型系列的ThetaEvolve,验证了方法的有效性不依赖于特定模型。

科学与工程问题的环境定义概览
Table 1: 科学与工程问题的环境定义概览
数学问题的定量结果对比
Table 2: 数学问题的定量结果对比
圆堆叠问题的结果
Table 3: 圆堆叠问题的结果
TriMul核竞赛运行时间对比
Table 4: TriMul核竞赛运行时间对比
MLA-Decode核竞赛运行时间对比
Table 5: MLA-Decode核竞赛运行时间对比
AtCoder启发式竞赛结果
Table 6: AtCoder启发式竞赛结果
单细胞去噪任务结果
Table 7: 单细胞去噪任务结果
消融实验结果
Table 8: 消融实验结果
Erdős最小重叠问题的阶梯函数对比
Figure 2: Erdős最小重叠问题的阶梯函数对比
第一自相关不等式的阶梯函数与自卷积对比
Figure 3: 第一自相关不等式的阶梯函数与自卷积对比
消融实验的奖励分布演化
Figure 4: 消融实验的奖励分布演化
查看结构化数据
任务指标本文基线提升
Erdős最小重叠问题(组合数学) 上界值(越小越好) 0.380876 0.380924 (AlphaEvolve) 改善0.000048,是AlphaEvolve改善幅度的16倍
第一自相关不等式(分析学) 上界 C_1(越小越好) 1.50286 1.50314 (ThetaEvolve) 改善0.00028,从零开始发现而非精炼现有构造
TriMul核竞赛 A100 运行时间µs(越小越好) 2198 4531 (1st human) 快51.5%,超越人类最佳
TriMul核竞赛 H100 运行时间µs(越小越好) 1161 1371 (1st human) 快15.3%,超越人类最佳
AtCoder Heuristic Contest ahc039 竞赛得分(越高越好) 567062 566997 (1st human) 超越人类第一名,从ALE-Agent的第5名提升至第1名
AtCoder Heuristic Contest ahc058 竞赛得分(越高越好) 848414228 847674723 (1st human) 超越人类第一名,从零开始达到SOTA
单细胞去噪 PBMC Score(越高越好) 0.71 0.64 (MAGIC) 得分提升10.9%,MSE从0.19降至0.15
单细胞去噪 Tabula Muris Score(越高越好) 0.73 0.64 (MAGIC) 得分提升14.1%,MSE从0.18降至0.14

局限与改进

本文存在多个局限性。首先,也是作者明确指出的最重要限制:当前方法仅适用于具有连续奖励的问题。对于稀疏奖励或二值奖励问题(如'是否证明了定理'),熵正则化目标的优势可能无法体现。作者将此列为最重要的未来工作方向。其次,在第二自相关不等式(AC2)上,TTT-Discover未能做出发现——其最佳构造证明了 C_2 ≥ 0.959,而AlphaEvolve已证明 C_2 ≥ 0.961,表明方法并非在所有问题上都能超越SOTA。第三,在圆堆叠问题上,TTT-Discover仅匹配现有SOTA而非超越,说明在某些问题上改进空间有限。第四,计算成本方面,虽然每问题仅需几百美元(约500美元),但仍需50步×512个rollout的大量采样,且需要可扩展的GPU评估基础设施(如Modal)。第五,MLA-Decode核的评估存在显著的AMD MI300X实例间方差,使得统计显著性难以确认——在三个实例中,最佳核各不相同,且均未能统计显著地超越人类第一名。最后,作者承认在单细胞去噪任务上,基准指标的改进可能无法转化为实际的生物学洞见,专家审稿也指出了这一点。

独立分析的弱点

TTT-Discover存在几个值得深入分析的弱点。第一,自适应 β 的KL散度约束实现细节在正文中被推迟到附录,这使得读者难以完整理解其工作机制和潜在限制——例如,KL约束的选择如何影响不同任务间的迁移性?第二,PUCT中的探索系数 c 是固定的超参数,未提供敏感性分析——在不同问题上,探索与利用的最佳平衡点可能不同。第三,每步512个rollout和8组×64的分组策略是固定的,未探索更大批量或不同分组策略的影响——更大的批量可能带来更稳定的梯度估计,但也增加成本。第四,方法依赖于问题描述诱导出连续奖励函数,对于无法定义连续奖励的领域(如创意写作、哲学推理),方法完全不适用。第五,MLA-Decode核未使用Triton进行细粒度优化,而是依赖 torch.compile(),这可能限制了进一步改进的空间。改进方向包括:将熵正则化目标扩展到离散奖励场景(例如通过奖励塑形或辅助连续信号);引入自适应批量大小,在训练后期增加采样量;探索更高效的复用策略减少不必要的重访。

未来方向

作者明确指出的最重要未来方向是将测试时训练扩展到具有稀疏或二值奖励的问题,以及非可验证领域。基于本文成果,还有多个值得探索的方向:(1)将TTT-Discover应用于更广泛的科学问题,如蛋白质结构预测的特定子任务、材料设计、药物分子优化等,这些领域通常具有连续奖励且有明确的人类专家基线;(2)研究测试时训练与测试时搜索的混合策略——何时应该训练权重,何时应该仅搜索?训练步数和采样预算如何动态分配?(3)探索多问题迁移——在单个问题上训练的经验能否迁移到相关问题?例如在GPU核优化中,在一种核类型上获得的经验能否加速另一种核类型的优化?(4)降低计算成本——当前方法需要25600次LLM调用(50步×512次),探索更高效的采样策略(如重要性采样、课程学习)可能显著降低成本;(5)研究发现能力的涌现——随着模型规模增大,TTT-Discover的改进是否会超线性增长?(6)将方法扩展到多步MDP——当前仅考虑单步(通过状态复用隐式多步),更显式的多步规划可能在复杂问题上带来优势。

复现评估

本文在可复现性方面做得相当出色。作者明确声明所有结果均使用开源模型gpt-oss-120b实现,代码已公开发布。所有发现的解决方案(包括Erdős问题的600片阶梯函数、自相关不等式的30000片阶梯函数、竞赛核代码等)均在代码库中发布并可验证。数学问题的结果可通过执行代码中的构造来直接验证。GPU核竞赛的结果已提交至GPUMode官方排行榜(TriMul A100/H100)。AtCoder竞赛的结果已提交至官方平台且公开可查。然而,复现仍存在几个挑战:(1)需要Tinker API(Thinking Machines提供的推理服务),该API可能需要单独申请;(2)GPU核评估需要Modal上的可扩展GPU基础设施;(3)完整的50步训练约需25600次LLM调用,成本约500美元;(4)某些任务的评估基础设施(如GPUMode的排行榜)可能需要额外设置。总体而言,复现难度为中等偏高,主要受算力和API访问限制,但代码和结果的开放性为验证和后续研究提供了良好基础。