Source-linked AI summary
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
TL;DR
Scientific discovery requires solving problems beyond existing knowledge, a setting where frozen-LLM search cannot itself learn new ideas. TTT-Discover performs reinforcement learning on experience from the specific test problem and prioritizes the most promising solutions. It achieves new state-of-the-art results in almost all attempted problems across four domains using an open model.
Problem
Scientific discovery requires solving a specific problem beyond existing knowledge rather than merely generalizing across familiar tasks.
Method
TTT-Discover performs reinforcement learning in an environment defined by one test problem, training the policy on accumulated search experience and prioritizing promising solutions.
Results
TTT-Discover achieves new state-of-the-art performance in almost all attempted problems across mathematics, GPU kernels, algorithms, and biology, including 50% faster TriMul kernels than the top human kernel on A100s.
Takeaways & Limitations
The results show that an open model can be adapted at test time to discover strong solutions across diverse scientific and engineering problems.
Takeaways & Limitations
Benchmark performance in the single-cell denoising application does not guarantee biological validity for downstream tasks.
Abstract
from arXiv · showhide
How can we use AI to discover a new state of the art for a scientific problem? Prior work in test-time scaling, such as AlphaEvolve, performs search by prompting a frozen LLM. We perform reinforcement learning at test time, so the LLM can continue to train, but now with experience specific to the test problem. This form of continual learning is quite special, because its goal is to produce one great solution rather than many good ones on average, and to solve this very problem rather than generalize to other problems. Therefore, our learning objective and search subroutine are designed to prioritize the most promising solutions. We call this method Test-Time Training to Discover (TTT-Discover). Following prior work, we focus on problems with continuous rewards. We report results for every problem we attempted, across mathematics, GPU kernel engineering, algorithm design, and biology. TTT-Discover sets the new state of the art in almost all of them: (i) Erdős' minimum overlap problem and an autocorrelation inequality; (ii) a GPUMode kernel competition (up to $2\times$ faster than prior art); (iii) past AtCoder algorithm competitions; and (iv) denoising problem in single-cell analysis. Our solutions are reviewed by experts or the organizers. All our results are achieved with an open model, OpenAI gpt-oss-120b, and can be reproduced with our publicly available code, in contrast to previous best results that required closed frontier models. Our test-time training runs are performed using Tinker, an API by Thinking Machines, with a cost of only a few hundred dollars per problem.
1 Introduction
Scientific discovery requires ideas beyond existing knowledge, making out-of-distribution generalization difficult for AI. TTT-Discover addresses this by training the model on problem-specific experience while searching, and reports state-of-the-art results across diverse continuous-reward problems.
- Scientific discovery requires ideas beyond both a model’s training data and humanity’s existing knowledge.
- Prior test-time search prompts a frozen LLM to generate many attempts, but the model itself cannot internalize new ideas from those attempts.
- TTT-Discover continues training the LLM on experience generated while solving the specific test problem.
- TTT-Discover uses a learning objective and search subroutine that favor the most promising solutions across mathematics, GPU kernels, algorithms, and biology.
- The method reports results for every attempted problem and achieves new state-of-the-art performance in almost all of them using an open model.
2 Preliminaries
The paper formalizes scientific discovery as finding a candidate solution whose reward exceeds the best known reward. It defines search environments around candidate states, problem-induced transitions, and continuous rewards, then considers exploration and reuse of prior solutions.
- A scientific problem induces an environment in which a candidate solution is the state and its quality is measured by a continuous reward.
- A discovery is a state s whose reward R(s) exceeds the best-known reward r_sota.
- Discovery problems seek a state with a large reward improvement over the state of the art.
- Search generates actions from the problem description and an initial solution, with the environment defining the transition to a new state.
- Starting from an empty solution encourages exploration, while state reuse warm-starts attempts from previous solutions to exploit promising directions.
- Evolutionary search reuses prior actions or states through heuristics involving fitness, diversity, mutation, and crossover.
3 Learning to Discover at Test Time
TTT-Discover adapts reinforcement learning to a single discovery problem, where the objective is one exceptional solution rather than average policy performance. It combines online training on accumulated attempts with state reuse and search choices that emphasize high-reward outcomes while retaining exploration.
- Online training: TTT-Discover trains the policy online on its own search attempts accumulated in a buffer, rather than only improving subsequent prompts.
- Discovery-specific objective: Discovery differs from standard reinforcement learning because success depends on finding one state-of-the-art solution, not maximizing average reward for repeated deployment.
- Naive RL limitations: Naive expected-reward optimization can favor safe actions, while starting every attempt from scratch limits the effective search horizon.
- TTT-Discover: The method introduces two components to address these shortcomings: a training objective and a state-reuse routine.
- Entropic objective: The entropic objective favors maximum-reward actions, approaching a max objective as β →∞, while adaptive β avoids instability early and vanishing advantages later.
- PUCT reuse: PUCT selects initial states using maximum child reward, reward-ranked priors, visit counts, and an exploration bonus to balance exploitation with under-visited states.
- TTT-Discover: TTT-Discover combines the entropic objective J_β(s) for training with PUCT for reuse.
- Implementation: The reported implementation uses gpt-oss-120b with LoRA training and batches of 512 rollouts per step.
4 Applications
TTT-Discover is evaluated across mathematics, GPU kernel engineering, algorithm design, and biology, achieving new state-of-the-art results in most attempted problems while also revealing concrete scope limits.
- Mathematics: TTT-Discover improves the Erdős’ Minimum Overlap upper bound to 0.380876, surpassing AlphaEvolve’s 0.380924 with a 600-piece asymmetric construction.Its improvement over AlphaEvolve is reported as 16 times larger than AlphaEvolve’s improvement over the previous state of the art.
- Mathematics: TTT-Discover proves the upper bound C1 ≤1.50286 using a 30,000-piece step function, while making no discovery on the second autocorrelation inequality.The best second-inequality construction certified 0.959, compared with AlphaEvolve’s tighter lower bound of 0.961.
- GPU kernel engineering: TTT-Discover reaches state-of-the-art TriMul performance across all GPU types, exceeding the best human submissions by more than 15% uniformly.On A100s, its best kernel is 50% faster than the top human kernel despite A100 runtime not being used as the reward during training.
- GPU kernel engineering: The TriMul kernels reduce memory traffic and launch overhead by fusing pointwise operations, while using FP16 and vendor libraries for the compute-heavy matmul.The matmul has O(N^3) complexity and is delegated to cuBLAS or rocBLAS to leverage TensorCores or MatrixCores.
- Algorithm engineering: For AtCoder heuristic competitions, the discovered algorithms would have placed first in both competitions, marginally improving the best human result on ahc039 and outperforming all submissions on ahc058.On ahc039, TTT-Discover improved the ALE-Agent solution from fifth place to first place.
- Biology: In single-cell denoising, TTT-Discover consistently improves results on both datasets by adding gene-adaptive transform ensembling, low-rank SVD refinement, and log-space polishing to MAGIC.These steps directly optimize the benchmark metric, but benchmark performance does not guarantee biological validity for downstream tasks.
5 Related Works
Related work develops continual and test-time training around instance-specific adaptation, while TTT-Discover targets one exceptional solution to one test problem. Its distinctive contribution is tailoring learning and reuse to discovery rather than generalization or average performance.
- Test-time training reflects two human-learning properties: personalization to an individual instance and learning without a strict training–testing boundary.
- Continual learning conventionally updates models as data distributions change, whereas test-time training formulates a potentially different learning problem for each test instance.
- Neighbor-based test-time training fine-tunes models on training examples near a test instance, increasing effective capacity and improving shifted-distribution generalization.
- Recent reasoning systems apply test-time training through generated curricula or augmentations, including AlphaProof and methods for few-shot reasoning tasks.
- MiGrATe, ThetaEvolve, and EvoTune combine per-instance reinforcement-learning updates with replay or reuse, while TTT-Discover tailors both objective and reuse rule to discovery.
- Unlike work that trains on one training example or an entire test set, TTT-Discover trains on the test problem itself to find one exceptional solution.
6 Future Work
The current method is restricted to problems with continuous rewards, motivating future work on sparse, binary, and non-verifiable settings.
- TTT-Discover currently applies only to problems with continuous rewards, leaving sparse or binary rewards and non-verifiable domains for future work.
A Training details
The experiments use fixed hyperparameters nearly throughout, with a KL penalty coefficient of 0.1 for most applications and 0.01 for algorithm engineering.
- The KL penalty coefficient is fixed at 0.1 for almost all applications and 0.01 for algorithm engineering.
A.1 Entropic utility objective
TTT-Discover uses an entropic utility objective with adaptive temperature control and archive search that emphasizes promising improvements. Its reuse mechanism favors optimistic expansion while limiting update concentration and archive growth.
- Entropic utility objective: The entropic objective reweights rollouts toward higher rewards, with its gradient producing an entropic advantage for policy updates.
- Adaptive β: Adaptive β is selected per initial state by enforcing a KL budget on the entropically tilted distribution.
- Batch estimator: The batch estimator starts from N rollouts with a uniform empirical distribution and computes reward-dependent reweighting for the update.
- Archive search: Archive search scores prior states with a PUCT-inspired rule, using visitation counts, reward-based priors, and optimistic values.
- Archive update: After expansion, the archive retains the top-2 children per parent and at most the top-1000 states globally while preserving initial seeds.
- Archive search: Unlike AlphaZero PUCT, the archive rule tracks maximum child reward, uses rank-based priors, and propagates visits to ancestors.
B.1 Circle Packing
The circle-packing construction initializes circles in staggered arrangements, then optimizes their centers and radii to maximize total radius while enforcing boundary and non-overlap constraints.
- Initialization: The procedure generates staggered or hexagonal initial arrangements with specified row and column patterns.It varies circles per row, uses refined horizontal and vertical spacing, and adds extra circles in one construction.
- Initialization: Initial circle centers and radii are flattened into one optimization-variable array.The construction assigns radii by row before collecting each center and radius into the flattened representation.
- Optimization: The objective minimizes the negative sum of radii, thereby maximizing the total packed radius.The objective is implemented as a negative total because the optimizer performs minimization.
- Constraints: Boundary constraints keep every circle inside the square while pairwise distance constraints prevent overlaps.The pairwise constraint requires center distance to be at least the sum of the two radii.
- Optimization: The constrained optimization uses Sequential Least Squares Programming with refined parameters, then extracts optimized centers, radii, and their sum.The returned quantities are the optimized geometry and the aggregate radius objective.
B.2 Autocorrelation Inequalities
The autocorrelation experiments construct nonnegative sequences and evaluate them with verifier-defined inequalities involving discrete autocorrelation norms and numerical integration.
- Construction: Initial sequences repeat a random value between 1,000 and 8,000 times, or load an available state-of-the-art construction.The verifier rejects invalid sequences and returns ∞ for invalid constructions.
- Verification: The first verifier computes an upper bound using 2n·max(f ∗f )/(Pf )2 and discrete autocorrelation.Inputs must be non-empty lists of non-negative floats clamped to [0,1000] with sum ≥0.01.
- Verification: The second inequality uses 2/(∥f ∗f ∥1 ·∥f ∗f ∥∞) with piecewise-linear integration for the L2 norm.The integration uses endpoint zeros over the normalized interval [−1/2,1/2].
B.3 Erdős’
The experiments initialize TTT-Discover with randomized constructions and task-specific kernel states before optimization or search.
- Initialization: TTT-Discover starts from random constructions of 40-100 samples centered around 0.5 with random perturbations.Sequences with more than 1000 values are filtered out before evaluation.
- Initialization: For trimul, the initial state is a Triton README matrix-multiplication kernel provided mainly for syntax.For MLA-Decode, a preliminary softmax prompt produces a correct but unoptimized kernel used as the initial state.
C.1 Kernel evaluation details
Kernel evaluation uses competition-aligned correctness and runtime verification, with target-hardware validation and a submission-selection protocol.
- Verification: The training verifier follows the official GPU Mode competition code for kernel correctness and runtime evaluation, with minor integration changes.Correctness is checked against a PyTorch reference program.
- Reporting: Some tasks could not be submitted to the official leaderboard because of a GPU Mode server infrastructure problem.For those tasks, the authors used a locally replicated environment reviewed and confirmed by the GPU Mode team.
- Selection: For TriMul H100, the procedure selects 20 kernels with the best verifier scores throughout training.For other tasks, it combines 20 top-scoring training kernels with 20 random correct kernels every 10 training steps.
- Selection: Selected kernels are reverified three times on target hardware, and the kernel with the smallest average runtime is submitted.This accounts for differences between training-verifier hardware and target hardware.
C.2 Analysis of best generated kernels
The best TriMul kernel combines fused operations with optimized memory access, while using mixed precision and cuBLAS for the compute-heavy matrix multiplication. Compared with human and earlier generated kernels, these choices improve performance but leave auto-tuning as a limitation.
- Kernel design: The kernel fuses input LayerNorm, input gating, output LayerNorm, output gating, and output projection to reduce memory I/O.The compute-heavy O(N^3) matmul is converted to fp16 and delegated to cuBLAS to use H100 TensorCores.
- Kernel design: The final kernel improves over earlier generated kernels by fusing more operations and optimizing memory access within the fused kernels.The middle-training kernel already fuses the same operations but accesses memory less efficiently around output LayerNorm, gating, and projection.
- Comparison with human kernels: Compared with the best human kernel, TTT uses similar fusion for input LayerNorm and input gating but additionally fuses output operations with output projection.The authors identify reduced auto-tuning of block size as a possible limitation and suggest the additional output fusion may explain the advantage.
- Implementation stages: The implementation contains separate fused stages for row-wise LayerNorm, projection and gating with an optional mask, and hidden-dimension LayerNorm followed by gating and final linear projection.The code uses fp32 reductions or accumulation alongside fp16 outputs and matrix multiplications.
C.3 TTT MLA-Decode kernels filtered with Triton kernels
This section reports results for TTT MLA-Decode kernels filtered to retain Triton implementations.
- Results: The section presents results for TTT MLA-Decode kernels filtered with Triton kernels.The supplied caption identifies the table as the results display for this filtered set.
- Results: Table 10 is the designated results table for the Triton-filtered MLA-Decode kernels.The caption provides the table’s scope but no numerical comparison.
D Algorithm Engineering
The algorithm-engineering evaluation trains on generated AtCoder cases and submits the highest-scoring programs to the official platform. Final scores use the original contest’s correctness, runtime, and compilation checks.
- Training evaluation: Training uses 150 AtCoder test cases generated with seeds 0 through 149 in an ALE-Bench C++20 container.A program receives non-zero reward only after passing all correctness checks within the 2-second limit.
- Final evaluation: The final evaluation submits the three highest-scoring programs from local training runs to the official AtCoder website.This procedure selects candidate programs for evaluation on the contest platform.
- Final evaluation: Final AtCoder scores apply the original contest’s scoring and validation process, including hidden-test correctness, time-limit, compilation, and runtime checks.Submissions use C++23 with GCC 15.2.0.
E Single cell analysis
The single-cell evaluation uses the Pancreas dataset for optimization and pbmc and tabula for generalization. It constrains Poisson loss while optimizing MSE, and tests denoising methods under fixed computational and baseline conditions.
- Benchmark setup: The OpenProblems benchmark provides pancreas, pbmc, and tabula datasets, with Pancreas used for scoring and the other two for generalization.Scores use normalized MSE and Poisson loss, averaged into the benchmark’s main metric.
- Benchmark setup: Verification rejects solutions with normalized Poisson below 0.97 or above 1, focusing optimization on MSE.This keeps the Poisson component within a specified range while improving the second metric.
- Evaluation constraints: The evaluation increases memory to 3GB and reduces execution limits to 400 seconds to encourage generalization.The prompt also asks solutions to draw inspiration from the OpenProblems supplementary materials.
- Denoiser design: The denoiser uses adaptive, gene-wise diffusion and can add log-space polishing targeted at the MSE representation.The implementation includes diffusion weighting, residual smoothing, variance matching, and optional final smoothing steps.
- Denoiser design: The method supports multiple variance-stabilizing transforms, dropout-based gene filtering, and optional diffusion controls such as final_smooth_weight and final_smooth_t.The output is a non-negative float64 denoised count matrix.
F Prompts
The prompts define search tasks over mathematical constructions and denoising procedures, while constraining the generated programs and evaluation calls. They also provide prior methods, targets, and examples of measured improvements to guide test-time discovery.
- Task formulations: Prompts ask models to generate non-negative step-function heights or coefficient sequences that optimize an evaluation function.The tasks include minimizing the Erdős minimum-overlap constant and generating sequences for autocorrelation inequalities.
- Search interface: The search functions can call evaluate_sequence repeatedly and must return the best sequence found within 1000 seconds, with non-negative entries.The prompt permits arbitrary search methods but imposes execution-time and output-validity constraints.
- Search guidance: The prompts encourage departing from prior algorithms, exploring different starting points, and adjusting heuristics or hyperparameters to obtain meaningful improvements.One prior approach uses coarse-to-fine stochastic search followed by projected gradient ascent.
- Execution constraints: Program rules require a named entry function, top-level helpers, limited scientific-library use, at most two CPUs, and no filesystem or network I/O.The required entry function differs by task, including propose_candidate and construct_function.
- Mathematical prompts: The Erdős prompt specifies h:[0,2]→[0,1], unit integral, discretization over [0,2], and a target C5≤0.38080 versus the current C5≤0.38092.The overlap objective is defined as C5 = max_k ∫ h(x)(1−h(x+k)) dx.
- Denoising prompt: Example denoising prompts describe noisy single-cell RNA-seq counts and report a lower-bound change from 0.6666666667 to 0.9235566275.The stated target is a lower bound close to at least 0.97 for a construction of length 1024.