Learning to Discover at Test Time — part 5/5
Cleaned-text code points [62247, 75345). Bibliography omitted here; source.md includes it.
We maintain an archive ℋ t \mathcal{H}_{t} of previously discovered states s s with reward R ( s ) ∈ R R(s)\in\mathbb{R} . To choose the next start state, we score each s ∈ ℋ t s\in\mathcal{H}_{t} by a PUCT-inspired rule, analogous to applying PUCT at a virtual root whose actions correspond to selecting a start state from the archive [ 55 , 63 , 65 , 64 ] : score ( s ) = Q ( s ) + c ⋅ scale ⋅ P ( s ) 1 + T 1 + n ( s ) , \text{score}(s)=Q(s)+c\cdot\text{scale}\cdot P(s)\frac{\sqrt{1+T}}{1+n(s)}, where n ( s ) n(s) is a visitation count, T T is the number of expanded parents so far, c > 0 c>0 is an exploration coefficient, and scale = R max − R min \text{scale}=R_{\max}-R_{\min} is the reward range over the archive. The prior P ( s ) P(s) is a linear rank distribution: P ( s ) = | ℋ t | − rank ( s ) ∑ s ′ ∈ ℋ t ( | ℋ t | − rank ( s ′ ) ) , P(s)=\frac{|\mathcal{H}_{t}|-\text{rank}(s)}{\sumop\displaylimits_{s^{\prime}\in\mathcal{H}_{t}}(|\mathcal{H}_{t}|-\text{rank}(s^{\prime}))}, where rank ( s ) ∈ { 0 , … , | ℋ t | − 1 } \text{rank}(s)\in\{0,\ldots,|\mathcal{H}_{t}|-1\} orders states by descending reward (rank 0 0 is the best state). The term Q ( s ) Q(s) uses the best one-step reachable reward m ( s ) m(s) : Q ( s ) = { m ( s ) n ( s ) > 0 R ( s ) n ( s ) = 0 . Q(s)=\begin{cases}m(s)&n(s)>0\\ R(s)&n(s)=0\end{cases}. After expanding parent p p and observing its best child reward y = max s ′ ∈ Child ( p ) R ( s ′ ) y=\max_{s^{\prime}\in\mathrm{Child}(p)}R(s^{\prime}) , we update: m ( p ) \displaystyle m(p) ← max ( m ( p ) , y ) \displaystyle\leftarrow\max(m(p),\,y) (direct parent only) n ( a ) \displaystyle n(a) ← n ( a ) + 1 ∀ a ∈ { p } ∪ Anc ( p ) \displaystyle\leftarrow n(a)+1\quad\forall a\in\{p\}\cup\mathrm{Anc}(p) (backprop visitation) T \displaystyle T ← T + 1 . \displaystyle\leftarrow T+1. For the archive update, we keep the top- 2 2 children per expanded parent (largest R R ) before inserting, then enforce a global size constraint by retaining the top- 1000 1000 states in ℋ t \mathcal{H}_{t} by R R , while always keeping the initial seed states.
##### Comparison to AlphaZero PUCT.
AlphaZero’s PUCT operates over a tree of state-action edges, selecting actions via a = arg max a [ Q ( s , a ) + c ⋅ P ( s , a ) ⋅ ∑ b N ( s , b ) / ( 1 + N ( s , a ) ) ] a=\arg\max_{a}[Q(s,a)+c\cdot P(s,a)\cdot\sqrt{\sumop\displaylimits_{b}N(s,b)}/(1+N(s,a))] , where Q ( s , a ) Q(s,a) is the mean value of simulations through edge ( s , a ) (s,a) , P ( s , a ) P(s,a) is a learned policy prior, and N ( s , a ) N(s,a) counts visits to that edge [ 65 , 64 ] . Our formulation differs in four ways: (i) Q ( s ) Q(s) tracks the maximum child reward rather than the mean, favoring optimistic expansion; (ii) P ( s ) P(s) is a rank-based prior over archived states rather than a learned action distribution; (iii) visitation counts backpropagate to all ancestors, so expanding any descendant reduces the exploration bonus for the entire lineage; and (iv) we block the full lineage (ancestors and descendants) from the current batch to encourage diversity, whereas AlphaZero uses virtual loss as a temporary penalty.
## Appendix B Mathematics
### B.1 Circle Packing
⬇
### B.2 Autocorrelation Inequalities
For autocorrelation inequalities, initial sequences are created by sampling a random value in [ 0 , 1 ] [0,1] and repeating it between 1,000 and 8,000 times (or loading a state-of-the-art construction when available). For the first inequality, the verifier computes the upper bound 2 n ⋅ max ( f ∗ f ) / ( ∑ f ) 2 2n\cdot\max(f*f)/(\sumop\displaylimits f)^{2} where f ∗ f f*f denotes discrete autocorrelation; it validates that inputs are non-empty lists of non-negative floats clamped to [ 0 , 1000 ] [0,1000] with sum ≥ 0.01 \geq 0.01 , and returns ∞ \infty for invalid constructions. For the second inequality, verifier computes the lower bound C 2 = ‖ f ∗ f ‖ 2 2 / ( ‖ f ∗ f ‖ 1 ⋅ ‖ f ∗ f ‖ ∞ ) C_{2}=\|f*f\|_{2}^{2}/(\|f*f\|_{1}\cdot\|f*f\|_{\infty}) using piecewise-linear integration for the L 2 L^{2} norm (Simpson-like rule with endpoint zeros) over the normalized interval [ − 1 / 2 , 1 / 2 ] [-1/2,1/2] . Each algorithm is run with 1 GB with 2 CPUs each and a timeout of up to 1100 seconds.
### B.3 Erdős’
We initialize TTT-Discover with random constructions of 40-100 samples around 0.5 with random perturbations. We filter out sequences with more than 1000 values in the verifier. Each algorithm is run with 1 GB with 2 CPUs each and a timeout of up to 1100 seconds.
## Appendix C Kernel engineering
For trimul, we provide the a matrix multiplication kernel that triton provides in README , mostly for syntax purposes For MLA-Decode, we first put a softmax kernel in a preliminary prompt to let the base model generate a correct but unoptimized MLA-Decode kernel, and then use that as the initial state with the earlier softmax example removed.
### C.1 Kernel evaluation details
Setup of verifier for training. We follow the exact same practice for evaluating kernel correctness and runtime as the original GPUMode competitions. Specifically, the verifier used in our training jobs uses the same code as the official GPU Mode Competition Github repository, with minor adjustment to integrate into our training codebase. The verification process includes a correctness check that compare output values between the custom kernel and a pytorch reference program under a designated precision, followed by runtime benchmarking of the custom kernel across multiple iterations. All the details in our verification procedure follow the official competition exactly, including the test cases used for correctness check and benchmarking, hyper-parameters such as matching precision and iterations used for timing, etc. We run our verifier on H100s for TriMul, and H200s for MLA-Decode, both from the Modal cloud platform.
Setup of environments for final report. For final report, we submit to the official TriMul A100/H100 leaderboard and report the runtime shown. For TriMul B200/MI300X and MLA-Decode MI300X tasks, due to an infra problem on GPU Mode’s server, we could not submit to the official leaderboard. For these tasks, we work with the GPU Mode team closely to set up our local environment, which replicates the official environment and gets GPU Mode team’s review and confirmation.
Selection protocol for best kernels. For TriMul H100 task, we select 20 kernels with the best verifier score throughout training. For other tasks, since our verifier hardware in training is different from the target hardware, we select 20 kernels with the best training scores plus 20 random correct kernels every 10 steps of training. Finally, we used our verifier with the target hardware to verify each selected kernels for three times, and submit the kernel with the smallest average runtime for final report.
### C.2 Analysis of best generated kernels
TriMul H100 kernels. The below code shows the best TriMul kernels discovered by TTT for H100 GPU. At the high level, the kernel correctly identifies a major bottleneck of the problem, which is the heavy memory I/O incurred by a series of elementwise operations, and then focuses on fusing them with Triton. Specifically, the kernel fuses: (i) operations in the input LayerNorm, (ii) elementwise activation and multiplication for input gating, and (iii) operations in the output Layernorm and output gating. As for the compute-heavy operation, which is an O ( N 3 ) O(N^{3}) matmul, the kernel converts its inputs to fp16 and delegate the computation to cuBLAS to effectively leverage the TensorCores on H100 GPU.
Compared with kernels generated early in training, the final kernel achieves a big improvement by (i) fusing more operations together, and (ii) deeper optimization of the memory access pattern inside fused kernels. For example, a kernel generated early fuses LayerNorm operations, but does not fuse the input gating process. A kernel generated in the middle of training fuses the same operations as the final kernel, but has less efficient memory access pattern in the fused kernel for output LayerNorm, gating, and output projection.
Compared with the best human leaderboard kernel, the TTT discovered kernel adopts a similar fusion strategy for the input LayerNorm and input gating. Different from human kernel, the TTT kernel does not perform as much auto-tuning of block size, which could be a limitation. However, the TTT kernel fuses the output LayerNorm and gating with output projection whereas the human kernel does not, which could explain the moderate advantage of the former.
⬇
### C.3 TTT MLA-Decode kernels filtered with Triton kernels
## Appendix D Algorithm Engineering
During the contest, AtCoder provides an official input generator, tester to evaluate program correctness, and a scoring function used for the final ranking. For training, we generate 150 test cases using seeds 0 through 149 from the input generator and run our program on each of these cases with an ALE-Bench provided C++20 container (yimjk/ale-bench:cpp20-202301) . A program receives a non-zero reward only if it passes all correctness checks and executes within the problem time limit (2 seconds) across all 150 test cases. The per-test case score is problem-specific and matches the scoring used in the AtCoder contest. For ahc039, we use ShinkaEvolve’s performance metric, which is determined by the score’s relative placement among the final contest’s scores, and for ahc058, we directly use the contest score.
For the final evaluation, we select the top three highest-scoring programs from our local training runs and submit them to the official AtCoder submission website. For our language, we specify C++23 (GCC 15.2.0). The submission is evaluated using the same scoring and validation process as the original contest, including checks for incorrect output, time limit violations, and compilation or runtime errors on AtCoder’s hidden test cases. The resulting score is used as the final evaluation.
For AHC training runs, we make a slight modification from our standard hyperparameters. For AHC039, we decrease the prompt length + thinking token limit to 22000 due to the large initial program. For AHC058, we similarly decrease the prompt length + thinking token limit to 25000 and found that a learning rate of 2 × 10 − 5 2\times 10^{-5} performed slightly better. For both AHC problems, we use a KL coefficient of 1 × 10 − 2 1\times 10^{-2} . Other hyperparameters are set to our standard values.
## Appendix E Single cell analysis
The OpenProblems benchmark provides three datasets: pancreas, pbmc and tabula. We select the Pancreas dataset to compute MSE and Poisson loss scores and use the other two datasets to assess generalization. MSE and Poisson loss scores are normalized with respect to the scores that no denoising and perfect denoising would get on this task. The main score metric in the OpenProblems denoising benchmark is the mean between the normalized MSE and the normalized Poisson. During verification, we reject all the solutions that obtain a normalized Poisson lower than 0.97 or larger than 1 so that we can focus only on improving a single metric, MSE.
In the prompt we also include instructions regarding what makes a solution taking inspiration from the Supplementary Materials of the OpenProblems paper [ 44 ] . For this specific applications, considering the size of the datasets, the memory limit is increased to 3GB. To force generalization, we reduce the time limits for the execution to 400 seconds.
We ran the OpenEvolve baseline with 25,600 samples. After sample 17,000, we observed the OpenEvolve database filling up with programs that timed out. Consequently, we selected the best program found up to that point.
Both TTT-Discover and the Best-of-25600 baselines are run with max tokens equal to 20,000.
Both MAGIC and the solution found by TTT-Discover are run with default parameters.
⬇
## Appendix F Prompts
Below we show example prompts from a sample step.
⬇
⬇
⬇
⬇
⬇
⬇
⬇
⬇
## Instructions for reporting errors
We are continuing to improve HTML versions of papers, and your feedback helps enhance accessibility and mobile support. To report errors in the HTML that will help us improve conversion and rendering, choose any of the methods listed below:
Click the "Report Issue" ( ) button, located in the page header.
Tip: You can select the relevant text first, to include it in your report.
Our team has already identified the following issues . We appreciate your time reviewing and reporting rendering errors we may not have found yet. Your efforts will help us improve the HTML versions for all readers, because disability should not be a barrier to accessing research. Thank you for your continued support in championing open access for all.
Have a free development cycle? Help support accessibility at arXiv! Our collaborators at LaTeXML maintain a list of packages that need conversion , and welcome developer contributions .