# Uncertainty-Aware Budget Allocation for Adaptive Test-Time Reasoning

arXiv: https://arxiv.org/abs/2605.26849 · HTML: https://mlsanity.com/paper/2605.26849/

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Placed in Test-time compute and efficient thinking, one of 17 directions in LLM reasoning.
- Builds on 18 of its 33 references that are in the map.
- First posted on arXiv 26 May 2026.

## Placed in

- LLM reasoning / [Test-time compute and efficient thinking](https://mlsanity.com/lr/test-time-compute-and-efficient-thinking/)

## Builds on (most cited first, 18 of 18)

- [Training Verifiers to Solve Math Word Problems](https://mlsanity.com/paper/2110.14168/) (2021)
- [Measuring Massive Multitask Language Understanding](https://mlsanity.com/paper/2009.03300/) (2020)
- [Self-Consistency Improves Chain of Thought Reasoning in Language Models](https://mlsanity.com/paper/2203.11171/) (2022)
- [Measuring Mathematical Problem Solving With the MATH Dataset](https://mlsanity.com/paper/2103.03874/) (2021)
- [Let's Verify Step by Step](https://mlsanity.com/paper/2305.20050/) (2023)
- [Training a Helpful and Harmless Assistant with Reinforcement Learning from Human Feedback](https://mlsanity.com/paper/2204.05862/) (2022)
- [GPQA: A Graduate-Level Google-Proof Q&A Benchmark](https://mlsanity.com/paper/2311.12022/) (2023)
- [Scaling LLM Test-Time Compute Optimally can be More Effective than Scaling Model Parameters](https://mlsanity.com/paper/2408.03314/) (2024)
- [Large Language Monkeys: Scaling Inference Compute with Repeated Sampling](https://mlsanity.com/paper/2407.21787/) (2024)
- [Inference Scaling Laws: An Empirical Analysis of Compute-Optimal Inference for Problem-Solving with Language Models](https://mlsanity.com/paper/2408.00724/) (2024)
- [Debate or Vote: Which Yields Better Decisions in Multi-Agent Large Language Models?](https://mlsanity.com/paper/2508.17536/) (2025)
- [SelfBudgeter: Adaptive Token Allocation for Efficient LLM Reasoning](https://mlsanity.com/paper/2505.11274/) (2025)
- [MUR: Momentum Uncertainty guided Reasoning for Large Language Models](https://mlsanity.com/paper/2507.14958/) (2025)
- [Reliability-Aware Adaptive Self-Consistency for Efficient Sampling in LLM Reasoning](https://mlsanity.com/paper/2601.02970/) (2026)
- [Optimal Self-Consistency for Efficient Reasoning with Large Language Models](https://mlsanity.com/paper/2511.12309/) (2025)
- [Hear Both Sides: Efficient Multi-Agent Debate via Diversity-Aware Message Retention](https://mlsanity.com/paper/2603.20640/) (2026)
- [Seer Self-Consistency: Advance Budget Estimation for Adaptive Test-Time Scaling](https://mlsanity.com/paper/2511.09345/) (2025)
- [Breaking the Pre-Sampling Barrier: Activation-Informed Difficulty-Aware Self-Consistency](https://mlsanity.com/paper/2602.09438/) (2026)

## Abstract (arXiv)

Sampling multiple responses improves language model reasoning, but uniform compute allocation is inefficient because 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 using uncertainty estimated from the initial samples themselves. In Phase-1, every question receives a small fixed number of generations. Their answer disagreement, measured by vote entropy, provides a difficulty signal while these generations contribute to the final vote. In Phase-2, the remaining budget is allocated by a marginal-greedy algorithm that optimally solves a concave coverage-maximization surrogate, concentrating samples on questions whose initial answers disagree. Across five open-weight models (1.5B--27B parameters) and five reasoning benchmarks of varying difficulty, UAB improves average accuracy by +2.3% over uniform allocation, and achieves gains of up to +5.5% on individual benchmarks, with the largest gains in low-resource settings. Moreover, UAB meets the target budget exactly and requires no auxiliary model or additional LLM calls. Code is publicly available at https://github.com/manhitv/UAB.