Uncertainty-Aware Budget Allocation for Adaptive Test-Time Reasoning

Abstract: Sampling multiple responses improves LLM reasoning, but uniform compute allocation is inefficient: easy questions are over-sampled while hard questions remain under-explored. We propose Uncertainty-Aware Budget Allocation (UAB), a concave integer optimization framework that reallocates a fixed sampling budget based on per-question uncertainty estimated at no additional inference cost. In Phase 1, every question receives one generation; its average negative log-likelihood (ANLL), extracted directly from output log-probabilities, serves as a difficulty signal while the generation contributes to the final vote. In Phase 2, the remaining budget is allocated by a marginal-greedy algorithm that solves a concave coverage-maximization surrogate exactly: uncertain questions receive more sampling budget while confident questions receive fewer additional samples. Evaluated on six open-weight and black-box models spanning 1.5B to 27B parameters and five reasoning benchmarks covering math, logic, and preference tasks, UAB outperforms baselines by up to +3% in average accuracy and up to +5% on individual benchmarks, with the largest gains in low-resource settings, requiring no auxiliary model or additional LLM call. Code is publicly available at https://github.com/manhitv/UAB.

• The paper introduces UAB, a training-free method that uses average negative log-likelihood from an initial generation to estimate question difficulty and greedily allocate remaining samples under a fixed budget.
• The method improves average accuracy over uniform self-consistency across six models and five benchmarks, with gains of up to 3% on average, 5% on individual tasks, and the strongest benefits at budgets of two to four samples.
• The paper shows UAB can reach the same accuracy with up to 27% fewer model calls, while its sensitivity bound explains why allocation remains effective despite imperfect probability estimates and its negligible overhead avoids auxiliary judge models.

Motivation and problem statement

Test-time compute scaling via repeated sampling and majority voting (self-consistency) improves LLM reasoning accuracy, but its marginal returns diminish with each additional sample, and the value of an extra sample varies sharply across questions. Uniform allocation—assigning every question the same number of generations $N$—wastes compute on questions that saturate after one correct sample while leaving genuinely hard questions under-explored. The inefficiency is most acute at small $N$, exactly where inference cost binds most. The paper asks: given a fixed total budget $B = N \times M$ for $M$ questions, how should samples be distributed to maximize accuracy?

Existing adaptive approaches each make a structural compromise that disqualifies them from this setting. Sequential early-stopping methods such as Adaptive-Consistency and ReASC optimize per-question stopping rules under an unbounded cap, so they cannot honor a fixed total budget a priori or reallocate savings across questions. Auxiliary-component methods (SEER, DSC, ACTSC, Sonata) require a separate judge model, trained estimator, or extra inference pass. Intra-trace methods (MUR, SelfBudgeter) allocate tokens within a single chain of thought rather than samples across questions. The authors position UAB as the first method that is simultaneously cross-question, fixed-budget, training-free, and inference-efficient in its difficulty estimation.

Method

UAB operates in two phases under total budget $B = N \times M$. In Phase 1, every question receives exactly one generation; the average negative log-likelihood (ANLL) of that generation,

$$s_i = -\frac{1}{L}\sum_{t=1}^{L} \log P(y_{i,t} \mid x_i, y_{i,<t}),$$

serves as a zero-cost difficulty signal extracted from log-probabilities already produced, and the generation itself contributes to the final vote. ANLL is mapped to a per-sample success probability via $p_i = e^{-s_i/T}$, where temperature $T$ controls allocation sharpness ($T \to \infty$ recovers uniform allocation; the default is $T{=}0.2$).

In Phase 2, the remaining $N$0 samples are allocated by a marginal-greedy procedure over the coverage objective $N$1 subject to $N$2, $N$3. This is a separable concave knapsack, for which greedy assignment by marginal gain $N$4 is provably exact—a classical result the authors restate with a KKT-style exchange-argument proof specialized to this setting. Uncertain questions maintain high marginal gains longer and receive more samples; confident questions saturate quickly.

Two theoretical results support the design. First, greedy exactness holds in $N$5 time. Second, a sensitivity proposition bounds objective degradation linearly: for any allocation, $N$6, and consequently the greedy allocation under estimated probabilities loses at most $N$7 relative to the optimum under true probabilities. This matters because ANLL is explicitly acknowledged to be a heuristic ranking signal rather than a calibrated probability estimator—the bound guarantees graceful degradation as long as estimation error stays bounded.

For black-box APIs that do not expose log-probabilities, the framework substitutes verbalized confidence scores (VCS): the model rates its own confidence 1–10 within the same generation, normalized to $N$8 and used directly as $N$9. The framework is otherwise signal-agnostic.

Main experimental results

Evaluation covers six models (Qwen2.5-1.5B/7B, Llama3.2-3B, GPT-OSS-20B, Gemma3-27B, plus Cohere Command-A via API) and five benchmarks (DeepScaler, GPQA Diamond, HH-RLHF, MMLU Formal Logic, MATH500), against Random, Length, Uniform (self-consistency), and LLM-Judge baselines at matched budgets, defaulting to $B = N \times M$0.

UAB achieves the highest average accuracy for every model. Representative averages at $B = N \times M$1:

Model	Uniform	LLM-Judge	UAB
Qwen2.5-1.5B	37.1	37.7	40.1
Llama3.2-3B	33.5	30.9	36.0
Qwen2.5-7B	48.5	49.0	51.2
GPT-OSS-20B	53.9	54.5	55.1
Gemma3-27B	60.8	61.3	61.8
Cohere (VCS)	66.0	66.0	67.2

Gains over Uniform reach up to $B = N \times M$2 on individual benchmarks (Formal Logic with Qwen2.5-1.5B) and up to $B = N \times M$3 on average, concentrated at low budgets ($B = N \times M$4–$B = N \times M$5); the advantage narrows but remains positive at $B = N \times M$6. The authors are candid that several individual cells fall within one standard deviation—for example Qwen2.5-7B on HH-RLHF—and rely on the consistent positive direction across all 15 model–benchmark cells rather than claiming per-cell significance. Notably, LLM-Judge underperforms Uniform in several settings, indicating binary easy/hard classification is too coarse for fine-grained allocation; on Llama3.2-3B its labels collapse entirely, causing accuracy degradation beyond $B = N \times M$7.

Cost-at-fixed-accuracy analysis shows concrete savings: on MATH500 with Qwen2.5-1.5B, UAB reaches 54.5% at $B = N \times M$8 versus $B = N \times M$9 for Uniform (19% fewer calls), widening to 27% at higher targets; on Formal Logic, UAB attains 43.7% at roughly $M$0 versus $M$1 for Uniform. In the large-model and black-box regime, gains shrink to $M$2–$M$3 over Uniform but remain consistent, and UAB improves over the single-call baseline by up to $M$4 (GPT-OSS-20B). Wall-clock overhead is negligible relative to Uniform (all runs under 2.5 minutes; UAB is even slightly faster than Uniform on long chain-of-thought MATH500, though 14% slower on short-generation Formal Logic due to the serial two-phase dependency).

Ablations

Allocation analysis validates the mechanism: allocated samples rise monotonically across ANLL deciles (roughly 3 samples for the most confident decile versus nearly double for the least confident), and correctness decreases monotonically with ANLL (Pearson $M$5–$M$6). Against annotated MATH500 difficulty levels L1–L5, UAB's allocation increases monotonically while LLM-Judge's is irregular.

The uncertainty-signal ablation is the most informative. Removing uncertainty entirely (uniform fallback) or excluding Phase-1 samples from the vote both hurt, confirming both components matter. Among free log-probability features, Total NLL falls below Uniform by conflating uncertainty with response length; length-aware Token Var and Max Token NLL beat Uniform but trail ANLL by 1–5 points, establishing length normalization as essential. Neither costly alternative surpasses the free signal: VCS loses every cell, and Vote Entropy ($M$7) ties on MATH500 but trails by 2–3% on Formal Logic because its binary entropy cannot rank uncertain questions finely. Threshold-exit variants—skipping questions deemed too hard or already confident—do not improve accuracy on average (best hard-threshold costs $M$8; best easy-threshold saves ~20% budget at $M$9), except for well-calibrated Qwen2.5-7B, where the easy-threshold saves 11.4% budget at only $B = N \times M$0. Plain UAB remains the default.

Limitations and open questions

The paper concedes four substantive limitations. First, the optimization maximizes a coverage surrogate (probability of at least one correct sample) while evaluation uses majority vote; the two share concave marginal structure and align empirically, but a formal bound relating them remains open, and extra samples on low-probability questions can in principle dilute a correct minority. Second, the Bernoulli model assumes i.i.d. samples, whereas same-prompt generations share systematic biases; the authors argue shared correlation preserves difficulty rankings uniformly across questions, but heterogeneous intra-question correlation could perturb the ordering when two questions have similar difficulty—an extreme case could invert it. Third, although $B = N \times M$1 generalizes across their benchmarks, new model families or domains may require re-validation. Fourth, extension to open-ended tasks where majority voting is ill-defined (e.g., summarization) is left unresolved. A practical constraint worth noting is that ANLL requires per-token log-probabilities unavailable from many commercial APIs, motivating the VCS substitute—which, however, underperforms ANLL in every ablation cell tested.

Conclusion

This paper casts cross-question sample allocation under a fixed inference budget as a separable concave knapsack solved exactly by marginal-greedy allocation, using ANLL from the first generation as a free difficulty signal. Across six models and five benchmarks, UAB consistently outperforms uniform self-consistency and heuristic baselines at matched budgets, with the largest gains (+3% average, +5% per-benchmark, up to 27% call savings at fixed accuracy) concentrated in the low-budget regime where inference cost matters most, at negligible wall-clock overhead and no auxiliary model. Its main theoretical contribution is a sensitivity bound showing the allocation degrades gracefully under imperfect probability calibration, which is essential given that ANLL is a ranking heuristic rather than a calibrated estimator.