Strategic Scaling of Test-Time Compute: A Bandit Learning Approach
Authors: Bowen Zuo
Yinglun Zhu
Venue: ICLR 2026 Poster (2026)
Keywords: Test-time scaling
bandit learning
large language models
pure exploration
Primary area: foundation or frontier models, including LLMs

Abstract: Scaling test-time compute has emerged as an effective strategy for improving the performance of large language models. However, existing methods typically allocate compute uniformly across all queries, overlooking variation in query difficulty. To address this inefficiency, we formulate test-time compute allocation as a novel bandit learning problem and propose adaptive algorithms that estimate query difficulty on the fly and allocate compute accordingly. Compared to uniform allocation, our algorithms allocate more compute to challenging queries while maintaining accuracy on easier ones. Among challenging queries, our algorithms further learn to prioritize solvable instances, effectively reducing excessive computing on unsolvable queries. We theoretically prove that our algorithms achieve better compute efficiency than uniform allocation and empirically validate their effectiveness on math and code benchmarks. Specifically, our algorithms achieve up to an 11.10\% performance improvement (15.04\% relative) on the MATH-500 dataset, up to 10.82\% (14.44\% relative) on the AIME25 dataset, and up to an 11.23\% performance improvement (15.29\% relative) on the LiveCodeBench dataset.

Published as a conference paper at ICLR 2026
STRATEGIC SCALING OF TEST-TIME COMPUTE: A BANDIT LEARNING APPROACH
Bowen Zuo Yinglun Zhu' University of California, Riverside (bzuo002, yzhu)@uer.... edu
ABSTRACT
Scaling test-time compute has emerged as an effective strategy for improving the performance of large language models. However, existing methods typically allocate compute uniformly across all queries, overlooking variation in query dif- ficulty. To address this inefficiency, as a novel bandit learning problem and propose adaptive algorithms that estimate query difficulty on the fly and allocate compute accordingly. Compared uniform allocation, our algorithms allocate more compute to challenging queries while maintaining accuracy on easier ones. Among challenging queries, our algorithms further learn to prioritize solvable instances, effectively reducing excessive comput- ing on unsolvable queries. We theoretically prove that our algorithms achieve better compute efficiency than uniform allocation and empirically validate their effective- ness on math and code benchmarks. Specifically, our algorithms achieve up to an 11.10% performance improvement (15.00 .4.% relative) on the MATH-500 dataset, up to 10.82% (14.44% relative) on the AIME25 dataset, and up to an 11.23% performance improvement (15.29% relative) on the LiveCodeBench dataset.
formulate test-time compute allocation
INTRODUCTION
Recent advances large language models (LLMs) have shifted attention from training-time compute (Kaplan et al., 2020: Hoffmann et al., 2022: Chowdhery et al., 2022) to test-time compute (Wei et al., 2023; Yao et al., 2023; Madaan et al., 2023; Agarwal et al., 2024; Muennighoff et al., 2025) as means of improving model performance. Test-time scaling methods such Best-of-N sampling (Brown et al., 2024; Snell al., 2024) and consistency checking (Wang et al., 2022) enhance output quality by generating multiple responses and selecting the most promising one. This selection process can be strengthened using high-quality reward oracles (Cobbe et al., 2021; Uesato et al., 2022: Lightman et al., 2023; Zhang et al., 2025a). These methods have achieved strong empirical gains without additional model training. For instance, as noted in OpenAl's ol release report (OpenAI, 2024), repeated sampling with 64 generations improves accuracy the 2024 AIME competition math dataset from 74.4% to 83.3% nearly 9% gain without any model updates.
Despite recent advances, most test-time scaling techniques still allocate compute uniformly across all queries (Brown et al., 2024; Snell et al., 2024), ignoring the inherent variability in query difficulty. This one-size-fits-al strategy inefficient: simple arithmetic questions receive the same compute as multi-step reasoning tasks, leading to wasted resources on easy queries and insufficient budget on hard ones. Ideally. one should allocate just enough compute to confidently solve easy queries and reallocate the remaining budget to harder ones. While recent work has begun exploring adaptive test-time strategies, existing methods either (1) focus on compute allocation within single query (Sun et al., 2024; Manvi et al., 2024; Tan et al., 2025), or (2) rely on two-stage procedures (Damani et al., 2024; Wang et al., 2025b) that require training an auxiliary model (or pre-compute allocation) in the first stage to guide later allocation decisions.
In this work, we introduce new serecciive--trattgee scaling of test-time compute-in which compute is adaptively allocated across set of queries based on their estimated difficulty. We formulate test time compute allocation as a fully adaptive pure-exploration-styl bandit learning
Project lead and corresponding author.

---

Published as a conference paper at ICLR 2026
0.73
660
E50
0.5
Elecery gan 268
dend 4.00
Oldeggoo +4000
675
Average Dorraste Budget
Avtrego Derrasse Budget
Antrago Dembute Buoget
Ancrago Dempute Buoget
Figure 1: Comparison between our algorithm and baselines. First: Accuracy comparison on MATH-500 with Llama-3. 3.10-3.1888 1-118-Innttrcct Second: Accuracy comparison on MATH-500 with Gemini-2, S-fllas--itte Third: Accuracy comparison on AIME25 with Qwen3-4B. Fourth: Cov- erage comparison on LiveCodeBench with DeepSeek-R1-Distinee -amma--aaam--8
problem (Bubeck al., 2009; Jamieson & Nowak, 2014; Locatelli et al., 2016; Zhu et al., 2020), treating each query as an action and allocating compute sequentially to maximize the number of queries answered correctly within fixed budget. Our adaptive algorithms estimate query difficulty the fly and prioritize compute for those most likely to henefit from additional inference. Empirically, method achieves up 11.10% absolute (15.04% relative) improvement the MATH-500 dataset (Lightman et al., 2023; Hendrycks et al., 2021), 10.82% absolute (14.44% relative) on AIME25 (AIME, 2025), and 11.23% absolute (15.29% relative) on LiveCodeBench (Jain al., 2024) all under the same compute budget as baselines. Fig. provides a high-level comparison between our method and various baselines.
Contributions. We summarize our main contributions below:
1. We formulate LLM test-time compute allocation as a novel bandit learning problem, bridging test-time scaling and bandit learning communities. This formulation grounds strategic test-time scaling in precise decision-theoree framework.
We propose a general algorithmic framework for strategic compute allocation, supporting flexible exploration strategies-includdin a novel entropy-basee rule. Our framework naturally extends to incorporate alternative aggregation methods and handles both streaming and token-constrained settings. We further provide theoretical insights into the efficiency gains of our adaptive approach.
We conduct extensive experiments on math and code benchmarks and show that our algorithms consistently outperform baselines. Further analyses demonstrate that our algorithms adaptively allocate compute to harder queries in standard settings, and to solvable queries in scenarios containing both solvable and unsolvable instances, effectively avoiding compute waste.
Paper organization. We introduce the problem setting in Section 2. Our solution is presented in Section 3, including the bandit formulation, algorithmic framework, extensions, and theoretical analysis. Empirical results are in Section 4, covering main results, analyses, and ablations. We conclude Section 5. Due to space limitations, we defer discussion related work, additional proofs, and extra experimental details and results the Appendix.
2 PROBLEM SETTING
Let p denote language model, which takes a query E X' input and generates a response y ~ p(- æ). Recent studies show that scaling up the test-time compute can significantly improve the performance of LLMs across a variety of tasks (Snell al., 2024). In this context, we consider the amount of test-time compute as the total number of responses generated by the language model. For example, given query x E X' and compute budget of N, the model can generate set N responses g(x; N) == {11, distribution p(- x). A reward oracle r x y - [0, 1] is used to evaluate the quality of each generation; the reward oracle can be instantiated by either ground truth verifier or a learned reward model (Cobbe et al., 2021; Uesato et al., 2022; Lightman et al., 2023; Zhang et al., 2025a). When the evaluation metric requires single response as the output, test- time compute methods such as the ofst-of-N algorithm (Brown et al., 2024) use the reward oracle to score each response and return the with the highest score. Specifically, given of responses g(x; N) (V1,"* ... ,yN) and
). where each response Yi ~ p(- .:) is sampled from the conditional
2

---

Published as a conference paper at ICLR 2026
letting r(z, Ue) denote the score of response yes the final output f(; N). == fg((=; is defined as: faz N) ; where Y" == arg max (x,
VE[N]
While scaling test-time compute improve performance, existing methods primarily focus uniform allocation of compute budget. Specifically, given a set of queries = {I1, .. Fn and total compute budget B := nB, existing approaches assign the same compute budget to each query 2; and generate the final outputs {211 f (21; B)),... (In, f ())): This uniform allocation is inefficient: it ignores differences in query difficulty and assigns the same compute to both easy and hard queries.
2.1 STRATEGIC TEST-TIME COMPUTE ALLOCATION
To address the limitations of uniform allocation, we study the problem of strategic test-time compute allocation-how adaptively allocate a total compute budget across of queries maximize the fraction of correctly answered queries. Let B denote the total compute budget and S = { Fn) be a of queries. Let Metric [0, 1] be evaluation metric and c(ar,) be the compute allocated to query The goal is to maximize the overall performance subject to a budget constraint:
max (o(m))))
(etric(a,,e(a;))) subject
Σ c( S.
(I)
We consider two popular evaluation metrics: coverage and accuracy. Given compute allocation c(z;), coverage evaluates whether any of the c(x;)) generations in g(z;; c(x;)) correctly answers query æ;, while accuracy evaluates whether the final output f(x;; c(r;)) is correct. These metrics defined as:
(overage c(xx)) [l[hee exists y g( c(z:)) that correctly answers query xi.} Accuracy(I; c(x:)) c((za)
== (x.; correctly answers query 2;-}
The key challenge in Eq. (1) is adaptively allocate compute budget c(x,) to each query æ; under uncertainty- that is, without knowing in advance the difficulty of each query or how much compute is needed answer it correctly. To isolate and address this challenge, we adopt the standard Best-of-N approach (Brown et al., 2024; Snell et al., 2024) for both compute counting (i.e., measuring the number generations per query) and final output selection.
METHODS
We present our approaches to solve the strategic test-time compute allocation problem introduced in Section 2.1. In Section 3.1, 1, we first formulate test-time compute allocation as a bandit learning problem. We then introduce algorithmic framework in Section 3.2, followed by extensions in Section 3.3 and theoretical analysis of compute efficiency in Section 3.4.
3.1 TEST-TIME SCALING AS BANDIT LEARNING
To address the challenge of strategic compute allocation under uncertainty, introduce novel bandit learning formulation tailored to LLM test-time compute objectives. Following the bandit terminology, we treat each query & E S an action, and interpret sampling action allocating one unit of compute to query x to obtain a randomly generated response y. After taking action v, the learner receives feedback from reward oracle the form of score r(z, y).
Our objective is to design an adaptive compute allocation algorithm that maximizes the fraction of queries that are correctly answered within fixed compute budget B. Assuming availability of sufficiently accurate reward oracle (e.g., ground truth labels), we approximate the correctness of response using a user-spe threshold y E [o, 1]: a response y to query x: is considered correct r(x,2) v) 7. Formally, the algorithm adaptively distributes the total compute budget B across all
We assume access to sufficiently accurate reward oracle in order to focus on the key challenge of adaptive compute allocation. This assumption is clearly satisfied in settings with ground truth labels, and is approximately satisfied by recently developed pracess reward models (Zhang al., 2025a).
3

---

Published as a conference paper at ICLR 2026
Algorithm 1 Strategic Test-Time Compute Allocation
Input: Query set S, total compute budget B, reward oracle r, per-round per-query compute budget
K, and elimination threshold y,
1: For each query E S, maintain response set ((r), the best-scoring response j(x),, and its associated reward f(z).
2: Initialize the active set A <- S to be the full query set.
3: while > 0 and LA| > 0 do
4: for x E A do
5:
- B - K. /1 The exploration rule: allocating compute all queries in the active A. We discuss extensions of the exploration rule in Section 3.3.
Generate K new responses {u: h Update g(x) + g(x) U fy: and B
6:
Get 1* f- arg (aX r(2.8)).
7:
rx, Ve-) F(x:) then
8:
Update j(z)
yı- and f(x: + -xx,00)).
9:
rf y/-) ) 7 then
// The elimination rule. Output: For each x S. output its response set g(x) and the best-scoring response ÿ(z). // Use
10:
Update A f A A {x).
g(x) for coverage evaluation and j(x) for accuracy evaluation.
queries through an allocation (c(x:))== optimizing the following objective:
max r(z;,y) u) 2 7),
max (o(2)))
where g( c(x;)) denotes the set of c(x;) responses generated for query x,,?
While our formulation is conceptually related to the bandit pure exploration problem (Bubeck et al., 2009; Jamieson & Nowak, 2014) and its thresholding bandit variants (Locatelli et al., 2016; Zhu et al., 2020), it fundamentally departs from the conventional objectives. Standard pure exploration settings aim to identify actions (queries) with high expected scores, which correspond-in our setting-to identifying subset of easy queries that can be reliably answered by the LLM. In contrast, our objective aims at generating at least one high- nigh-quality (correct) response for each query, regardless of its expected score. To our knowledge, this not only introduces novel bandit formulation but also opens the door to further exploration of bandit-based LLM test-time compute allocation.
3.2 OUR ALGORITHMIC FRAMEWORK
Based on the bandit formulation, we next present our algorithmic framework in Algorithm Given query set S, Algorithm initializes active set A S that contains active queries that have not yet been confidently answered. For each query x E S, it maintains response set g(x), the best-scoring response ixx) observed far, and its corresponding reward score f(x), evaluated by the reward oracle r. Algorithm proceeds in rounds, and operates based on two key components: an exploration rude and elimination rule:
The exploration rule. At each round, Algorithm I explores all queries in the active set, i.c., for each active query x E A, it generates new responses { and updates the response set g(x) + g(x) {b ===1 We discuss extensions to this simple exploration rule in Section 3.3.
The elimination rule For each explored query let yı- denote the response that achieves the highest score among newly generated responses, i.e., 1 = arg maxxg[K] r(x, ys). the reward r(z, ye- ) is greater than the previously observed best score f(z), then Algorithm 1 (1) updates its maintained best-scoring response )(x) yı- and the corresponding reward f(z) = r r(z, y and (2) eliminates query x from the active set if the score r(x, ye-) is also greater or equal to the elimination threshold 7.
Al Algorihhm terminates when the compute budget exhausted (i.e., B = 0) or when all queries have been eliminated from the active set (i.e., A 0). For each query x Algorithm outputs
When the evaluation metric is Accuracy, one must further explicitly select and output the correct response.
4

---

Published as a conference paper at ICLR 2026
its maintained response set g(x) for coverage evaluation, and its best-scoring response (xx) for accuracy evaluation.
Reward oracles. Reward oracles have become core component in test-time compute techniques, even for the vanilla uniform Best-of-N algorithm (Brown et al., 2024; Snell et al., 2024). Common reward oracles include outcome reward models (ORMs, Cobbe et al. (2021), 1)) and process reward models (PRMs, Uesato et al. (2022); Lightman et al. (2023); Zhang et al. (2025a)). For tasks with easy or automatic verification, such as math and code generation, ground truth (GT) labels can serve exact reward oracle. We emphasize that Algorithm uses the same number of reward oracle calls as the uniform Best-of-N algorithm, which relies on the reward oracle to select the final output.
Hyperparameters. Algorithm 1 takes two hyperparameters as input: the per-round per-query compute budget K and user-specified elimination threshold 7, The hyperparameter per-round per-query compute budget controls the granularity level of the budget allocation: a smaller value of K leads to more fine-grained budget allocation with an increased number of allocation rounds. The climination hyperparameter y decides when to eliminate a query from the active set A. The value of be determined based expert knowledge or based on cross-vas vosssvvaiaationn on separate training set. These hyperparameters offer additional levels of flexibility for Algorithm 1. We conduct ablation studies of these hyperparameters in Section 4.4 and Appendix D.2.1.
3.3 EXTENSIONS OF ALGORITIIM
Algorithm 1 with different aggregation strategies. While main discussion centers Best- of-N, the proposed framework is flexible and can accommodate alternative aggregation strategies. Prior work (Wang al., 2025a) has shown that Self-Consistengy (SC) is often more effective for reasoning models-succ as Qwen3-4B-44e
to their tendeney to produce logically coherent outputs. To incorporate SC into Algorithm 1, we make two modifications: (I) the selection rule (line 8) now uses SC instead of a reward model (PRM), and (2) the climination rule (line 10) is updated to eliminate a query once certain proportion of its collected responses converge to the same answer (e-g., when over 50% agree). When using SC, the reliance on PRMs can be eliminated altogether. Experiments in Section 4.2 confirm that our algorithm remains effective when using SC as the aggregation rule.
Algorithm 1 with different exploration rules. While the base version of Algorithm I (ELIMINATION) explores all active queries uniformly at each round, our framework supports more targeted exploration strategies inspired by the pure exploration bandit literature. For example, Upper Confidence Bound (UCB) prioritizes queries with high empirical reward plus an uncertainty bonus (Kalyanakrishnan et al., 2012; Jamieson et al.. 2014), while gap-based sampling (GAP) focuses queries near the elimination threshold y, allocating compute inversely proportional to the estimated reward gap (Locatelli et al., 2016). We also propose a novel entropy- rule (ENTrOPY) that selects queries with more diverse response patterns, as measured by empirical entropy, and encourages exploration of under-sampled queries. Experiments in Section 4.3 show that ENTROPY is particularly effective across extremely difficult query sets. We defer formulations of these strategies to Appendix B. For all the exploration rules, to prevent over-allocation of compute on difficult queries, additionally introduce per-query cap max_sam sample E N, which limits the number of generated responses for any individual query.
Algorithm 1 with fine-grained token controls. The default version of Algorithm models compute cost as the number of response generations. However, it can be easily extended to track and control token-level usage. At each iteration, the algorithm can record token consumption and stop once the total token budget is reached. Alternatively, one impose fine-grained token caps per generation. We evaluate this variant in Section 4.4 and find that Algorithm continues to outperform baselines under the same token budget.
Algorithm 1 with streaming queries. In streaming settings, queries arrive sequentially, i.e., only the current query 2; is accessible at round To adapt Algorithm to this setting, we modify line to focus solely 2; while keeping the rest of the framework unchanged. Also, to prevent over- allocation of compute on difficult queries, we additionally introduce a per-query cap max sample e N, which limits the number of generated responses for any individual query. This constraint enforces
5

---

Published as a conference paper at ICLR 2026
local trade-off between exploration and exploitation and promotes balanced compute usage across the query stream. Ablation results in Section 4.4 show that this streaming variant remains competitive with our original method.
3.4 THEORETICAL INSIGHTS ON COMPUTE EFFICIENCY
A key strength of Algorithm 1 (and its variants in Section 3.3 with different exploration rules) lies in their ability adapt compute based on estimated query difficulty: easier queries get fewer samples, while harder ones are given more when needed.
To better understand the advantage of Algorithm over uniform compute allocation, we consider the following probabilistic model. For each query E S, we model the correctness of the LLM's response
single, independent generation as a Bernoulli random variable with parameter A₂ E (0, 1). That is, X Bernoulli( A.), where X if the LLM answers the query correctly and X 0 otherwise.
To ensure the reward oracle is compatible with this probability model and the specified threshold 7, we make the following assumption.
Assumption For any query x E and any randomly generated response y. y correctly answers x and only if the reward oracle r assigns score r(z, y) 2 %
Assumption ensures that the climination decision is aligned with the reward oracle and the threshold, allowing us to focus the analysis of the adaptive design in Algorithm 1. This assumption is satisfied by the ground truth reward oracle, and holds approximately when the reward model is sufficiently accurate- condition empirically supported by recent advances in high-quality process reward models (Zhang et al., 2025a).
Suppose K = O1), we derive the following quantitative comparison between Algorithm 1 and uniform compute allocation.
Theorem 1. Assume Assumption and fix any 5 10, To output correct responses for all queries S with probability at least Algorithm requires a total budgel Baurs O(Eres L.), and this matches the information-
lower bound up to logarithmic factors. contrast, uniform factorrs
allocation strategy requires budget Bunif = f(maxzes
(S) A to achieve the same guarantee.
Theorem I highlights the clear efficiency advantage Algorithm over uniform allocation: the budget Bours required by Algorithm 1 is always smaller than that of uniform allocation Bunif and in certain cases, Bunif can be nearly |S| times larger than Bours- To give a concrete example, suppose |S] n and A, = i/n. We then have Bours (m), yet Bunif = s(n²), which is n times larger than Bours. This theoretical separation underscores the benefit of strategic compute allocation: by adapting to query difficulty, Algorithm 1 can dramatically reduce overall compute usage.
One may further consider the classical two-stage explore-ther www-tteenncmmiie (ETC) strategy, which allocates fixed exploration budget to each query before committing additional compute based on the observed outcomes (Lattimore & Szepesvári, 2020). Despite its partial adaptivity, ETC remains fundamentally limited because its first stage performs uniform allocation. We next show that this structural constraint already leads to strictly larger budget complexity in heterogeneous regimes.
Proposition 1. Assume Assumption 1, and fix any 8 (0, 1). Consider any two-stage explore-then- commit algorithm that allocates a fixed exploration budget m to every query in the first stage. To output correct responses for all queries in S with probability at least 1 - 6, any such algorithm must incur total compute BETC Si(Ees max(m, 3.)
In particular, when m is nontrivial and many queries are easy (i.e., when x is small), ETC incurs an unavoidable mifoorm-exploration overhead of ((m) per easy query, whereas Algorithm I only spends ((1/A.). This establishes strict separation between ETC and our fully adaptive allocation. Consistent with this theory, our empirical results Section 4 also show that the ETC baseline (denoted as Two STAGE) substantially outperformed by our algorithm.

---

Published as a conference paper at ICLR 2026
3.55
0.30
910
1993
Oholecy 240
0.2
64
0m
333
PM (oua)
000-04 Dudn
Larv5 50.
(a) Llama-3.. 3.2-1B
(b) Llama-3. -32-1B
(c) Llama-3 1amma-3..11 3-BB
(d) Gemini-2.5
Figure 2: Results on MATH-500 with different models.
Accuracy
results of
Llama-3 3. In-nrstttct and Gemini-2. 2 5--5ashh-lie
are presented in Fig. 1.
EXPERIMENTS
We describe experimental setups in Section 4.1, present main results Section 4.2, offer further analysis in Section 4.3, and report ablations in Section 4.4. Additional experimental details and results are deferred to Appendix D.
4.1 EXPERIMENTAL. SETUP
Datasets. We examine the performance of our algorithms on standard math and code benchmarks: MATH-500 and AIME25 (Lightman et al., 2023; Hendrycks et al., 2021; AIME, 2025) and Live- CodeBench (Jain et al., 2024). MATH-500 contains 500 math questions, AIME25 contains 30 difficult math questions, and the LiveCodeBench contains 479 code execution questions that were collected from 5/1/2023 to 12/1/2023, From MATH-500, we further construct one challenging subset: MATH- 500-Hard-166 16, which contain questions that cannot be correctly answered after allocating 16 units of compute. Intuitively, this subset consists of the most difficult queries in the MATH-500 dataset.
Baselines. We compare our algorithms with the uniform Best-of-N baseline (Brown et al., 2024), referred to as UNIFORM, and two-stage baseline, referred to as TWo STAGE (Damani et al., 2024; Wang al., 2025b). The TWO STAGE baseline first uniformly allocates compute to estimate problem difficulty (stage 1) and then allocates the remaining compute proportionally (stage 2); in other words, operates an explore-then-commi style. For the Two STAGE algorithm, we vary the stage compute ratio from (25%, 50%, 75%) } and report the best results. We report the performance of our Algorithm (ELIMINATION) and its variants introduced in Section 3.3.
Models and metrics. We conduct experiments with commonly used LLMs of various sizes, in- cluding 3
and Llama-3. Instrnst (Grattafiori et al., 2024), as well as more recently developed reasoning models DeepSee--R1--istil--- 2025), Qwen3-4B (Yang et al., 2025), and Gemini-2.55 1tw.ii2.55-laasseeee 5--lass--Li (Comanici et al., 2025). For MATH-500, we use different reward oracles depending on the base model family: for Llama models, we use the PRM Qwen2. 5-Math-PRM-78 (Zhang al., 2025a); for Gemini-2 2.5.5lash--iie an LLM-aa-a-judge oracle based on Gemini-3, wwweiiniin...ff..... ground-truth (GT) oracle when available." For AIME25, of baselines and our methods, as recent work shows that SC is more effective for reasoning models (Wang al., 2025a). For Bench, since correctness can be deterministically verified by code execution, we use the GT reward oracle and report only the coverage metric (equivalent to accuracy). We conduct experiments under average compute budgets of (4, 8, 16,32) and report results averaged over 4 random runs, with shaded regions in plots representing +0.5 standard deviations.
(DeepSeek-A) et al.,
use In addition, we also report results under the
use the Self- Self-Consistency (SC) variants
4.2
MAIN RESULTS
MATH-500 results. Fig. 2 presents experimental results on the MATH-500 dataset across two LLMs and evaluation metrics (the accuracy result of Llama-3. --B-Insttrct MATH-500 is pre- sented in the first plot of Fig. 1). Across all configurations, all variants of our Algorithm consistently
When using the GT reward oracle, accuracy and coverage equivalent. In this case, only report coverage.
7

---

Published as a conference paper at ICLR 2026
0.90
0.70
0.7
0.6
0.80
Coveregg
0.80
Covernge
Effciency galn 1.8 131.0-
0.55
0.75
Effiiency gan - 4.00 (8 22.00)
0.50
enrry pain
Unftmm SC
0.4
0.70
5m joum)
Elm (Dars)
E 3m--TT0 (Durs)
Lin For--5T
Unifom-57 Eln- EI (Dans)
16 Avorage Compute Budget
32
16
32
16 Average Compute Budget
32
Averag o Compute Budget
(a) Qwen3-48
(b) Llama- 2.B--...s-1----ruc
(c) Llama- 1-88 Instruct
Figure 3: Results on AIME25 (left) and LiveCodeBench (middle and right). Accuracy re- sult of Qwen3-4B on AIME25 presented at the middle plot Fig. I; coverage result of -itbseek-RR-Dittll---
on LiveCodeBench is presented at right plot of Fig.
outperform both baselines. Under the accuracy metric, when the average compute budget 16, our method achieves 2.50% absolute improvement (7.37% relative) on Llama-3. 3.ataaa...11------ 2-1B-Instruct; this corresponds to 1.82x efficiency gain as shown on plot (a): UNIFORM takes 1.82x compute to achieve the same performance. For Llama-3. 1-8B-Inse Instruct,
observe a 1.40% absolute improve- ment (4. 11% 1% relative), with a 2x efficiency gain (Fig. 1, first plot). For the coverage metric, when the average compute budget 16, method yields 10.70% absolute improvement (17.95% relative) on Llama-3. #rB-Innstrct resulting in a 2 efficiency gain (Fig. 2, plot (b)).* When the average budget is we observe an 11.10% absolute gain (15.04% relative) on Llama-3 1-8B-Instruct, yielding x.93x efficiency gain (Fig. 2, plot (c)).
We further evaluate our method on the frontier model Gemini 5-Flash-Lite with reasoning enabled. Given its strong base performance, we use smaller average compute budgets of 2,, 4]. Because existing PRMs are not reliable at this performance level, we instead use Gemini-3 0-Flssh as an LLM L-a-aa--ajjudee oracle for demini-21 11..2..faass--itttt. Under the accuracy metric, our method improves performance by 3.4% absolute (3.66% relative), corresponding gain (Fig. 1, second plot). Under the coverage metric, obtain 2.4% absolute improvement (2.64% relative), again exceeding a 2x efficiency gain (Fig. 2. plot (d)).
> 2.00 efficiency
AIME25 results. Fig. 3 reports results on the AIME25 dataset with Qwen3-4B (accuracy result of Qwen3-4B on AIME25 is presented at the third plot of Fig. 1). Following Yang et al. (2025), we enable the reasoning ability of Qwen3-4B and set the max token length 38,912. Under the accuracy metric (Fig. 1, third plot), we are able to gain an absolute of 3.33% performance gain (4.60% relative) when the average compute budget is 8 on Qwen3-45 yielding 4.00x efficiency gain. For the coverage metric (Fig. 3, left plot), when the average compute budget is 16, we observe 10.82% absolute gain (14.44% relative) on Qwen3-4B, yielding a 4.00x efficiency gain.
LiveCodeBench results. Fig. 3 (middle and right) presents results on the LiveCodeBench with Llama models of different sizes, and the fourth plot of Fig. presents results of Depssek--RI-Dstill--L
-B. As described in Section 4.1, we use the GT reward oracle and re- port coverage, which is equivalent accuracy in this setting. We report results for the ELIMINATION variant only, as UCB and GAP behave identically to EL.MINATION under the GT oracle. Across all compute budgets, our method consistently outperforms uniform allocation. With an average compute budget of 16, Llama-3. 2-12 2-1B-Instrucc achieves 6.47% absolute improvement (11.63% relative), corresponding to 1.98x efficiency gain (Fig. 3, middle plot). With an average compute budget of 8, Llama-3. 1 -ns-Inst achieves 7.41% absolute improvement (14.40% relative), corresponding to a 3.11x efficiency gain (Fig. 3, right plot). With an average compute budget of achieves 9.97% absolute improvement (12.30% relative),
8, -iiieek------sssl--llaaaal- corresponding to a 4.00x efficiency gain (Fig. 1, fourth plot).
MATH-500-Hard result. Fig. 4 (left) presents experimental results the MATH-500-Han 16 dataset, which was constructed to include the most challenging questions the MATH-500 benchmark. We evaluate performance using the GT reward oracle, as PRM-based scores are less
Under the GT reward oracle, report only the performance of FLIMINATION, other variants yield similar results. See Appendix D.2.1 for full comparisons.
8

---

Published as a conference paper at ICLR 2026
100
100
Uniften Tan-Stage Ellm UCB Entropy
Average Compute Burgee
PRM
05.14
80
Quastions 32
60-
60
50.56
2.
40
Solvbbic
6.53
Unsetabble
16
Average Compute Budget
Easy Group
Hard Group
Averago Computs Budget
Figure 4: MATH-500- Hard result and analyses with Llama-3. 3. 1..-1B-IIstruue at an average compute 32. Left: Coverage MATH-500-HAREH 16. Center: Allocation behavior of Algorithm for easy vs. hard groups. Right: Allocation behavior of ENTROPY for solvable vs. unsolvable groups.
reliable these difficult questions. On this dataset, baseline methods and the vanilla version of our Algorithm (BLIMINATION) doesn't performance well, as many questions can not be solved by the base LLM. However, we find alternative exploration rules introduced in Section 3.3 p.r-puticullarly EITROPY-achies significantly better results. These findings highlight the benefits of incorporating more nuanced exploration strategies, such as those developed in Section 3.3, for effective compute allocation on challenging benchmarks. We defer more experiment results to Appendix D.2.2.
4.3 ANALYSIS ON THE ADVANTAGES OF STRATEGIC COMPUTE ALLOCATION
We conduct further empirical analyses to illustrate the benefits strategic compute allocation in two settings: (1) on standard datasets containing both easy and hard queries, and (2) on challenging datasets containing both solvable and unsolvable queries. All experiments are conducted using Llama-3. IIInttrrct with average compute budget of 32.
Strategic allocation on standard datasets. In the first analysis, we partition the MATH-500 dataset into two subsets: queries that can be correctly answered with most 32 units of compute (easy group), and those that cannot (hard group). Intuitively, the easy group consists of questions that require less than 32 units compute to solve, while the hard group includes questions that would benefit from additional compute. In the middle plot of Fig. 4, we visualize the compute allocation of Algorithm 1 under both PRM and GT reward oracles. Compared to uniform allocation, our algorithm allocates fewer resources to easy queries and more hard ones. This demonstrates the ability of Algorithm to strategically allocate compute-
effort for harder queries that need most.
Strategic allocation challenging datasets. In the second analysis, we consider the MATH-500- Hard-16 16 dataset and divide it into solvable queries and unsolvable ones, where the latter cannot be correctly answered even after allocating 500 units of compute. In such settings, effective allocation should prioritize the solvable subset, as investing in unsolvable queries leads to wasted compute. The right plot of Fig. 4 shows that under 32-unit compute budget, ENTROPY allocates on average to solvable queries, and larger fraction of them receive more than 32 samples. This demonstrates that method learns concentrate compute on tractable instances, avoiding waste on queries unlikely to be resolved.
compute
To understand why ENTROPY behaves this way, inspect model outputs these challenging questions in detail. We observe that unsolvable queries often yield invalid responses (e.g., incomplete or poorly formatted), leading lower entropy across generations. In contrast, solvable queries tend to produce more diverse and well-formed outputs, resulting in higher entropy. We defer detailed experiments and analysis of ENTROPY extension Appendix D.2.2.
4.4 ABLATION STUDIES
Algorithm with different K. All main experiments in Section 4.2 use the default setting K Smaller values of K enable finer- grainer adaptive allocation and are generally preferred for max- imizing performance. In Fig. 5 (left), we conduct ablation study with K E {1, 2, 4, the LiveCodeBench dataset using Llama-3, 2-1B-Innttrun Across values of K, our method consis- tently outperforms the uniform baseline. While larger K reduces allocation granularity-mal

---

Published as a conference paper at ICLR 2026
0.7
0.85
Coverage
0.5
Unter-GT
ElmET (0-1)
timE1 (0-2)
unform-GT
Unitsm-0
Elm-TT (Kx6)
tim (Durs)
mim-P (Dura)
0.3
16
Elmee 00-10
0.70
Eim-St (Durs Taken Maichell
16
16
Bim-Ph (Dura. ctrmnnngg
Average Compute Budget
Average a Compute Budget
Average Compute Budget
Figure 5: Ablation studies for Algorithm 1.
Left: Effects of K on LiveCodeBench Token control on LiveCodeBench with
with 3. 2-18-Instruct. Center: -iireeek------sssll-------- Llama-3. 2-1B- 3.2-B- 1-11-IInstrut..
Right:
Streaming variant MATH-500 with
performance closer uniform allocation under tight budgets the gap narrows the average compute budget increases. These results show that Algorithm is robust to the choice of K.
Algorithm with token controls. We evaluate Algorithm in token-controlled setting, where compute is measured by total token usage rather than the number of generations. To ensure comparability, we match the average token budget used by uniform allocation and discard ex- cess samples when needed. As shown in Fig. 5 (middle), on the LiveCodeBench dataset with -nn-See--RR-Distll---- equivalent token budgets.
Algorithm still outperforms uniform baseline even under
Algorithm with streaming queries. We also test Algorithm a streaming setting, where queries arrive sequentially and the full query pool is not available in advance. Using the variant described in Section 3.3, we evaluate performance in this setting and report results Fig. 5 (right). Our method performs comparably to its pool-based counterpart and still outperforms UNIFORM with pool access, demonstrating its effectiveness under streaming constraints.
5 CONCLUSION
We introduce new perspective on LLM test-time scaling by formulating strategic compute allocation bandit learning problem. We develop adaptive algorithms that estimate query difficulty on the fly and allocate compute to maximize the fraction of correctly answered queries under a fixed compute budget. Our framework is flexible and extends naturally incorporate alternative aggregation and exploration strategies, as well as to support both streaming and coken-constraing settings. We provide theoretical guarantees that strategic compute allocation improves compute efficiency over uniform allocation, and we empirically demonstrate substantial performance improvements up to 11.10% on MATH-500, 10.82% on AIME25, and 11.23% on LiveCodeBench. These findings underscore the potential of bandit-based compute allocation for more effective test-time scaling.
REPRODUCIBILITY STATEMENT
We present algorithmic framework and our ideas for choosing hyperparameters in Section 3. We also discussed algorithmic extensions formulation in Appendix We list the datasets, models, and our baseline methods in Section 4. We report extra hyperparameters, our experiment setting, and also our prompts used for all experiments in Appendix D.
10

---

Published as a conference paper at ICLR 2026
REFERENCES
Rishabh Agarwal, Avi Singh, Lei M. Zhang, Bernd Bohnet, Luis Rosias, Stephanie Chan, Biao Zhang. Ankesh Anand, Zaheer Abbas, Azade Nova, John D. Co-Reyes, Eric Chu, Feryal Behbahani, Aleksandra Faust, and Hugo Larochelle. Many-shot in-context learning, 2024. URL https: / //arxiv.org/abs//444.. 11018.
Shipra Agrawal and Navin Goyal. Analysis of thompson sampling for the multi armted bandit problem. In Conference on learning theory, pp. 39-1. JMLR Workshop and Conference Proceedings, 2012.
AIME Aime problems and solutions, 2025. URL httpssss/arrrofrroooloooolvieeeo com/wiki/ wwtp:pph/AAMMMPbbbba/.. owr.oobeeeeenansoooooooo.o..
Jean- Yves Audibert and Sébastien Bubeck. Minimax policies for adversarial and stochastic bandits.
In COLT, pp. 217-226, 2009.
Peter Auer, Nicolo Cesa-Bianchi, and Paul Fischer. Finite-time analysis of the multiarmed bandit problem. Machine learning, 47:235-256, 2002.
Amanda Bertsch, Maor Ivgi, Uri Alon, Jonathan Berant, Matthew R. Gormley, and Graham Neubig. In-context learning with long-context models: An in-depth exploration, 2024. URL https:
//arxiv.org/abs/220000 00200.
Bradley Brown, Jordan Juravsky, Ryan Ehrlich, Ronald Clark, Quoc V. Le, Christopher Ré, and Azalia Mirhoseini. Large language monkeys: Scaling inference compute with repeated sampling, 2024. URL https:///axivv wwwpsiiwwwg///b.eeeeeeeeee arrg/abb/24077 21787.
Tom B. Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, Sandhini Agarwal, Ariel Herbert- Gretchen Krueger, Tom Henighan, Rewon Child, Aditya Ramesh, Daniel M. Ziegler, Jeffrey Wu, Clemens Winter, Christopher Hesse, Mark Chen, Eric Sigler, Mateusz Litwin, Scott Gray, Benjamin Chess, Jack Clark, Christopher Berner, Sam McCandlish, Alec Radford, Ilya Sutskever, and Dario Amodei. Language models are few-shot learners, 2020. URL https: //arxiv. .rgg /a//aabs///2/// 14165.
Sébastien Bubeck, Rémi Munos, and Gilles Stoltz. Pure exploration in multi-armed bandits problems. In International conference on Algorithmic learning theory, pp. 23-37. Springer, 2009.
Sébastien Bubeck and Nicolò Cesa-Bianchi. Regret analysis of stochastic and nonstochastic multi- armed bandit problems, 2012. URL https: / orgg hargg/abs/12/ 11111.4.7222 5721.
Olivier Chapelle and Lihong Li. An empirical evaluation of thompson sampling. Advances in neural information processing systems, 24, 2011.
Dingyang Chen, Qi Zhang. and Yinglun Zhu. Efficient sequential decision making with large
language models. Empirical Methods in Natural Language Processing, 2024.
Xinyun Chen, Maxwell Lin, Nathanael Schärli, and Denny Zhou. Teaching large language models to self-debug. arXiv preprint arXivv2330..05/22 2023.
Aakanksha Chowdhery, Sharan Narang, Jacob Devlin, Maarten Bosma, Gaurav Mishra, Adam Roberts, Paul Barham, Hyung Won Chung, Charles Sutton, Sebastian Gehrmann, Parker Schuh, Kensen Shi, Sasha Tsvyashchenko. Joshua Maynez, Abhishek Rao, Parker Barnes, Yi Tay. Noam Shazeer, Vinodkumar Prabhakaran, Emily Reif, Nan Du Ben Hutchinson, Reiner Pope, James Bradbury, Jacob Austin, Michael Isard, Guy Gur-Ari, Pengcheng Yin, Toju Duke, Anselm Lev- skaya, Sanjay Ghemawat, Sunipa Dev, Henryk Michalewski, Xavier Garcia, Vedant Misra, Kevin Robinson, Liam Fedus, Denny Zhou, Daphne Ippolito. David Luan, Hyeontaek Lim, Barret Zoph, Alexander Spiridonov, Ryan Sepassi, David Dohan, Shivani Agrawal, Mark Omernick, Andrew M. Dai, Thanumalayan Sankaranarayan Pillai, Marie Pellat, Aitor Lewkowycz, Erica Moreira, Rewon Child, Oleksandr Polozov, Katherine Lee, Zongwei Zhou, Xuezhi Wang, Brennan Saeta, Mark Diaz, Orhan Firat, Michele Catasta, Jason Wei, Kathy Heiir--ellsttem. Douglas Eck, Jeff Dean, Slav Petrov, and Noah Fiedel. Palm: Scaling language modeling with pathways, 2022. URL https://arxiv www.riiwww..o.bbooaeee www.ooggg///....oooooooooo....g/ /2004.//22....2 02311.
11

---

Published as a conference paper at ICLR 2026
Wei Chu, Lihong Li, Lev Reyzin, and Robert Schapire. Contextual bandits with linear payoff functions. In Proceedings of the fourteenth international conference on artificial intelligence and statistics, pp. 208-214. JMLR Workshop and Conference Proceedings, 2011.
Karl Cobbe, Vineet Kosaraju, Mohammad Bavarian, Mark Chen, Heewoo Jun, Lukasz Kaiser, Matthias Plappert, Jerry Tworek, Jacob Hilton, Reiichiro Nakano, Christopher Hesse, and John Schulman. Training verifiers to solve math word problems, 2021. URL https://arxiv.orr abs/2110. 14168.
Gheorghe Comanici, Eric Bieber, Mike Schaekermann, Ice Pasupat, Noveen Sachdeva, Inderjit Dhillon, Marcel Blistein, Ori Ram, Dan Zhang, Evan Rosen, al. Gemini 2.5: Pushing the frontier with advanced reasoning, multimodality, long context, and next generation agentic capabilities. arXiv preprint arXiv:2507.0626/ 2025.
Mehul Damani, Idan Shenfeld, Andi Peng, Andreea Bobu, and Jacob Andreas. Learning how hard to think: Input-adaptive allocation of Im computation. arXiv preprint arXiv:2410.04707 2024.
Dee,Seek-AI, Daya Guo, Dejian Yang. Haowei Zhang, Junxiao Song, Ruoyu Zhang, Runxin Xu, Qihao Zhu, Shirong Ma, Peiyi Wang, Xiao Bi, Xiaokang Zhang. Xingkai Yu, Yu Wu, Z. F. Wu, Zhibin Gou, Zhihong Shao, Zhuoshu Li, Ziyi Gao, Aixin Liu, Bing Xue, Bingxuan Wang, Bochao Wu, Bei Feng. Chengda Lu, Chenggang Zhao, Chengqi Deng, Chenyu Zhang, Chong Ruan, Damai Dai, Deli Chen, Dongje Ji. Erhang Li, Fangyun Lin, Fucong Dai, Fuli Luo, Guangbo Hao, Guanting Chen, Guowei Li. H. Zhang, Han Bao. Hanwei Xu, Haocheng Wang, Honghui Ding, Huajian Xin, Huazuo Gao, Hui Qu, Hui Li, Jianzhong Guo, Jiashi Li, Jiawei Wang. Jingchang Chen, Jingyang Yuan, Junjie Qiu, Junlong Li, J. L. Cai, Jiaqi Ni, Jian Liang, Jin Chen, Kai Dong, Kai Hu, Kaige Gao, Kang Guan, Kexin Huang, Kuai Yu, Lean Wang, Lecong Zhang, Liang Zhao, Litong Wang, Liyue Zhang, Lei Xu, Leyi Xia, Mingchuan Zhang, Minghua Zhang, Minghui Tang, Meng Li, Miaojun Wang, Mingming Li, Ning Tian, Panpan Huang, Peng Zhang, Qiancheng Wang. Qinyu Chen, Qiushi Du. Ruiqi Ge. Ruisong Zhang, Ruizhe Pan, Runji Wang. R. Chen. R. Jin, Ruyi Chen. Shanghao Lu, Shangyan Zhou, Shanhuang Chen, Shengfeng Ye, Shiyu Wang, Shuiping Yu, Shunfeng Zhou, Shuting Pan, S. Li, Shuang Zhou, Shaoqing Wu, Shengfeng Ye, Tao Yun, Tian Pei, Tianyu Sun, T. Wang. Wangding Zeng, Wanjia Zhao, Wen Liu, Wenfeng Liang, Wenjun Gao, Wenqin Yu, Wentao Zhang, W. L. Xiao, Wei An, Xiaodong Liu, Xiaohan Wang, Xiaokang Chen, Xiaotao Nie, Xin Cheng, Xin Liu, Xin Xie. Xingchao Liu, Xinyu Yang, Xinyuan Li, Xuecheng Su, Xuheng Lin, X. Q. Li, Xiangyue Jin, Xiaojin Shen, Xiaosha Chen, Xiaowen Sun, Xiaoxiang Wang, Xinnan Song, Xinyi Zhou, Xianzu Wang. Xinxia Shan, Y. K. Li, Y. Q. Wang. Y. X. Wei, Yang Zhang, Yanhong Xu. Yao Li, Yao Zhao, Yaofeng Sun. Yaohui Wang, Yi Yu, Yichao Zhang, Yifan Shi, Yiliang Xiong, Ying He, Yishi Piao, Yisong Wang, Yixuan Tan, Yiyang Ma, Yiyuan Liu, Yongqiang Guo, Yuan Ou, Yuduan Wang, Yue Gong, Yuheng Zou, Yujia He, Yunfan Xiong, Yuxiang Luo, Yuxiang You, Yuxuan Liu, Yuyang Zhou, Y. X. Zhu, Yanhong Xu, Yanping Huang, Yaohui Li, Yi Zheng, Yuchen Zhu Yunxian Ma, Ying Tang, Yukun Zha, Yuting Yan, Z. Z. Ren, Zehui Ren, Zhangli Sha, Zhe Fu, Zhean Xu, Zhenda Xie, Zhengyan Zhang, Zhewen Hao, Zhicheng Ma, Zhigang Yan, Zhiyu Wu, Zihui Gu, Zijia Zhu, Zijun Liu, Zilin Li, Ziwei Xie, Ziyang Song, Zizheng Pan, Zhen Huang, Zhipeng Xu, Zhongyu Zhang, and Zhen Zhang. Deepseek-rl: Incentivizing reasoning capability in Ilm via reinforcement learning, 2025. 12948.
URL
Yihan Du, Wei Chen, Yuko Kuroki, and Longbo Huang. Collaborative pure exploration in kernel bandit. arXiv preprint arXiv:21 10..1111.,117,,, 2021.
Eyal Even-Dar, Shie Mannor, and Yishay Mansour, Pac bounds for multi-armed bandit and markoy decision processes. In Computational Learning Theory: 15th Annual Conference on Computational Learning Theory, COLT 2002 Sydney, Australia, July 8-10, 2002 Proceedings 15, pp. 255-270. Springer, 2002.
Eyal Even-Dar, Shie Mannor, Yishay Mansour, and Sridhar Mahadevan. Action elimination and stopping conditions for the multi-armed bandit and reinforcement learning problems. Journal of machine learning research, 7(6), 2006.
Xidong Feng, Ziyu Wan, Muning Wen, Stephen Mareus McAleer, Ying Wen, Weinan Zhang, and Jun Wang. Alphazero- like tree-search can guide large language model decoding and training. arXiv preprint arXiv:2309.171 17179, 2023.
12

---

Published as a conference paper at ICLR 2026
Tanner Fiez, Lalit Jain, Kevin G Jamieson, and Lillian Ratliff. Sequential experimental design for transductive linear bandits. Advances in neural information processing systems, 32, 2019.
Dylan Foster and Alexander Rakhlin. Beyond ucb: Optimal and efficient contextual bandits with regression oracles. In International conference on machine learning, pp. 3199-3210. PMLR, 2020.
Dylan J Foster, Sham M Kakade, Jian Qian, and Alexander Rakhlin. The statistical complexity of interactive decision making. arXiv preprint arXiv:22 12. 13487, 2021.
Aurélien Garivier, Hédi Hadiji, Pierre Menard, and Gilles Stoltz. Kl-ucb-switch: optimal regret bounds for stochastic bandits from both a distribution-depeedent and a distribution-fre viewpoints. Journal of Machine Learning Research, 23 2022.
Zhibin Gou, Zhihong Shao, Yeyun Gong, Yelong Shen, Yujiu Yang, Nan Duan, and Weizhu Chen. Critic: Large language models can self-correct with tool-interactive critiquing. arXiv preprint vrXiv:2305.11738 2023.
Aaron Grattafiori, Abhimanyu Dubey, Abhinav Jauhri, Abhinav Pandey, Abhishek Kadian, Ahmad Al-Dahle, Aiesha Letman, Akhil Mathur, Alan Schelten, Alex Vaughan, et al. The llama 3 herd of models. arXiv preprint arXiv:2407.2178. 2024.
Dan Hendrycks, Collin Burns, Saurav Kadavath, Akul Arora, Steven Basart, Eric Tang, Dawn Song, and Jacob Steinhardt. Measuring mathematical problem solving with the math dataset. NeurlPS, 2021.
Jordan Hoffmann, Sebastian Borgeaud, Arthur Mensch, Elena Buchatskaya, Trevor Cai, Eliza Rutherford, Diego de Las Casas, Lisa Anne Hendricks, Johannes Welbl, Aidan Clark, al. Training compute-optimal large language models. arXiv preprint arXiv:2203. 1555, 2022.
Naman Jain, King Han, Alex Gu, Wen-Ding Li, Fanjia Yan, Tianjun Zhang, Sida Wang, Armando Solar-Lezama, Koushik Sen, and Ion Stoica. Livecodebench: Holistic and contamination free evaluation of large language models for code, 2024. URL https: //arxiv. org aprrg//bb/240333 07974.
Kevin Jamieson and Robert Nowak. Best-arm identification algorithms for multi-armed bandits in the fixed confidence setting. In 2014 48th annual conference on information sciences and systems (CISS), pp. 1-6. IEEE, 2014.
Kevin Jamieson, Matthew Malloy, Robert Nowak, and Sébastien Bubeck. lil' ucd: An optimal exploration algorithm for multi-armed bandits. In Conference on Learning Theory, pp. 423-439. PMLR, 2014.
Shivaram Kalyanakrishnan, Ambuj Tewari, Peter Auer, and Peter Stone. Pac subset selection in stochastic multi-armed bandits. In ICML, volume 12, pp. 655-662, 2012.
Jared Kaplan, Sam McCandlish, Tom Henighan, Tom B. Brown, Benjamin Chess, Rewon Child, Scott Gray, Alec Radford, Jeffrey Wu, and Dario Amodei. Scaling laws for neural language models, 2020. URL htttts:///aaaeeee httl/oor///bbbo httpg//bb/20000 2001. 08361.
Zohar Karnin, Tomer Koren, and Oren Somekh. Almost optimal exploration in multi armed bandits. In International conference on machine learning, pp. 1238-12040 1246. PMLR, 2013.
Julian Katz-Samuels, Lalit Jain, Kevin G Jamieson, et al. An empirical process approach to the union bound: Practical algorithms for combinatorial and linear bandits. Advances in Neural Information Processing Systems, 33: 10377-100082, 2020.
Emilie Kaufmann and Shivaram Kalyanakrishnan. Information complexity in bandit subset selection. In Conference on Learning Theory, pp. 228-55. PMLR, 2013.
Woosuk Kwon, Zhuohan Li, Siyuan Zhuang, Ying Sheng, Lianmin Zheng, Cody Hao Yu, Joseph Gonzalez, Hao Zhang, and Ion Stoica. Efficient memory management for large language model serving with pagedattention. In Proceedings of the 29th Symposium on Operating Systems Principles, pp. 611-626, 2023.
Tor Lattimore and Csaba Szepesvári. Bandit algorithms. Cambridge University Press, 2020.
13

---

Published as a conference paper at ICLR 2026
Lihong Li, Wei Chu, John Langford, and Robert E Schapire. A contextual- approach to personalized news article recommendation. In Proceedings of the 19th international conference on World wide web, pp. 661-670, 2010.
Lisha Li, Kevin Jamieson, Giulia DeSalvo, Afshin Rostamizadeh, and Ameet Talwalkar. Hyperband: A novel bandit-based approach to hyperparameter optimization. Journal of Machine Learning Research, 18(185):1-53 2018.
Hunter Lightman, Vineet Kosaraju, Yura Burda, Harri Edwards, Bowen Baker, Teddy Lee, Jan Leike, John Schulman, Ilya Sutskever, and Karl Cobbe. Let's verify step by step, 2023. URL
https://arxxw.orr//bb///22 wtt.or/////b..////ee 20050.
Andrea Locatelli, Maurilio Gutzeit, and Alexandra Carpentier, An optimal algorithm for the thresh- http//awww..o.ggb///////aaa wwwg///g//////ooooo//o/o//oe/sr/ 1605. 08671.
olding bandit problem, 2016. URL https: //axxiv
Aman Madaan, Niket Tandon, Prakhar Gupta, Skyler Hallinan, Luyu Gao, Sarah Wiegreffe, Uri Alon, Nouha Dziri, Shrimai Prabhumoye, Yiming Yang, et al. Self-refine: Iterative refinement with self-feedbace Advances in Neural Information Processing Systems, 36:46534 46594, 2023.
Rohin Manvi, Anikait Singh, and Stefano Ermon. Adaptive inference-time compute: Lims can predict if they can do better, even mid-generation, 2024. URL
hwtp:/wwwwww.wgb/////bba/ 02725.
Marius Mosbach, Tiago Pimentel, Shauli Ravfogel, Dietrich Klakow, and Yanai Elazar. Few- shot fine-tuning vs. in-context learning: A fair comparison and evaluation, 2023. URL https: //arxiv.org/abs/23555 16938.
Niklas Muennighoff, Zitong Yang. Weijia Shi, Xiang Lisa Li, Li Fei-Fei, Hannaneh Hajishirzi, Luke Zettlemoyer, Percy Liang, Emmanuel Candès, and Tatsunori Hashimoto. s1: Simple test-time scaling, 2025. URL https: httpaaaiiv..ogggbb/.///// 2501. 19393.
OpenAI. Learning to reason with 1lms, 2024. tearning-to-reason-iii thmsl
URL https://opemai..om com/ innex/
Mark Rucker, Yinglun Zhu, and Paul Mineiro. Infinite action contextual bandits with reusable data exhaust. In International Conference on Machine Learning, pp.
PMLR, 2023.
Daniel Russo, Benjamin Van Roy, Abbas Kazerouni, Ian Osband, Zheng Wen, et al. A tutorial on thompson sampling. Foundations and Trends in Machine Learning, 11(1):1-96, 2018.
Chengshuai Shi, Kun Yang, Zihan Chen, Jundong Li, Jing Yang, and Cong Shen. Efficient prompt optimization through the lens of best arm identification. arXiv preprint arXiv:2402.09773, 2024.
Charlic Snell, Jachoon Lee, Kelvin Xu, and Aviral Kumar. Scaling I1m test-time compute optimally can be more effective than scaling model parameters, 2024. URL https::/arxxv..rr//aaaaa 03314.
Hanshi Sun, Momin Haider, Ruiqi Zhang, Huitao Yang, Jiahao Qiu, Ming Yin, Mengdi Wang, Peter Bartlett, and Andrea Zanette. Fast best-of-n decoding via speculative rejection. arXiv preprint arXiv:2410.20290, 2024.
Zhendong Tan, Xingjun Zhang, Chaoyi Hu, Yancheng Pan, and Shaoxun Wang. Adaptive rectification sampling for test-time compute scaling, 2025. URL https::/aaxxi.oor//aatttt 013900001331 01317.
William R Thompson. On the likelihood that one unknown probability exceeds another in view of the evidence of two samples. Biometrika, 25(3/4):25--94, 1933.
Jonathan Uesato, Nate Kushman, Ramana Kumar, Francis Song, Noah Siegel, Lisa Wang, Antonia Creswell, Geoffrey Irving, and Irina Higgins. Solving math word problems with process-and outcome- based feedback. arXiv preprint arXiv:2211.14275. 2022.
Soffa S Villar, Jack Bowden, and James Wason. Multi-armed bandit models for the optimal design of clinical trials: benefits and challenges. Statistical science: a review journal of the Institute of Mathematical Statistics, 3092):199, 2015.
14

---

Published as a conference paper at ICLR 2026
Junlin Wang, Shang Zhu, Jon Saad-Falcon, Ben Athiwaratkum, Qingyang Wu, Jue Wang, Shuai- wen Leon Song, Ce Zhang, Bhuwan Dhingra, and James Zou. Think deep, think fast: In- vestigating efficiency of verifier-free inference-time-scalii methods, 2025a. URL https: //arxiv.org////// http://///b//.....///4/ 14047.
Xinglin Wang. Shaoxiong Feng, Yiwei Li, Peiwen Yuan, Yueqi Zhang, Chuyi Tan, Boyuan Pan, Yao Hu, and Kan Li. Make every penny count: Difficulty-ddaptive confisesstt for cost-efficient reasoning, 2025b. URL
13457.
Xuezhi Wang, Jason Wei, Dale Schuurmans, Quoc Le. Ed Chi, Sharan Narang, Aakanksha Chowdh- ery, and Denny Zhou. Self-consistency improves chain of thought reasoning in language models. arXiv preprint arXhv:203.11117,, 2022.
Jason Wei, Xuezhi Wang, Dale Schuurmans, Maarten Bosma, Brian Ichter, Fei Xia, Ed Chi, Quoe Le, and Denny Zhou. Chain-of-thought prompting elicits reasoning in large language models, 2023. URL https: http://xxxii.orr/////////////// 11903.
Yang, Anfeng Li, Baosong Yang, Beichen Zhang, Binyuan Hui, Bo Zheng, Bowen Yu, Chang Gao, Chengen Huang, Chenxu Lv, Chujie Zheng, Dayiheng Liu, Fan Zhou, Fei Huang, Feng Hu, Hao Ge, Haoran Wei, Huan Lin, Jialong Tang, Jian Yang, Jianhong Tu, Jianwei Zhang, Jianxin Yang. Jiaxi Yang. Jing Zhou, Jingren Zhou, Junyang Lin, Kai Dang, Keqin Bao, Kexin Yang, Le Yu, Lianghao Deng, Mei Li, Mingfeng Xue, Mingze Li, Pei Zhang, Peng Wang. Qin Zhu, Rui Men, Ruize Gao, Shixuan Liu, Shuang Luo, Tianhao Li, Tianyi Tang, Wenbiao Yin, Xingzhang Ren, Xinyu Wang, Xinyu Zhang. Xuancheng Ren, Yang Fan, Yang Su, Yichang Zhang, Yinger Zhang, Yu Wan, Yugiong Liu, Zekun Wang, Zeyu Cui, Zhenru Zhang, Zhipeng Zhou, and Zihan Qiu. Qwen3 technical report, 2025. URL https: /awpxvv.vvvv///////////////a/aa org.abs. 2505. 09388..
Shunyu Yao, Dian Yu, Jeffrey Zhao, Izhak Shafran, Thomas L. Griffiths, Yuan Cao, and Karthik Narasimhan. Tree of thoughts: Deliberate problem solving with large language models, 2023. URL https://arxii org. 2b35/ 10601.
Zhenru Zhang, Chujie Zheng, Yangzhen Wu, Beichen Zhang, Runji Lin, Bowen Yu, Dayiheng Liu, Jingren Zhou, and Junyang Lin. The lessons of developing process reward models in mathematical reasoning. arXiv preprint arXiv:2501.07300/ 2025a.
Zhenru Zhang, Chujie Zheng, Yangzhen Wu, Beichen Zhang, Runji Lin, Bowen Yu, Dayiheng Liu, Jingren Zhou, and Junyang Lin. The lessons of developing process reward models in mathematical reasoning, 2025b. URL /ttpp:://arxiiii http://awxvv..../abba/a/ abs 2501 07301.
Yinglun Zhu and Paul Mineiro. Contextual bandits with smooth regret: Efficient learning in continu- ous action spaces. In International Conference on Machine Leurning, pp. 27594-27590. PMLR, 2022.
Yinglun Zhu and Robert Nowak. On regret with multiple best arms. Advances in Neural Information Processing Systems, 33:9050-9000 9066,, 2020.
Yinglun Zhu and Robert Nowak. Pareto optimal model selection in linear bandits. In International Conference on Artificial Intelligence and Statistics, pp. 6793-6813 13. PMLR, 2022.
Yinglun Zhu, Sumeet Katariya, and Robert Nowak. Robust outlier arm identification. In International Conference on Machine Learning, pp. 11566-11577 11555. PMLR, 2020.
Yinglun Zhu, Dongruo Zhou, Ruoxi Jiang, Quanquan Gu, Rebecca Willett, and Robert Nowak. Pure exploration in kernel and neural bandits. Advances in neural information processing systems, 34: 11618-11630, 2021.
Yinglun Zhu, Dylan J Foster, John Langford, and Paul Mineiro. Contextual bandits with large action spaces: Made practical. In International Conference on Machine Learning, pp. 27428-222222 2-4533.. PMLR, 2022a.
Yinglun Zhu, Julian Katz-Samuels, and Robert Nowak. Near instance optimal model selection for pure exploration linear bandits. In International Conference on Artificial Intelligence and Statistics, pp. 6735-6769. PMLR, 2022b.
15

---

Published as a conference paper at ICLR 2026
A RELATED WORK
Test-time compute techniques. Scaling test-time compute (TTC) has emerged as powerful class of methods for improving the performance of large language models, typically without requiring additional parameter updates. In-context learning (Brown et al., 2020), including its scaling to many-shot regimes (Agarwal et al., 2024; Bertsch et al., 2024), and prompting-bar -based methods such as Chain-of- Chain-of-Thought (Wei et al., 2023) and Tree-of- Thought (Yao et al., 2023; Feng et al., 2023), have demonstrated that carefully designed test-time techniques can match or even surpass finetuned models (Mosbach et al., 2023). Self-reflection (Madaan et al., 2023) is another popular technique for leveraging TTC to improve performance: by prompting the LLM to iteratively refine its own generations, the model can produce higher-quality responses across range of tasks (Chen et al., 2023; Gou et al., 2023). Muennighoff et al. (2025) further demonstrates that simply increasing the number of generated "thinking" tokens leads to substantial performance gains.
Repeated sampling methods-moo most notably Best-of-N (Brown et al., 2024; Snell et al., 2024; Wang et al., 2022)- have become popular for scaling test-time compute, especially when combined with high-quality reward models (Cobbe et al., 2021; Uesato et al., 2022; Lightmnn et al., 2023; Zhang et al., 2025a). Building on this line of work, recent-and in some cases concurrett-effort. have proposed adaptive variants of Best-of-N that dynamically allocate compute for given query (Sun et al., 2024; Manvi et al., 2024; Tan al., 2025). However, these methods focus on adaptive allocation within an individual query, without considering opportunities to redistribute compute across a set of queries. contrast, we study strategic compute allocation across multiple queries, introducing an additional layer of optimization for example, deciding when to transfer unused budget from casier queries to harder ones. The problem setting in Damani et al. (2024) is closely related, as they also consider multi-query compute allocation. However, their approach relies on a two-stage schedule that requires training an additional model in the first stage to guide compute distribution, incurring extra compute cost. In contrast, we formulate the problem and develop fally adaptive algorithms that learn to allocate compute on the fly, without any additional training overhead. Furthermore, we provide the first theoretical result that provably demonstrates the advantage of strategic test-time compute allocation over uniform allocation.
novel bandit learning task
Bandit learning and pure exploration. Bandit learning is fundamental framework for sequential decision making under uncertainty, where an agent must choose among a set of actions (or arms) to optimize long-term objective with limited feedback (Bubeck & Cesa-Bias Desa-Bianchi, 2012: Lattimore & Szepesvári, 2020). Popular algorithms include Upper Confidence Bound (UCB, Auer et al. (2002); Audibert & Bubeck (2009); Chu et al. (2011); Zhu & Nowak (2020, 2022); Garivier et al. (2022)), which selects the action with the highest upper confidence bound: Thompson Sampling (Thompson, 1933; Chapelle & Li, 2011: 1; Agrawal & Goyal, 2012; Russo et al., 2018), which selects the action with the highest sampled reward from the posterior; and inverse gap weighting strategies (Foster & Rakhlin, 2020; Foster et al., 2021: Zhu et al., 2022a; Zhu & Mineiro, 2022; Rucker et al., 2023), which sample actions with probabilities inversely proportional to their estimated reward gaps. Bandit algorithms have been widely applied in domains such as online recommendation systems (Li et al., 2010), clinical trials (Villar al., 2015). hyperparameter tuning (Li et al., 2018), and more recently applications with LLMs (Shi et al., 2024; Chen et al., 2024).
Pure exploration (Bubeck et al., 2009; Jamieson & Nowak, 2014), also known as the best arm identification (BAI) problem, is key subfield of bandit learning that aims to identify high-perforning arms using few samples as possible. Core algorithms include successive elimination (Even- Dar et al., 2002, 2006; Karnin et al., 2013), UCB-based strategies (Kalyanakrishnan et al., 2012: Kaufmann & Kalyanakrishnan, 2013: Jamieson et al., 2014), and gap-based sampling methods (Locatelli et al., 2016). Recent extensions generalize these techniques to more expressive function classes, including linear models (Fiez et al., 2019; Katz-Samuels et al., 2020; Zhu et al., 2022h), kernel functions (Du et al., 2021), and neural networks (Zhu et al., 2021). Our work introduces novel pure-exploration-st// bandit formulation, tailored to LLM test-time compute allocation-a setting not previously explored in this context. We treat each query as a bandit action and adaptively allocate compute to maximize the fraction of queries correctly answered under a fixed compute budget. This formulation enables the use of classical bandit techniques such elimination rules, confidence bounds, and gap- based sampling. In addition, we propose a new entropy-basee based sampling strategy (Section 3.3) that prioritizes queries with diverse response patterns. While our formulation is
16

---

Published as a conference paper at ICLR 2026
conceptually related to the thresholding bandit problem and its variants (Locatelli al., 2016; Zhu & Nowak, 2020), it departs fundamentally in its objective. Thresholding bandits aim to identify actions (queries) whose expected reward exceeding given threshold. In contrast, goal is to generate at least one high-quality response for each query, regardless of its expected score.
B ALGORITHM I WITH DIFFERENT EXPLORATION RULES
Our main algorithmic framework (Algorithm 1) presented with a simple exploration rule that explores all queries within the active set (lines 4-5). In practice, this rule can be flexibly extended to incorporate diverse exploration objectives. Motivated by developments in the bandit pure exploration literature, we introduce several alternative exploration rules in the following. We use g(x) to denote the response set to query z, and (x) = [g(x)] to denote the number generations so far.
Upper confidence bound (UCB). For any active query E A. let f(z) Emeeg(a) r(z, v:)/N(x) denote the empirical average reward based on previously collected re- sponses. Let À be a hyperparameter. At cach round, the UCB exploration rule selects the query based the following criteria:
arg max F(x) + (N 1/2
aEA
This exploration rule follows the principle of optimism in the face of uncertainty (Kalyanakrishnan et al., 2012; Jamieson et al., 2014), and prioritizes on selecting queries that are more likely to be (1/x)-1// solved (i.e., those with higher average rewards). The term AN(x) 1/22 is used to construct the upper confidence bound of the reward.
GAP. For any active query x E A, let F(z) E Egeetta) r(z, 31) (x)) denote the empirical average reward based on previously collected responses. each round, the GAP exploration rule selects the query based on the following criterion:
arg min (7 -(x)) (x)
=EA
This exploration rule prioritizes queries whose estimated reward is close the elimination threshold y, with preference toward less-explored queries. The weighting term N(x)-1/2 ensures that compute is allocated inversely proportional to the reward gap from the elimination threshold (Locatelli et al., 2016).
ENTROPY. For any active query x E A, let {vx) be the set of distinct responses in g(x), and define the empirical probability of observing response U as (&( := {(i yi uk Vi E g(z))|/N(()) -E px(x)) log px(x)) denote the entropy of the empirical response distribution p(z). Let À > 0 be a hyperparameter. At each round, the ENTROPY exploration rule selects the query based on the following criterion:
Let H (x)
arg max H(x) + 1/xx aEAA
This exploration rule, proposed in our work, prioritizes queries that elicit more diverse set of responses, indicated by higher entropy. The term AN encourages exploration of under-explored queries by balancing the trade-off between response diversity and sample count.
C SUPPORTING RESULTS FROM SECTION 3.4
C.I PROOF OF THEOREM
Theorem Assume Assumption and fix any (0, To output correct responses for all queries S with probability at least 1 - Algorithm requires a total budget Bours Ô(Σrs ).), and this matches the information-theorete lower bound up to logarithmic factors. In contrast, a uniform (S) allocation strategy requires budget Bunif = A(max I. to achieve the same guarantee.
17

---

Published as a conference paper at ICLR 2026
Proof. We first prove that Bours = O A.) under the elimination rule of Algorithm 1. Note that under Assumption I, query is eliminated and only it is correctly answered." For the upper bound, we denote 8 8/|S]. n2 = 2. log and consider the following event:
1
Ex == (query will be correctly answered within fhr random generations).
We know that Eg happens with probability at least δ as the probability of E, is upper bounded by 8:
(1 -a) ns S As Ds log 8.
P P(E,) 21-6, As a result, the with probability of at least - 6. Algorithm 1 and its variants in Appendix B output correct responses for all queries with compute budget Eves f2 Exes A. log 5
union bound over x E leads to P(UDES 2x) 1- 2111211111
EacsP(E,
Is OEErs 2.).
For the lower bound, we quantify the amount of compute ng needed to correctly answer any fixed query æ E S with probability at least 1- 8. Recall that each independent generation answers correctly with probability A₂ (0, Therefore, after nı independent generations, the probability that none of them is correct equals
P(query æ is not correctly answered within generations) =(1 AA)))
an algorithm outputs a correct response for all queries in S with probability at least 6, then in particular it must output a correct response for this fixed query with probability at least
which requires (1 Aa)"= n= I 8. Taking logarithms gives ng 2 -log(1-A.) Using the inequality
log[1/6]
- -log(1 w) I I-u for all E (0, 1), we further obtain
1 1 log
ng
1-Ax log 1 & In particular, if A < 1 c for all x E S and some fixed constant E (0, 1), then the lower bound simplifies to 52 22 ES A ), which matches our upper bound up to logarithmic factors.
Summing the above inequality over all r yields the compute lower bound Q
As for uniform compute allocation, it allocates the same amount of compute for each query E S. We thus only need to quantify the amount of compute needed to correctly answer = arg min A... the hardest query within S, with high probability. The above lower bound analysis shows that one needs D - log 2t) compute to correctly answer the ã with probability at least 8. Since uniform allocation assigns the same budget to every query, the total compute must satisfy Bunir =
2 [S] log ) = Ammaxrs A.
C.2 PROOF PROPOSITION
Proposition 1. Assume Assumption 1, and fix any 8 E (0, 1). Consider any two-stage explore-ther commit algorithm that allocates a fixed exploration budget m to every query in the first stage. To output correct responses for all queries in S with probability at least - ô, any such algorithm must incur total compute BETС S'EEes max(m, J.
Proof. Suppose the ETC algorithm allocates m generations to every query x E S in the exploration stage. Let n denote the total number of generations allocated query across both stages. Then
we have ng 2 m for all E S. On the other hand, by the same per-query lower bound argument used in the proof of Theorem 1. to correctly answer query with probability at least is necessary that
l [ log })) Combining these two necessary conditions yields =
1. log }). Summing over all x E gives BETC DiEres max(m, A. i(E max(m, ).))
O(max(m, &r
1. log !))
SSince the elimination rule works for all variants of Algorithm I introduced in Appendix B, the guarantee also holds for these variants.
18

---

Published as a conference paper at ICLR 2026
D OTHER DETAILS AND RESULTS FOR EXPERIMENTS
D.1 ADDITIONAL DETAILS ON EXPERIMENTAL SETUPS
D.1.1 ADDITIONAL HYPERPARAMETERS
We conduct all experiments 2023) for LLM response generation, with temperature of 0.6.
NVIDIA RTX 6000 Ada GPUs. We use vLLM (Kwon al.,
The extended exploration rules introduced Appendix require additional hyperparameters. For UCB, we set À = 1; for ENTROPY, we set À = 3. For UCB, GAP, and ENTROPY, we additionally set max samples hyperparameter to prevent over-allocation of compute on difficult queries. The values of the max_samples hyperparameter will be introduced along with the experiment results.
D.1.2 ADDITIONAL DETAILS ON MATH-500-HARD HARD DATASETS
As discussed Section 4.1, we construct the MATH-500-Hard dataset by removing queries from MATH-500 that can be solved with 16 units of compute. On top of this, construct another dataset by removing queries that can be solved with units of compute, and we call this dataset MATH-500-Harde After removing these relatively easy queries, HATH-500-Harde contains 71 challenging queries, and MATTH5500-Hard-I contains 56 challenging queries. We further divide MATH-500-Hard queries into two subsets: the subset that cannot be solved after allocating M compute units (Unsolvable) and the rest (Solvable). We set M = 500 for Llama-3.2- 3.2--B-IInstreee Instruutt, and M 350 for Llama-3. 1e-1--8Immme Instruct. For normalized coverage, we calculate based on the percentage of solved questions over all solvable questions. On these MATH-500-Hamm datasets, we select the best from max_samples E (36, 48, 64] for UCB, GAP, and ENTROPY. The setting of max_samples is discussed in the same way Section 3.3.
1.1.3 PROMPTS SELECTION
Prompts for MATH-500. We include four in-context examples in the prompt for MATH-500 for Llama models. An illustrative example used in experiments is shown below.
MATH prompts
Problem: det 2 and det 12, then find det(A B).
Solution: We have that
det A B) = det(A) det(BB = (2) (12) 24
Final Answer: The final answer is 24. hope it correct.
Problem:
factual test question}
Solution:
Prompts for Benehh We use the prompt provided on the official GitHub of Live- CodeBench (Jain al., 2024). The prompt used in experiments is shown below.
19

---

Published as a conference paper at ICLR 2026
LiveCodeBench CoT prompt
You are given Python function and assertion containing imput to the function. Complete the assertion with a literal (no unsimplified expressions, no function calls) containing the output when executing the provided code the given input, even if the function is incorrect incomplete. Do NOT output any extra information. Execute the program step by step before arriving an answer, and provide the full assertion with the correct output in [ANSWER] and [/ANSWE]] tags, following the examples.
[PYTHON] def performOperationn "hi") == ?? [/PYTHON]
==s return "b" s "a" assert perform@peratione
[THOUGHT] Let's execute the code step by step: The function performoperation defined, which takes a single argument S. 2. The function is called with the argument "hi", so within the function, s is initially "hi". 3. Inside the function, s is concatenated with itself, so becomes "hihi" 4. The function then returns new string that starts with "b", followed by (now "hihi"), and ends with "a". 5. The return value is therefore "bhihia". [/THOUGHT] [ANSWER] assert performOperation(s "hi") == "bhihia" [/ANSWER]
AIME prompt The same format MATH-500 prompts, but without any in-context examples.
AIME prompt example
Problem:
{actual test question)
Solution:
Output segmentation for PRMs. The output from our models follow a structured format shown below. We segment each solution into individual reasoning steps using the marker Step ##, and feed the resulting step-level segments to the PRM Owen2.5-Math-ee MrMattt-P----77 7B for evaluation. Prior work commonly uses n step delimiter and finds it effective (Zhang al., 2025b). In our setting, however, we find that Step ## better matches the formatting of our outputs and yields higher overall performance; we therefore adopt Step ## for segmentation throughout (see Table 1).
Expected generation from both Llama models about MATH questions
Step 1: [Description of first step]
Step 2: [Description of second step]
Step 3: [Description of third step] Step...
The final answer is: $)
Table 1: Effects of segmentation keywords on MATH-500 with Llama-3, 3. 1-8B--B--- -B-Instruct
Method
8
16
32
Uniform (delimiter: ##)
60.26%
63.30%
66.56%
67.80%
Uniform (delimiter: \,n),n)
59.06%
61.50%
64.36%
64.76%
Elim (delimiter: ##)
62.50%
65.90%
67.96%
69.46%
Elim (delimiter: \n)n)
61.96%
64.66%
65.66%
66.20%
D.2 ADDITIONAL EXPERIMENTAL RESULTS
D.2.1 ADDITIONAL EXPERIMENTS ON MATH-500 AND LIVECODEBENCH
Results of extended exploration rules MATH-500. Fig. 6 provides the coverage performance of UCB and GAP on the MATH-500 dataset; the max samples hyperparameter for both UCB and GAP are provided in Table 2.
20

---

Published as a conference paper at ICLR 2026
0.7
Coverage
Coverage
0.6
0.5
Uniform GT
Two Stage PRM
Uniform GT
Two Stage PRM
0.4
Elim PRM
UCB:PRM
0.65
Elim-PRM
UCB PRM
Gap-PRM
Elim-GT
Gap-PAM
Elim-GT
Gap GT
Gap
4
16 Average Compute Budget
32
16
32
Average Compute Budget
Figure 6: Results MATH-500 with extended exploration rules.
Left: results with
Llama-3.24 2-0--3..--BIIIstuuce Right: results with Llama-3. 3. 1-8t-Inst
Table 2: The choice of max_samplee by model-scene and compute budget on MATH-500.
Scene
Avg. budget 4
Avg. budget 8
Avg. budget 16
Avg. budget 32 (max_samples)
(max_samples)
(max_samples)
(max_samples)
Llama-3. 1 w/ GT
40
40
120
300
Llama-3.1 w/ PRM
12
40
80
120
Llama-3. 3. w/ GT
40
40
120
120
1..a3..22 w/ PRM
12
12
60
60
Online compute allocation on MATH-500. We further evaluate a streaming variant of Algorithm in Table where queries arrive sequentially. this setting, compute allocated on a per-query basis over 500 rounds, rather than jointly across the entire dataset. For the streaming variant, we tune the max max_samples hyperparameter via binary search over the grid (15 + 5k .,....377 3,3 Even under the streaming constraint, our algorithm still outperforms uniform allocation with pool access.
Table 3: Results of the streaming variant of Algorithm on MATH-500 with L1ama-3. 2--.1..-1---nruueee
Average compute budget
Method
4
16
32
Coverage
Uniform-PRM
41.80%
50.75%
59.60%
68.80%
Elim-PRM
43.70%
53.75%
63.45%
71.45%
Elim-PRM (Streaming)
42.65%
52.85%
62.95%
71.05%
Accuracy
Uniform-PRM
35.00%
40.70%
46.45%
49.45%
Elim-PRM
37.15%
43.70%
48.90%
52.60%
Elim-PRM (Streaming)
36.55%
43.10%
48.55%
52.20%
Token-controlled evaluation. To isolate the effect of allocation from differences in total generation length, token-controlled variant of Algorithm 1. Specifically, we first record the total number of generated tokens under UNIFORM allocation, and then run Algorithm while enforcing the same total token budget booken-macheed) As shown in Table 4 (MATH-500) and Table 5 (LiveCodeBench), Algorithm remains consistently better than under UNIFORM with the same number of tokens.
Choices of the threshold 7. The elimination threshold 7 controls which queries are confidently answered and can be removed from the active set. Since higher reward scores generally indicate higher-quality responses, setting a high threshold y E [0, 1] is natural. We conduct ablations with y (0.97, 0.98, 0.99, 1.0) MATH-500 and report the results in Fig. 7. We observe that 7 1.0
21

---

Published as a conference paper at ICLR 2026
Table 4: Token matched evaluation on MATH-500 with 3--
Average compute budget
Method
8
16
Coverage
Uniform-PRM
41.80%
50.75%
59.60%
68.80%
Elim-PRM
43.70%
53.75%
63.45%
71.45%
Elim-PRM (Token Matched)
42.70%
52.10%
60.90%
69.95%
Accuracy
Uniform-PRM
35.00%
40.70%
46.45%
49.45%
Elim-PRM
37.15%
43.70%
48.90%
52.60%
Elim-PRM (Token Matched)
36.50%
42.60%
48.05%
51.85%
Table Token matched evaluation on LiveCodeBench with DeppSek-RI-Distlll-Lamm
Average compute budget
Method
4
8
16
32
Coverage
Uniform-CC -T
73.43%
81.05%
86.01% %
89.87%
Elim-GT
84.66%
91.02%
94.00%
94.00%
Elim-GT (Token Matched)
84.50%
91.65%
94.00%
94.00%
performs slightly better, likely because Qwen2. 5-Math to answers it deems correct-a property specific to this PRM. Importantly, across all tested values, our method consistently outperforms the uniform allocation baseline, indicating that Algorithm is robust to variations in ?
assigns a deterministic score of 1.0
0.55
0.50
Accuracy
0.45
Uniform RRMM
(y=
Elim-PRM (y 0.98)
0.35
Elim-PRM (y 0.97)
Average Compute Budget
Figure 7: Effect of y on performance for Llama-3, 2-B--Insttt
Effect of PRM quality. study the impact PRMs quality on the effectiveness of our method, we conducted additional experiments using relatively weak PRM, Qwen2. --M.MMhh777-PPRM000 in contrast to the strong PRM, Qwen2.5-M6 5-Math-PRMM-7 employed in the main experiments. As shown in the Fig. 8, our method consistently outperforms the baselines across all compute budgets under both weak and strong reward models. Notably, performance improves with a stronger oracle, suggesting that our algorithm benefits directly from higher-quality reward signals.
D.2.2 ADDITIONAL EXPERIMENTS AND ANALYSES ON H0TH-500-HARI
Fig. 9 presents additional results on SATHH-50--Hart datasets, including experiments with the GAP algorithm and experiments with the Llama-3 -8B-Insttrut model. These experiments show that
22

---

Published as a conference paper at ICLR 2026
62.0
0.68
Accuracy
60.0
0.66
58.0
0.64
56.0
Uniform- PRM
Uniform-PPP PRM
Wo--StagePRRM
0.60
Peo-Staag-PRR
54.0
Elim-PRM (Ours)
Ellm-RMM (Ours)
16
32
8
16
32
Average Compute Budget
Average Compute Budget
(a) Weak PRM: Qwen2.5-H8 -aatha78-PRM---
(b) Strong PRM: Qwen2. 5-Math-PRM-77
Figure 8: Effect of PRM quality MATH-500 (texttLLlama 3.B-BB-----B.. (-------nnttt) We com-
pare performance using a weaker PRM (Qwen2.5-Mar 5-Math-78--PMBON versus a stronger PRM
5aMth-PRMMM across compute budgets.
40
Uniform
30
Uniform
Normalized Coverage (%)
Stage
Normalizee Coverage (%)
Two-Stage
Elim
Eilm
30
UCB
UCB
20
Entropy
Entropy
20
10
16 Average Compute Budget
32
16 Average Compute Budget
32
(a) Coverage with 3. 2- M---IIstruutt on 4ATH-500-Hard-5
(b) Coverage with Llama-3. 1.ama-3.2111 2-11-1BB1 IIInnstrcc on MATH-500 H01H-500-Hard-16
35
Uniform
Uniform
Normalized Coverage (%6)
Two -tage
Normalizzed Coverage (%)
Tivo-Stage
UCB
20
UCB
21
Gap
Gap
Entropy
Entropy
10
16 Average Compute Budget
32
8
16
32
Average Compute Budget
(c) Coverage with Llama-3. -B-Instruct on MATH-500-Hard-8
(d) Coverage with Llama-3. 1-8B-Instruct on MATH-500-Harde
Figure 9; Coverage comparisons on MATH-500-Harr datasets with two language model different sizes: Llama-3,
and Llama-3, 3. 1--B-IIntruue
our algorithms introduced Section 3.3, particularly ENTROPY and UCB, achieve early advantages and outperform both UNIFORM and ELIMINATION.
Why ENTROPY works well on challenging datasets? We conduct detailed analysis on under- stand why ENTROPY performs particularly well on MATH-500- Hard datasets. Upon inspecting model responses to challenging queries, observe that unsolvable queries more likely yield invalid outputs (e.g., incomplete or improperly formatted), resulting lower entropy among their generated responses; contrast, solvable queries tend generate more diverse outputs (see Table and Table
23

---

Published as a conference paper at ICLR 2026
for statistics computed from 64 responses per query). Since ENTROPY prioritizes queries with higher entropy, naturally allocates more compute to those that are more likely to be solvable solvable-explainngg explaining its strong empirical performance to other challenging benchmarks, provided that invalid responses can be reliably identified. In such settings, ENTROPY offers effective means shift compute toward promising queries and achieve better performance under limited compute budget.
challenging problems. We expect this behavior to generalize
Table 6: Aggregated statistics by query group on MATH-500-Hard-8.
Query group
# #questions
Entropy (mean)
Invalid answers (%)
Unsolvable
49
4.26
19.45%
Solvable
22
4.52
12.45%
Table 7: Aggregated statisties by query group MATH-500-Hardd-16
Query group
# #questions
Entropy (mean)
Invalid answers (%
Unsolvable
43
4.33
18.39% 13. 10%
Solvable
13
4.53
Hyperparameter study of for ENTROPY and UCB. We study the effect of the exploration coefficient À in our Entropy- and UCB-based extension rules on the HATH-500-Harde dataset. Results Tables and suggest that ENTROPY benefits from larger X, while UCB prefers smaller
Table 8: Effect of À on Entropy-based Allocation.
choice / Avg compute budget
8
16
32
Entropy (A 1.0)
0.0%
5.4%
8.9%
17.9%
Entropy (A = 2.0)
0.0%
5,4%
8.9%
17.9%
Entropy (A 3.0)
0.0%
5,4%
8.9%
21.4%
Table 9: Effect on UCB-based Allocation.
choice / Avg compute budget
8
16
32
UCB (A 1.0)
0.0%
1.8%
3.6%
16.1%
UCB = 2.0)
0.0%
1.8%
3.6%
14.3%
UCB () = -3.0)
0.0%
1.8%
1.8%
12.5%
E
THE USE OF LARGE LANGUAGE MODELS (LLMs)
L.MMs were used to polish the writing of this paper.
24


Reviews:
[{"review_id": "tGcsWRjw5s", "reviewer": "Reviewer_pUPq", "summary": "This paper studies online algorithms for optimizing test-time compute allocation across a set of queries. Instead of uniformly allocating compute, the authors formulate the problem as a bandit learning task and develop an adaptive framework that learns to prioritize solvable instances, thereby avoiding excessive computation on unsolvable queries. The paper provides theoretical analysis showing improved compute efficiency over uniform allocation and demonstrates empirical gains on math and coding datasets.", "questions": "Please refer to the above weakness part.", "limitations": "", "rating": 4, "confidence": 3, "soundness": 3, "presentation": 4, "contribution": 2, "strengths": "1. The paper is clearly written and easy to follow.\n2. The proposed algorithmic framework is simple yet effective. It is modular and flexible—different components such as exploration and exploitation strategies can be adjusted to fit various scenarios.\n3. The work presents both theoretical and empirical results, and the experiments show consistent improvements over uniform allocation baselines.", "weaknesses": "1. While simplicity is a strength, the algorithmic contribution feels somewhat incremental. The framework extends a single-query setting to multiple queries in a relatively straightforward manner, and the theoretical analysis relies on standard probabilistic arguments. The results mainly show that the proposed method outperforms uniform allocation—which is expected—but do not clarify whether it outperforms other adaptive baselines such as ETC from a theoretical standpoint, or whether the approach is optimal or near-optimal.\n2. The assumption of having access to an accurate reward model and threshold parameter $\\gamma$ seems restrictive. It would be valuable to discuss how the method behaves or could be adapted under model inaccuracies.\n3. The current analysis focuses on a high-budget regime where every query can, in principle, be solved correctly. In realistic scenarios with a limited budget $B$, it would be more informative to characterize how many queries are expected to be solved correctly given small $B$, or how the algorithm trades off between query coverage and accuracy.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "HWTht57j2b", "reviewer": "Reviewer_q1r5", "summary": "This paper tackles the problem of test-time compute allocation for large language models (LLMs). Traditional approaches such as Best-of-N sampling uniformly allocate the same amount of compute to each query, regardless of difficulty. The authors argue that this is inefficient, as some queries are trivially solvable while others are much harder. They reformulate the problem as a bandit learning task, where each query is an “arm” and compute allocation is treated as exploration. Their proposed algorithms estimate query difficulty adaptively and allocate compute dynamically across queries. The core method progressively eliminates solved queries, redistributing compute to unsolved or more uncertain ones. Extensions include exploration variants, support for streaming queries, and token-level compute control. The authors provide theoretical results showing superior compute efficiency over uniform allocation and validate the method empirically on math and code benchmarks, achieving improvements under the same compute budget.", "questions": "Please refer to the weaknesses above.", "limitations": "", "rating": 2, "confidence": 4, "soundness": 2, "presentation": 2, "contribution": 2, "strengths": "- The idea of using bandit on allocating compute budget seems relatively novel.\n- A theory is provided.", "weaknesses": "- The major concern is that the authors should use stronger PRMs as reward \"oracles\", e.g. Reasonflux-PRM [1]. On the one hand, Qwen2.5-Math-PRM-7B is relatively old compared with the latest SOTA PRMs. On the other hand, Qwen3-1.7B and DeepSeek-R1-Distill-Llama-8B are reasoning models, thus using Qwen2.5-Math-PRM-7B, which is not trained for reasoning models, as the reward oracle may lead to suboptimal results. Reasonflux-PRM is a more recent SOTA PRM that can be used as the reward oracle to further validate the effectiveness of the proposed method on both non-reasoning and reasoning models.\n- Another concern on experiments is that self-consistency [2] is also a widely used algorithm to aggregate multiple response without PRMs, it will also be helpful to see how the proposed method will perform with that, such as [3].\n- The current paper fails to provide a systematic related work. While it is understandable that test-time scaling algorithms are a relatively new concept and most work can be classified as concurrent and does not need to be compared, it is not acceptable not to formally mention them.\n\n\n[1] ReasonFlux-PRM: Trajectory-Aware PRMs for Long Chain-of-Thought Reasoning in LLMs.\n\n[2] Self-Consistency Improves Chain of Thought Reasoning in Language Models.\n\n[3] Scaling LLM Test-Time Compute Optimally Can be More Effective than Scaling Parameters for Reasoning.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "G60BWfZDUK", "reviewer": "Reviewer_7ofK", "summary": "This paper addresses the inefficiency of uniform test-time compute allocation for LLMs and proposes a bandit learning framework for strategic compute distribution. It formulates test-time compute allocation as a bandit problem, treating each query as an action and allocating compute sequentially to maximize the number of correctly answered queries. Theoretical analysis and experiments validates the algorithm’s superior compute efficiency over uniform allocation.", "questions": "Please refer to the \"Weakness\" part.", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "1) The paper formulates test-time compute allocation as a bandit problem and demonstrates its superiority over uniform allocation through both theoretical and empirical analyses.  \n2) The framework supports multiple exploration strategies, aggregation methods, and settings, ensuring broad applicability.", "weaknesses": "1) The computational overhead of calculating Entropy and UCB should also be taken into account; comparing only the inference budget is insufficient. Explicit end-to-end run-time comparisons with each baseline are necessary.  \n2) The impact of the Entropy and UCB hyper-parameter $\\lambda$ has not been thoroughly analyzed; ablation experiments showing how $\\lambda$ affects performance across different tasks should be provided.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)