TTT-Discover: Learning to Discover at Test Time

matt_d1 pts0 comments

TTT-Discover: Learning to Discover at Test Time | ADRS — AI-Driven Research for Systems All posts

This post is part of the AI-Driven Research for Systems (ADRS) blog series, where we explore how AI can be applied to systems research. We feature exciting work on TTT-Discover this week!

TTT-Discover keeps training the LLM on a single test problem instead of prompting a frozen one. AlphaEvolve and OpenEvolve stuff past attempts into better prompts; the weights never move. TTT-Discover runs RL on the test problem itself, using an objective built for discovery rather than average performance. Everything runs on the open gpt-oss-120b for a few hundred dollars a problem.

📄 Paper · 💻 Code · 🔍 Webpage & Demo

More from ADRS:

✍️ Previous Blogs: https://ucbskyadrs.github.io/

📝 ADRS Paper: https://arxiv.org/abs/2510.06189

👩‍💻 ADRS Code: github.com/UCB-ADRS/ADRS

💬 Join us : join.slack.com/t/adrs-global and Discord

Follow us : x.com/ai4research_ucb

Distribution shift. Kernel-runtime distribution on the GPUMode TriMul task (H100), reproduced from the paper. The search baseline (grey) samples from a frozen model: its mass stays pinned at the slow left peak, topping out at 5,352 µs. TTT-Discover (orange) runs RL on the test problem itself — each update shifts the sampling distribution rightward, from the initial policy π₀ through mixed precision (π₁₀) and operator fusion (π₂₅) to deeper fusion (π₅₀), ending past the best human submission (1,371 µs) at 1,161 µs.

The Problem

Science and engineering are full of problems where the goal is to beat the best known result: a faster GPU kernel, a tighter bound on an open math problem, a higher-scoring scheduling algorithm. The current best already exists, at the top of a leaderboard or in a 2016 paper, and anything short of beating it counts for nothing. These are discovery problems.

They are hard for LLMs almost by definition. The record-beating solution appears in no training set, so the model has to generalize past everything it has seen. The standard workaround is search. Methods like AlphaEvolve and OpenEvolve sample a frozen model thousands of times, store the best attempts in a buffer, and feed them back into ever-richer prompts. The prompts and the solution tend to improve, but the model never does. Yet those attempts are exactly the data the problem was missing: hundreds of solutions to this specific out-of-distribution problem, which existed nowhere before the search started. Putting them in a prompt is the weakest way to use them. The stronger way is to train on them.

The catch is that standard RL is built for the wrong goal. It maximizes expected reward, because in normal RL the policy is the product: it will be deployed and needs to be reliably good. In discovery the policy is disposable. All that matters is that it produces one record-beating solution, even if 999 out of 1,000 attempts fail, and average-reward training would smooth that one-in-a-thousand behavior away. Restarting every attempt from a blank slate also caps how much structure a single attempt can build.

On TriMul, a GPU-kernel task from the GPUMode competition (runtime in µs, lower is better), best-of-N sampling with the same model and the same 25,600 attempts, no training at all, produced 5,352 µs.

TTT-Discover keeps the training and fixes the objective. It runs RL on the single test problem, with a learning objective that chases the maximum rather than the mean, and a reuse rule that keeps building on the most promising solutions. Everything runs on the open gpt-oss-120b for a few hundred dollars a problem.

Entropic Objective and PUCT reuse

Two changes turn standard RL into a discovery method.

The entropic objective. Standard RL maximizes expected reward:

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

TTT-Discover instead maximizes an entropic objective where rollouts are weighted exponentially by reward with a parameter β:

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

As β gets smaller we recover the standard expected-reward objective; as β grows the update is dominated by the highest-reward rollouts, so the gradient chases the max rather than the mean. In practice a fixed β is brittle, so it's set adaptively per state (see the paper for this detail).

Here a figure that describes the intuition behind the use of β:

The entropic objective — high-level intuition. Each panel shows the same set of rollouts, weighted by e^βR for increasing β. At β = 1 the weights are nearly uniform and the update behaves like standard average-reward RL. As β grows, weight concentrates on the highest-reward rollouts, until at β = 12 the best rollout dominates the gradient almost entirely. This is the intuition behind the objective: rather than improving the average, the policy update chases the tail — the single record-setting rollout is what matters for discovery.

PUCT reuse. To extend the effective horizon and balance exploring new states against re-expanding promising ones, states are selected AlphaZero-style:

a* = argmax_a...

discover problem objective adrs reward test

Related Articles