PaperlayerView exact revision on arXiv

Source-linked AI summary

Surprisal-Guided Selection: Compute-Optimal Test-Time Strategies for Execution-Grounded Code Generation

Jarrod Barnes

arXiv:2602.07670v1cs.LGcs.AI

TL;DR

The paper asks whether gradient-based test-time adaptation is preferable to search for dense-reward verifiable execution-grounded tasks, and compares these strategies under controlled compute. Across KernelBench, search substantially outperforms minimal adaptation, while surprisal-guided selection recovers oracle-level performance by targeting correct samples in the distribution tail.

  • Problem

    The paper asks whether adaptation is the right compute allocation strategy for dense-reward VEG tasks or whether simple search suffices.

  • Method

    The study compares Best-of-N sampling, batch test-time training, and related selection strategies under matched compute with a shared deterministic KernelBench evaluator.

  • Results

    Best-of-N achieves 90% task success (18/20) at K = 64, versus 30.6% for TTT’s best checkpoint, while surprisal-guided selection reaches 80% versus 50% for most-confident selection.

  • Takeaways & Limitations

    For dense-reward VEG tasks, compute should favor sample diversity and intelligent selection rather than gradient adaptation.

  • Takeaways & Limitations

    The experiments use one 120B-parameter model and only KernelBench L1 tasks; transfer to other architectures, sizes, and harder levels remains untested.

Abstract

from arXiv · show

Test-time training (TTT) adapts language models through gradient-based updates at inference. But is adaptation the right strategy? We study compute-optimal test-time strategies for verifiable execution-grounded (VEG) tasks, domains like GPU kernel optimization where a deterministic evaluator provides dense, continuous reward signals. Using KernelBench as our testbed and a 120B-parameter model (GPT-OSS-120B with LoRA adaptation), we find that search outperforms minimal adaptation (1-5 gradient steps): Best-of-N sampling achieves 90% task success (18/20 tasks) at K=64 across the full KernelBench L1 eval set while TTT's best checkpoint reaches only 30.6% (3-seed mean), with TTT's "equivalent K" falling below 1, worse than single-sample inference. The failure mode is over-sharpening: gradient updates collapse diversity toward mediocre solutions rather than discovering optimal ones. Our main contribution is surprisal-guided selection: selecting the highest-surprisal (lowest-confidence) correct sample yields 80% success vs. 50% for most-confident selection, a 30% improvement. Extending to surprisal-guided-top3 matches oracle performance at 100%. This zero-cost strategy, validated through length-controlled analysis, recovers oracle performance. For dense-reward VEG tasks, compute should be allocated to sample diversity and intelligent selection rather than gradient adaptation. The surprisal-guided selection principle may generalize to other execution-grounded domains where optimal solutions occupy the distribution tail.

1. Introduction

The paper asks whether gradient-based adaptation or search is compute-optimal for dense-reward verifiable execution-grounded tasks. Experiments show search preserves diversity and, with surprisal-guided selection, finds tail solutions that adaptation misses.

  • Motivation: VEG tasks use deterministic evaluators that provide ground-truth, continuous feedback without human or AI judgment.KernelBench evaluates correctness and runtime speedup, with speedup ranging from 0x to 10x+.
  • Search versus adaptation: Best-of-N search achieves 90% task success (18/20) at K =64, whereas TTT’s best checkpoint reaches 30.6% and equivalent K < 1.The comparison uses GPT-OSS-120B with LoRA adaptation across all 20 KernelBench L1 evaluation tasks.
  • Failure mode: TTT fails through over-sharpening: gradient updates collapse diversity toward mediocre early successes instead of reaching optimal kernels in the distribution tail.The mechanism is described as concentration toward the mode, which misses rare high-quality solutions.
  • Surprisal-guided selection: Surprisal-guided selection reaches 80% success versus 50% for confidence-guided selection, a 30% improvement with zero additional compute.Selecting among correct samples uses the highest-surprisal, lowest-log-probability candidate.
  • Surprisal-guided selection: Surprisal-guided-top3 matches oracle performance at 100%, supporting selection from the high-surprisal tail for dense-reward VEG tasks.The principle may generalize where rare, high-quality solutions occupy low-probability regions.
  • Search versus adaptation: Best-of-N saturates at K =16 with 99.9% success on 5-task subsets and 90% on the full 20-task evaluation.The reported implication is to invest in sample diversity rather than 1-5 gradient updates.

2. Related Work

Related work frames the paper against test-time training, distribution-sharpening accounts of reinforcement learning, confidence-based selection, and kernel-optimization systems. The paper instead studies how to allocate test-time compute for capable policies in execution-grounded settings.

  • Test-Time Training versus Search: TTT-Discover reports strong results with approximately 50 adaptation steps on sparse-reward discovery tasks, while this paper compares search with 1-5 step adaptation on dense-reward VEG tasks.The paper attributes the differing dynamics partly to reward density, while qualifying the explanation as likely.
  • Distribution Sharpening and Over-Fitting: The paper connects TTT’s diversity collapse to distribution sharpening and prior observations that execution-reward reinforcement learning can collapse onto narrow ideas.Its extension is to dense execution-grounded optimization, where optimal solutions may already exist in the policy distribution.
  • Selection Strategies and Model Confidence: Prior selection work favors confidence or reward-model reranking, whereas this paper operationalizes surprisal as a selection strategy for execution-grounded code generation.The reported gain is 30 percentage points with zero additional compute.
  • Verifiable Execution-Grounded Tasks: VEG tasks provide deterministic feedback for filtering correct samples before surprisal-based selection, with KernelBench measuring correctness and continuous speedup.Related VEG domains include assembly superoptimization and formal theorem proving.
  • Kernel Optimization: Prior kernel-optimization systems emphasize train-time reinforcement learning, evolutionary search, memory, or hierarchical policies; this paper studies test-time compute allocation for a capable policy.The comparison includes systems such as CUDA-L2, Magellan, AccelOpt, Dr. Kernel, CUDA-L1, and QiMeng-Kernel.

3. Method

The method separates train-time policy preparation from matched-compute test-time strategy comparisons using a shared deterministic evaluator. It compares sampling, adaptation, and self-distillation while controlling the starting checkpoint and decoding conditions.

  • Dual-Loop Architecture: The dual-loop design trains a base policy with RLVR and compares test-time strategies under matched compute using the same evaluator.The outer loop uses 80 KernelBench training tasks, while the inner loop evaluates held-out tasks.
  • Outer Loop: The shared base checkpoint comes from GRPO with LoRA training on 80 KernelBench L1 tasks and provides 98.4% correctness with 0.87x mean speedup.This checkpoint initializes all test-time strategies.
  • Inner Loop: The inner loop compares Best-of-N search, batch TTT with Best-of-Adaptation checkpoint selection, and SDPO with or without execution feedback.Best-of-N can use oracle, surprisal-guided, or random selection.
  • Controlled Comparison: All strategies share temperature 0.25, maximum length 1024, and the same base checkpoint; the variable is sampling diversity versus gradient adaptation.The evaluator defines correctness through functional equivalence and reward as baseline time divided by kernel time.
  • Test-Time Training: Test-time adaptation processes 5 tasks jointly, samples K =32 rollouts per task, and updates a rank-16 LoRA adapter for 1-5 steps.Each step produces 160 samples and uses a GRPO gradient update.
  • Best-of-Adaptation: Best-of-Adaptation selects the checkpoint with maximum fast 1 validation score across the adaptation trajectory rather than assuming the final checkpoint is best.Early stopping with P =1 matched oracle selection in the experiments.
  • SDPO: SDPO replaces scalar reward advantages with token-level self-distillation conditioned on execution feedback, using β = 1.0 for distillation strength.The teacher scores tokens sampled by the student.

4. Experimental Setup

The experiments compare Best-of-N sampling with batch test-time training under matched compute, using shared evaluation and multiple KernelBench L1 subsets.

  • Budget matching: 320 rollouts are used per comparison, either independently for Best-of-N or as gradient signal for batch TTT.Temperature, token limit, evaluator mode, checkpoint, and system prompt are held constant; only rollout use differs.
  • Baselines: Best-of-N samples K = 64 candidates per task and selects the highest fast 1 result.All K samples undergo shared correctness checking before selection.
  • Metrics: fast 1 is the primary metric because it requires both functional correctness and speedup > 1x.Correctness and mean speedup are also reported as secondary metrics.
  • Evaluation sets: The evaluation covers two five-task subsets and an extended full 20-task KernelBench L1 evaluation.The extended evaluation includes ten additional tasks and reports Best-of-N, per-sample, and logprob analyses.
  • Selection strategies: Five selection strategies are compared among K = 64 samples per task, including surprisal-guided strategies motivated by an inverse confidence–quality relationship.The highest-surprisal correct samples are described as potentially creative, hardware-aware optimizations.
  • Scaling: Best-of-N scaling saturates at K = 16 with 99.9% success on the reported scaling curve.The figure contrasts this scaling behavior with TTT BoA at 30.6% and K = 1 random sampling at 53.3%.

5. Results

Under matched compute budgets, Best-of-N search decisively outperforms minimal TTT adaptation on dense-reward VEG tasks. TTT’s over-sharpening collapses diversity, while surprisal-guided selection recovers strong performance from correct samples.

  • Main result: 99.9% success at K =16 shows that Best-of-N performance saturates with a modest sampling budget.Beyond K =16, marginal gains are near-zero.
  • Why adaptation fails: TTT gradient updates over-sharpen the distribution around early mediocre successes, destroying diversity needed to find optimal kernels in the low-probability tail.Performance peaks after 1–2 steps, then regresses; seed variance and negative NLL–speedup correlations provide additional evidence.
  • Surprisal-guided selection: Surprisal-guided-top3 matches oracle performance at 100% success by evaluating three high-surprisal correct samples and selecting the fastest.The selection signal depends on sufficient within-task logprob variance, which appears in 9 of 20 tasks.
  • Robustness and evaluation: The full 20-task result reaches 90% success, with failures attributed to a task-level ceiling and a model capability gap rather than search itself.The evaluation uses five performance trials per kernel instead of the full 50-trial KernelBench protocol.

6. Discussion

VEG tasks may constitute a distinct test-time compute regime because dense continuous rewards and deterministic evaluation favor sampling diversity and intelligent selection over gradient adaptation. The paper identifies reward density as a key variable, while noting that its TTT setup leaves stronger adaptation mechanisms and budgets untested.

  • Dense continuous rewards and deterministic evaluation make sample diversity with intelligent selection a potentially more efficient strategy than gradient adaptation.VEG speedup rewards range continuously from 0x to 10x+, and execution feedback can make an AI teacher redundant.
  • 1–2 adaptation steps mark the minimum signal threshold before over-sharpening degrades diversity in dense-reward kernel optimization.This threshold is measured from 160 diverse samples and may depend on reward density.
  • ∼50 steps in TTT-Discover versus 1–2 optimal steps here highlights different adaptation dynamics between sparse-reward discovery and dense-reward VEG tasks.The paper frames these findings as complementary and identifies reward density as the key variable determining optimal adaptation duration.
  • The TTT comparison remains limited because the experiments use vanilla GRPO, 320 rollouts, and no entropy regularization or reuse buffers.Whether TTT-Discover’s mechanisms and larger budget prevent over-sharpening in dense-reward VEG domains remains untested.
  • The comparison targets budgeted inference with fast-1 success, not open-ended discovery that returns the best solution across all adaptation steps.These objectives differ from TTT-Discover’s goal of finding a new state-of-the-art solution.
  • Zero-evaluation discovery remains a future direction requiring physics-grounded world models and internal simulations of code–hardware interactions.The paper presents this as a research direction rather than a demonstrated capability.
  • An entropy-regularized baseline, deployable surprisal-first ranking, and reward sparsification experiments are proposed to test the scope of the conclusions.These experiments would examine whether over-sharpening is specific to vanilla GRPO, operationalize selection before correctness filtering, and test reward-density effects.

7. Limitations

The evaluation is narrow in model, task level, statistical power, protocol, and selection coverage. These boundaries constrain how broadly the reported surprisal and adaptation dynamics should be generalized.

  • All experiments use one 120B-parameter model and L1 tasks, leaving transfer to other architectures and harder L2/L3 levels untested.The paper states that harder levels may exhibit different dynamics.
  • The TTT limitation is specific to the vanilla GRPO configuration used in these experiments.The paper distinguishes this setup from TTT-Discover’s configuration.
  • The primary 80% versus 50% selection comparison uses 10 task-seed pairs, limiting statistical power for binary tests.The exact sign test reports p = 0.125, supplemented by continuous speedup analysis.
  • Surprisal-guided selection requires multiple correct samples and sufficient intra-task logprob variance, conditions that may fail on harder or low-coverage tasks.On 11/20 L1 tasks with logprob standard deviation below 1.0, all strategies perform equivalently.
  • The fast-proxy protocol uses 5 timing trials per kernel, so rankings could shift under the full 50-trial KernelBench protocol.Selected kernels showed consistent rankings on H100 hardware, but larger full-protocol replication is still needed.
  • The inverse confidence–quality relationship may be domain-specific and could perform poorly when the distribution mode represents optimal behavior.The paper calls for analysis linking training-distribution coverage to test-time solution quality.

8. Conclusion

Across KernelBench L1, search outperforms minimal test-time adaptation, while surprisal-guided selection improves how sampled solutions are chosen. The results attribute TTT’s weakness to over-sharpening that collapses diversity away from promising tail solutions.

  • 90% task success at K = 64 (18/20 tasks) for Best-of-N exceeds TTT’s 30.6% best-checkpoint success, with TTT equivalent K < 1.The comparison covers all 20 KernelBench L1 evaluation tasks and indicates adaptation underperforms single-sample inference.
  • Best-of-N performance saturates at K = 16 with 99.9% success on 5-task subsets, supporting investment in sample diversity over gradient updates.
  • 80% success for highest-surprisal correct samples versus 50% for most-confident selection yields a 30% improvement with zero additional compute.
  • Surprisal-guided-top3 matches oracle performance at 100%.
  • Gradient updates collapse the policy toward mediocre solutions, destroying diversity needed to find optimal kernels in the distribution tail.
  • For dense-reward VEG tasks, compute should be allocated to sample diversity and intelligent selection rather than gradient adaptation.The principle may generalize where rare, high-quality solutions occupy low-probability regions of the model distribution.
  • The assumption that test-time training provides universal benefits does not hold in this dense-reward, deterministic-evaluation regime.

A. Experimental Configuration

The experimental configuration specifies an RLVR-trained GPT-OSS-120B model with GRPO and LoRA, alongside recorded evaluation and TTT trajectory tables.

  • The RLVR training configuration uses GPT-OSS-120B, GRPO, rank-16 LoRA, learning rate 1e-5, batch size 8, and group size 8.Training uses 80 KernelBench L1 training tasks, temperature 0.25, maximum 1024 tokens, and normalized rewards.
  • The evaluation section identifies the RLVR final checkpoint at step 40.
  • Tables 5 and 6 document RLVR training progression and the batch TTT trajectory for seed 42 on Subset 1.

B. RLVR Training Progression

The training progression establishes a strong base policy, then shows that batch test-time adaptation peaks early before regressing while compute-matched comparisons account for rollout and backpropagation costs.

  • 98.4% correctness with 0.87x mean speedup at step 40 establishes a strong base policy for test-time evaluation.
  • 42.5% aggregate performance peaks at step 2 before regressing in the batch TTT trajectory for seed 42 on Subset 1.Task dynamics are heterogeneous: some tasks peak early and regress, while others saturate or improve later.
  • Compute accounting enables comparisons across methods, including rollout counts and backpropagation overhead.A full 9-step TTT trajectory takes approximately 267 minutes for 1,440 rollouts.
  • fast 1, requiring correctness and speedup above 1x, is the operative metric because raw speedup magnitude varies across tasks.
  • Table 8 reports Best-of-N selected speedup for K =64 on Subset 1 with seed 42.

F. SDPO Self-Distillation Experiments

The SDPO experiments test whether execution feedback improves self-distillation at frontier scale. In this KernelBench L1 setting, feedback provides no lift over prompt-only distillation, while TTT over-sharpening persists across learning rates.

  • SDPO conditions each student rollout on the task prompt, an available correct batch solution, structured execution feedback, and a correction instruction.The student’s original code is omitted; only its execution outcome is included in the teacher context.
  • SDPO with full execution feedback reaches 26.3%, below the 30.4% achieved by prompt-only SDPO across all three seeds.The result is consistent with the finding that execution feedback provides no lift over prompt-only self-distillation in this setting.
  • At 120B scale, prompt-only self-distillation succeeds, but explicit feedback provides no lift in dense-reward VEG domains.The authors note that smaller models or harder L2/L3 tasks might benefit more from structured feedback.
  • Over-sharpening persists across TTT learning rates spanning three orders of magnitude.
  • The experiments include surprisal quartile and isolated per-task TTT analyses alongside SDPO results.
Loading 2602.07670v1…