Uncertainty-Aware Budget Allocation for Adaptive Test-Time Reasoning

• URL: http://arxiv.org/abs/2605.26849v1
• Date: Tue, 26 May 2026 11:06:58 GMT
• Title: Uncertainty-Aware Budget Allocation for Adaptive Test-Time Reasoning
• Abstract summary: Uncertainty-Aware Budget Allocation (UAB) is a concave integer optimization framework that reallocates a fixed sampling budget based on per-question uncertainty estimated at no additional inference cost.<n> 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.
• Score: 13.41454380481593
• License: http://arxiv.org/licenses/nonexclusive-distrib/1.0/
• Abstract: Sampling multiple responses improves language model 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.

Related papers

• Robust Average-Reward Markov Decision Processes: Minimax-Optimal Learning via Plug-in Reductions [51.50375419691955]
  Distributionally robust Markov decision processes provide a principled framework for sequential decision making under model uncertainty.<n>We study how many samples are necessary and sufficient to learn an $varepsilon$-optimal robust policy under the average-reward criterion.
  arXiv  Detail & Related papers  (2026-08-06T19:49:48Z)
• Refining Over Resampling: Test-Time Self-Correction for LLM Reasoning [3.8191765711200065]
  We propose a verifier-free breadth--depth refinement framework that uses test-time compute to both explore and improve candidate solutions.<n>Our method consistently improves over greedy decoding, majority voting, verifier-based best-of-$N$, beam search, and look decoding across multiple open-weight models.
  arXiv  Detail & Related papers  (2026-08-06T06:38:37Z)
• Instance-Optimal Estimation with Multiple LLM Judges on a Budget [84.31744861038106]
  We formalize this question as *budgeted heteroskedastic multi-judge estimation*.<n>Given $K$ prompt-response pairs, $J$ judges with known costs, and unknown query-judge variances, the goal is to estimate a bounded score vector while minimizing an $ell_p$-error.<n>We prove that EST-IVWE matches the oracle IVWE rate up to lower-order terms in the budget.
  arXiv  Detail & Related papers  (2026-05-22T08:26:08Z)
• Adaptive Test-Time Compute Allocation for Reasoning LLMs via Constrained Policy Optimization [18.737087162461563]
  Test-time compute scaling has become a powerful lever for improving large language model performance.<n>Yet deploying these techniques under finite inference budgets requires a decision that current systems largely ignore.<n>We formalize this as a constrained optimization problem (maximize expected accuracy subject to an average compute budget) and solve it with a two-stage Solve-then-Learn pipeline.
  arXiv  Detail & Related papers  (2026-04-16T10:39:22Z)
• ODAR: Principled Adaptive Routing for LLM Reasoning via Active Inference [60.958331943869126]
  ODAR-Expert is an adaptive routing framework that optimize the accuracy-efficiency trade-off via principled resource allocation.<n>We show strong and consistent gains, including 98.2% accuracy on MATH and 54.8% on Humanity's Last Exam.
  arXiv  Detail & Related papers  (2026-02-27T05:22:01Z)
• Efficient Evaluation of LLM Performance with Statistical Guarantees [11.703733256169214]
  We propose Factorized Active Querying (FAQ) for benchmarking large language models.<n>FAQ adaptively selects questions using a hybrid variance-reduction/active-learning sampling policy.<n>FAQ delivers up to $5times$ effective sample size gains over strong baselines on two benchmark suites.
  arXiv  Detail & Related papers  (2026-01-28T04:59:20Z)
• CarBoN: Calibrated Best-of-N Sampling Improves Test-time Reasoning [62.56541355300587]
  We introduce a general test-time calibration framework that adaptively modifies the model toward high-reward reasoning paths.<n>Within this framework, we propose CarBoN, a two-phase method that first explores the solution space and then learns a calibration of the logits.<n>Experiments on MATH-500 and AIME-2024 show that CarBoN improves efficiency, with up to $4times$ fewer rollouts to reach the same accuracy.
  arXiv  Detail & Related papers  (2025-10-17T14:04:37Z)
• Additive Distributionally Robust Ranking and Selection [0.8283940114367679]
  We propose a simple additive allocation (AA) procedure that aims to exclusively sample the $k + m - 1$ previously hypothesized critical scenarios.<n>We prove that AA is consistent and, surprisingly, achieves additivity in the strongest sense.<n>Results offer new and counterintuitive insights into the additive structure of DRR&S.
  arXiv  Detail & Related papers  (2025-09-07T17:36:29Z)
• Steering LLM Thinking with Budget Guidance [48.65894557568655]
  Budget guidance is a method for steering the reasoning process of LLMs toward a target budget without requiring any fine-tuning.<n>Our approach introduces a lightweight predictor that models a Gamma distribution over the remaining thinking length.<n>This signal is then used to guide generation in a soft, token-level manner, ensuring that the overall reasoning trace adheres to the specified thinking budget.
  arXiv  Detail & Related papers  (2025-06-16T17:57:05Z)
• Asymptotically Optimal Linear Best Feasible Arm Identification with Fixed Budget [55.938644481736446]
  We introduce a novel algorithm for best feasible arm identification that guarantees an exponential decay in the error probability.<n>We validate our algorithm through comprehensive empirical evaluations across various problem instances with different levels of complexity.
  arXiv  Detail & Related papers  (2025-06-03T02:56:26Z)
• Self-Evaluation Guided Beam Search for Reasoning [61.523627290397556]
  We introduce a stepwise self-evaluation mechanism to guide and calibrate the reasoning process of Large Language Model (LLM) We propose a decoding algorithm integrating the self-evaluation guidance via beam search. Our approach surpasses the corresponding Codex-backboned baselines in few-shot accuracy by $6.34%$, $9.56%$, and $5.46%$ on the GSM8K, AQuA, and StrategyQA.
  arXiv  Detail & Related papers  (2023-05-01T02:37:59Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.