@ai4research_ucb: TTT-Discover 通过在测试问题本身上训练来推动科学发现 [ADRS Blog #25] 我们介绍 TTT-Discover,一个…

X AI KOLs Following 工具

摘要

TTT-Discover 是一个框架,它在测试时使用强化学习对单个问题训练大型语言模型,在GPU内核工程中创下新纪录,并改进了数学界限。

TTT-Discover 通过在测试问题本身上训练来推动科学发现 [ADRS Blog #25] 我们介绍 TTT-Discover,一个框架,它使用强化学习在单个测试问题上持续训练大型语言模型,而不是提示一个固定的模型。基于开源的 gpt-oss-120b 运行,TTT-Discover 在GPU内核工程和组合学中创下新纪录,包括在四种GPU类型上超越了顶级的人类TriMul内核,并改进了AlphaEvolve先前处理的两个界限。 阅读博客:https://ucbskyadrs.github.io/blog/ttt ADRS 博客系列:https://ucbskyadrs.github.io TTT 论文:https://test-time-training.github.io/discover.pdf TTT 代码:https://github.com/test-time-training/discover… 项目页面:https://test-time-training.github.io/discover/
查看原文
查看缓存全文

缓存时间: 2026/08/17 02:16

TTT-Discover:在测试时学习发现 | ADRS —— 面向系统的AI驱动研究
来源:https://ucbskyadrs.github.io/blog/ttt/

本文是AI驱动研究系统(ADRS)博客系列的一部分,探讨AI如何应用于系统研究。本期我们重点介绍TTT-Discover这项激动人心的工作!TTT-Discover通过在单个测试问题上持续训练LLM,而非使用固定模型进行提示。AlphaEvolve和OpenEvolve试图将过去的尝试结果优化为更好的提示词,但模型权重从未改变。TTT-Discover则在测试问题本身上运行强化学习,采用专为“发现”而非平均性能设计的目标函数。所有实验均在开源的gpt-oss-120b模型上进行,每个问题成本仅需数百美元。

📄论文(https://test-time-training.github.io/discover.pdf)· 💻代码(https://github.com/test-time-training/discover)· 🔍项目主页与演示(https://test-time-training.github.io/discover/)

更多ADRS内容:


分布偏移。GPUMode TriMul任务(H100)上的内核运行时分布(引自论文)。搜索基线(灰色)从固定模型中采样:其质量始终停留在较慢的左侧峰值,最高仅达5,352微秒。TTT-Discover(橙色)直接在测试问题上运行强化学习——每次更新都将采样分布向右移动,从初始策略π0经过混合精度(π10)、算子融合(π25)到更深层融合(π50),最终超越人类最佳提交结果(1,371微秒),达到1,161微秒。

问题背景

科学与工程领域充斥着需要超越已知最佳结果的问题:更快的GPU内核、更紧的开放数学问题界限、更高分的调度算法。当前的最佳方案已然存在,或位于排行榜顶端,或出自2016年的论文,任何未能超越它的尝试都毫无意义。这些都属于发现型问题

对LLM而言,这类问题几乎天生就难以解决。突破性的解决方案从未出现在训练集中,因此模型必须超越其见过的所有数据进行泛化。常规的变通方法是搜索。AlphaEvolve和OpenEvolve等方法从固定模型中采样数千次,将最佳尝试存储在缓冲区中,并将它们反馈到日益复杂的提示中。提示词和解决方案确实得到了优化,但模型本身从未改变。然而,这些尝试恰恰是该问题缺失的数据:数百个针对此特定分布外问题的解决方案,在搜索开始前它们无处可循。将它们放入提示词是利用这些数据最弱的方式。更强的方式是基于它们进行训练。

问题在于,标准强化学习的目标设定是错误的。它最大化期望奖励,因为在常规强化学习中,策略本身就是最终产物:它将被部署,并且需要表现得可靠良好。但在发现型问题中,策略是可消耗的。唯一重要的是它能产生一个打破纪录的解决方案,即使1000次尝试中有999次失败,而平均奖励训练会平滑掉这种千分之一概率的行为。每次尝试都从零开始也限制了单次尝试能够积累的结构信息量。

在TriMul(GPUMode竞赛中的GPU内核任务,运行时间单位为微秒,数值越低越好)上,使用相同模型和相同的25,600次尝试进行最优N采样(无训练)的结果为5,352微秒。

TTT-Discover保留了训练过程,并修正了目标函数。它在单个测试问题上运行强化学习,采用追逐最大值而非平均值的学习目标,以及建立在最有希望解决方案基础上的重用规则。所有实验均在开源的gpt-oss-120b模型上进行,每个问题成本仅需数百美元。

熵目标与PUCT重用

两项改动将标准强化学习转变为一种发现方法。

熵目标。标准强化学习最大化期望奖励:

max_θ  E_{y ~ π_θ}[R(y)]

TTT-Discover则最大化一个熵目标,其中轨迹的权重按奖励呈指数分布,指数参数为β:

max_θ  log E_{y ~ π_θ}[e^{β R(y)}]

当β减小时,我们回归到标准的期望奖励目标;当β增大时,更新过程将由最高奖励的轨迹主导,因此梯度追逐最大值而非平均值。实践中,固定的β值可能不稳定,因此它会针对每个状态自适应设置(详细信息请参阅论文)。

下图说明了使用β背后的直觉:

熵目标——高层直觉。每个面板显示同一组轨迹,权重按 e^βR 计算,β值递增。当β=1时,权重几乎均匀,更新行为类似标准的平均奖励强化学习。随着β增长,权重逐渐集中到最高奖励的轨迹上,直到β=12时,最佳轨迹几乎完全主导了梯度。这就是该目标背后的直觉:策略更新追逐的是尾部——那个唯一打破纪录的轨迹才是对发现真正重要的东西。

PUCT重用。为扩展有效的时间范围,并在探索新状态与重新展开有希望的状态之间取得平衡,状态选择采用AlphaZero风格:

a* = argmax_a  Q(s,a) + c · P(s,a) · √N(s) / (1 + N(s,a))

关键区别在于Q值的计算。AlphaZero将Q(s,a)设为通过该节点的模拟的平均奖励;而此处使用的是其最佳后代奖励:

Q(s,a) = max_{y ∈ desc(s,a)} R(y)

因此,一个曾经催生过更快内核的内核,无论后续有多少失败尝试,都将保留其贡献价值。

论文描述了TTT-Discover在四个领域的应用:内核工程、数学、算法工程和生物学。这里我们简要介绍前两个领域。在内核工程方面,我们聚焦于优化TriMul内核(在AlphaFold等模型中重要的操作);在数学方面,我们关注两个AlphaEvolve在2025年处理过的组合数学问题:埃尔德什最小重叠问题和第一个自相关不等式。

内核工程

在提交时,TTT-Discover找到的TriMul内核在所有四种GPU类型上均超越了人类最佳成绩。A100的运行时间从4,531微秒降至2,198微秒,H100从1,371微秒降至1,161微秒。训练过程仅在H100上计时内核,因此A100的结果来自奖励函数从未要求过的泛化能力。

获胜内核将问题视为内存瓶颈,并融合了参考实现分开运行的三组操作:输入层归一化、输入门中的sigmoid和逐元素乘法、输出层归一化与输出门。对于O(N³)的矩阵乘法,它转换为FP16并将工作交给cuBLAS。GPUMode的组织者称这与最佳人类选手使用的策略相同,但执行得更好,因为大多数人类提交融合的复杂操作较少。

跨四种GPU的TriMul结果。最佳人类提交与TTT-Discover内核在各GPU类型上的运行时间对比(微秒,越低越好)。该内核在所有场景下均击败最佳人类:A100上−51%,H100上−15%,B200上−12%,MI300X上−38%。训练仅在H100上计时内核:A100、B200和MI300X的结果来自内核在训练硬件之外的泛化能力。

加性组合数学中的两个问题

我们处理了两个同样出现在AlphaEvolve论文中的数学问题。

第一个是埃尔德什最小重叠问题,自1955年开放至今。将1到2n的数字分成两个等量堆,计算每个间隔k分隔的配对数量,并找到使最常见间隔尽可能罕见的分法。记该计数为M(n);问题在于当n增长时,M(n)/n的极限是多少。人类最佳界限由Haugland在2016年给出,为0.380927。

第二个是第一个自相关不等式。将一个非负函数与其自身卷积:该自卷积的峰值总是高于C1乘以该函数质量的平方。任何峰值足够平坦的函数都能证明C1存在一个更低的上界,因此搜索的目标就是找到这样的函数。

两者都归结为同一类优化:在范数有界的阶梯函数上,最小化相关性达到的最大值。Swinnerton-Dyer的一个结论允许用密度函数代替埃尔德什分划,而任何由此构造出的结果都能自证其界。除了评估函数值外,无需额外证明。

TTT-Discover在这两个问题上都创下了新纪录。

在不等式问题上,它达到了1.50287,超越了ThetaEvolve的1.50314(数值越低越好)。两者之间的差距大于数字本身的差距:ThetaEvolve改进了AlphaEvolve的1,319块构造,而TTT-Discover从随机函数开始,从头构建了一个30,000块的构造。

在埃尔德什问题上,AlphaEvolve已将Haugland的界限改进至0.380924。TTT-Discover达到了0.380876,改进幅度是AlphaEvolve自身改进幅度的16倍。这个600块的构造也是非对称的。此前的所有纪录保持者,包括Haugland的51块和AlphaEvolve的95块,都是对称的。

埃尔德什问题的构造。证明埃尔德什最小重叠问题(1955年开放)界限的阶梯函数。上:AlphaEvolve的95块对称构造,创下了0.380924的前纪录。下:TTT-Discover的600块构造——与之前所有纪录保持者不同,它是非对称的——达到0.380876,改进幅度是AlphaEvolve对Haugland长期界限自身改进幅度的16倍。

结论

论文中的核心对比是TTT-Discover与最优25,600次采样。相同的模型,相同的轨迹数,相同的采样预算。权重是否更新是唯一的变量,而它决定了上述所有结果。同一套方法,未经修改,在内核工程、组合数学和算法工程领域均创下纪录,而这些领域此前曾被前沿模型流水线所占据。

尝试使用

TTT-Discover使用起来非常简单!首先,定义你的问题和环境:

from ttt_discover import Environment, BaseRewardEvaluator, State, DiscoverConfig, discover

# Define your reward function
class YourReward(BaseRewardEvaluator):

    def get_reward(self, code: str, state: State) -> float:
        # ...add logic here for computing reward

        return {
            "reward": reward,
            "correctness": 1.0,
            "raw_score": raw_score,
            "msg": f"Success; raw_score={raw_score}",
            "result_construction": [], # Could reuse
            "stdout": "", # No stdout
        }

class YourEnv(Environment):
    reward_function = YourReward
    state_type = State # You may define your own state if you wish

    def get_question(self) -> str:
        state_ctx = self.initial_state.to_prompt(100, metric_name="performance")

        return f"""You are an expert mathematician specializing in combinatorial problems and computational geometry. Your task is to ... {state_ctx}."""

然后,只需训练模型!

config = DiscoverConfig(
    env_type=YourEnv,
    experiment_name="test-run",
    wandb_project="",
)

# Run discovery
discover(config)

引用

@article{ttt-discover2026,
  title   = {Learning to Discover at Test Time},
  author  = {Yuksekgonul, Mert and Koceja, Daniel and Li, Xinhao
             and Bianchi, Federico and McCaleb, Jed and Wang, Xiaolong
             and Kautz, Jan and Choi, Yejin and Zou, James
             and Guestrin, Carlos and Sun, Yu},
  journal = {ICML},
  year    = {2026}
}

为ADRS博客系列贡献内容!

AI驱动研究系统(ADRS)倡议是一项开放的协作事业,旨在探索AI如何加速科学发现本身,从进化算法到优化现实系统。如果你曾构建、优化或实验过AI驱动的研究工具,我们很乐意听取你的经验。请在ADRS博客系列中分享你的经验、见解或案例研究。

👉 通过电子邮件联系我们:[email protected]

💬 加入我们:join.slack.com/t/adrs-global (https://join.slack.com/t/adrs-global/shared_invite/zt-3fgme22n5-PKYyAc9aIeTyX5iSQTKIoA) 和 Discord (https://discord.gg/mz8gAnpt)

相似文章

模块化TTT:将测试时训练重新构想为可组合模块

Hugging Face Daily Papers

本文介绍了模块化TTT,这是一个将测试时训练的内部学习器表示为有向无环图的框架,能够对组件进行系统的消融和组合。作者在1000亿个token上训练了4.1亿和14.5亿参数规模的模型,取得了与GatedDeltaNet相当的性能。

用 LLM 优化 LLM:面向测试时扩展的智能体发现方法

Hugging Face Daily Papers

本文提出了 AutoTTS,这是一种环境驱动的框架,通过将测试时扩展(TTS)策略的发现过程形式化为控制器合成,自动发现用于大型语言模型(LLM)的测试时扩展策略。该框架在数学推理基准测试上展示了更优的准确率-成本权衡,且计算开销极小。

测试时训练破坏安全护栏

arXiv cs.LG

本文识别了测试时训练(TTT)的三种威胁模型,攻击者可利用这些模型绕过LLM的安全过滤器,实现高攻击成功率。研究结果表明,TTT引入了新的漏洞,破坏了现有的安全护栏。