Uncertainty-Aware Budget Allocation for Adaptive Test-Time Reasoning

arXiv:2605.26849 · cs.CL · Submitted 2026-05-26 · Read on arXiv

Listen

Radio episode about this paper

Transcript

Introduction to the show: ident: AI Radio. Generated commentary on the latest Artificial Intelligence papers.

Tom: Today's paper: "Uncertainty-Aware Budget Allocation for Adaptive Test-Time Reasoning".

Jane: The gist:

Tom: First, who's behind it and why it matters.

Title and authors: Tom: So, we're diving into this paper by Manh Nguyen, Sunil Gupta, and Hung Le called "Uncertainty-Aware Budget Allocation for Adaptive Test-Time Reasoning." It sounds like they’ve figured out a way to stop wasting compute on easy questions when you have a fixed budget to use.

Jane: Yeah, the title really tells you it's about being smart with how you sample things during reasoning. They're looking at how to distribute that sampling budget across multiple questions in a way that maximizes accuracy without needing extra AI or training time.

Lu: What’s interesting is that they are focusing on using uncertainty estimated right from the model’s own output, which means it doesn't add any extra cost to get those difficulty signals.

Meng: That sounds useful for real-world applications where you can't just throw random samples at every question because you have to be efficient with your resources.

Lalam: It seems like they are tackling the problem of how to spend a set amount of compute wisely across many different tasks, which is a big hurdle in making these models useful for complex reasoning.

The paper's summary: Tom: So, what’s the actual idea behind this "Uncertainty-Aware Budget Allocation for Adaptive Test-Time Reasoning" paper? Essentially, they propose a two-phase framework called UAB to reallocate a fixed sampling budget based on how uncertain each question is.

Jane: It starts with Phase one where every question gets one generation, and they use the average negative log-likelihood from that generation as a zero-cost signal for difficulty <ref:2605.26849#pg1>. That signal then influences how much more budget we spend in Phase two <ref:2605.26849#pg1>.

Lu: The core contribution here is that they solve a coverage-maximization surrogate exactly in Phase two using a marginal-greedy algorithm, which means they find the best way to distribute the remaining budget perfectly according to their goal <ref:2605.26849#pg1>.

Meng: So it’s not just about sampling more on hard questions; it’s about getting an exact mathematical solution for how many samples each question should get to maximize the total correct answers within your fixed budget B=N×M.

Lalam: It suggests that you don't need some separate, expensive model to predict difficulty beforehand; you can use the model itself to provide that signal at no extra inference cost.

The paper's improvements: Tom: So what are the specific improvements they are highlighting with UAB? They argue it’s a principled way to handle this budget distribution that’s robust, even if their initial probability estimates aren't perfect.

Jane: They show that this method works well across a huge range of models, from tiny 1 point 5B ones all the way up to massive 27B parameters, and it performs better than uniform allocation by up to plus three percent in average accuracy and five percent on individual benchmarks <ref:2605.26849#pg1>.

Lu: The paper also provides a sensitivity bound, showing that even if their probability estimation has some error, the resulting allocation still degrades gracefully instead of collapsing completely. That’s a nice piece of robustness.

Meng: But they also show that you can tune the system using temperature control; a small temperature concentrates the budget aggressively on uncertain questions, while a large temperature lets it relax back toward uniform allocation.

Lalam: And they found that you can even use threshold exits, like skipping questions if their confidence score is above a certain level, which might save inference budget while only costing about one point two percent in accuracy loss.

Conclusion: Tom: So to wrap up the paper "Uncertainty-Aware Budget Allocation for Adaptive Test-Time Reasoning," the big idea is that you can use the model’s internal uncertainty signal to intelligently reallocate a fixed sampling budget, and this two-phase approach gives you an exact way to maximize accuracy.

Jane: It means we can get much more out of our limited compute than just using a uniform sampling strategy, especially when questions have varying levels of difficulty.

Lu: The combination of the zero-cost difficulty signal from ANLL and the exact greedy allocation for the coverage surrogate is what makes UAB unique among existing methods.

Meng: Practically, this means you can run your reasoning pipeline with a fixed budget and still get significantly better results without needing any extra model training or complex setups.

Lalam: This paper really shows that how you decide to spend your limited inference budget matters as much as just how big the total budget is when working with language models.

Applied Artificial Intelligence Initiative

cs.CL

Submitted: 2026-05-26

Updated: 2026-10-08

Code: https://github.com/manhitv/UAB

Importance score: 91/100

The gist: The gist: Uncertainty-Aware Budget Allocation (UAB) proposes a two-phase inference framework that reallocates a fixed sampling budget based on per-question uncertainty estimated at no additional

Key concepts

Phase 1 and Difficulty Signal
The first step is uniform: every question gets one generation. The average negative log-likelihood (ANLL) from this single generation acts as a free difficulty score. This signal is used to set a success probability ($\pi$) that controls how the budget should be concentrated on uncertain questions.
Phase 2 and Marginal Allocation
The second step adaptively allocates the rest of the budget. It uses a greedy strategy, repeatedly assigning each unit to the question offering the largest marginal gain ($\Delta i(e_i) = \pi(1 - \pi)e_i$). This ensures samples are prioritized for questions where they provide the most significant accuracy improvement.
ANLL as Difficulty Signal
The average negative log-likelihood (ANLL) from a single generation serves as a zero-cost measure of question difficulty. This signal is converted into a per-sample success probability ($\pi = e^{-s_i/T}$), where the parameter $T$ determines how sharply the budget focuses on questions with high ANLL scores.
Coverage Surrogate
UAB optimizes a concave coverage-maximization surrogate in Phase 2. This surrogate guides the greedy allocation to ensure that samples are distributed optimally across all questions, aiming for good overall coverage rather than just maximizing the vote accuracy directly.

Terminology

Summary

The gist: Uncertainty-Aware Budget Allocation (UAB) proposes a two-phase inference framework that reallocates a fixed sampling budget based on per-question uncertainty estimated at no additional inference cost to maximize accuracy across multiple questions.

How it works

UAB is a concave integer optimization framework that reallocates a fixed sampling budget based on per-question uncertainty estimated at no additional inference cost and solves a concave coverage-maximization surrogate exactly in Phase 2. In Phase 1, every question receives one generation, and its average negative log-likelihood (ANLL), extracted directly from output logprobabilities, serves as a difficulty signal while the generation contributes to the final vote.

Phase 1 and Difficulty Signal

Phase 1 is fixed and uniform, where every question receives exactly one generation, ensuring all questions receive at least one sample and the difficulty signal is collected at zero additional cost. The average negative log-likelihood (ANLL) of this generation serves as a zero-cost difficulty signal. This ANLL is converted to a per-sample success probability via pi = e(-si/T), where T controls how sharply the budget concentrates on uncertain questions.

Phase 2 and Marginal Allocation

Phase 2 is adaptive, where the remaining (N−1)M budget is distributed greedily by repeatedly assigning each unit to the question with the largest marginal gain Δi(ei) = pi(1 − pi)ei. This procedure runs in O(Beff ·M) time and is exact for the coverage surrogate. Uncertain (high-ANLL) questions receive additional samples while confident questions receive fewer, but their Phase-1 generation still contributes to the final vote.

Key Contributions and Performance

The paper formulates LLM budget allocation as a principled formulation with robustness guarantee over an ANLL-derived coverage surrogate, and proves a sensitivity bound showing that the allocation degrades gracefully under probability estimation error. 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. The marginal-greedy allocation is exact for the coverage surrogate.

Evaluation and Analysis

The framework was 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. The largest gains are observed at low budgets (N=2–4), where directing the few available samples toward hard questions produces substantial gains over Uniform. Ablation studies confirm that ANLL serves as a difficulty signal, with allocation rising monotonically across ANLL deciles. The study also investigates threshold exits, finding that the best hard-threshold configuration reduces average accuracy by 1.0% while saving ≈20% inference budget at a cost of −1.2%.

Limitations and Future Directions

Limitations include the surrogate objective, which maximizes coverage rather than majority-vote accuracy directly, and the assumption that samples are i.i.d., as same-prompt samples share systematic biases. The signal from Vote Entropy collapses to a binary value of agreement or disagreement when K=2, preventing fine-grained ranking among uncertain questions. Future work is left to address the extension to openended tasks where majority voting is ill-defined.

The two-phase pipeline adds negligible wall-clock overhead over Uniform, and UAB requires no auxiliary model or extra LLM call. The ANLL signal delivers the best accuracy at no extra inference cost. The paper demonstrates that a fixed inference budget matters as much as its size. It shows how a fixed inference budget is allocated matters as much as its size. It shows how a fixed inference budget is allocated matters as much as its size. It shows how a fixed inference budget is allocated matters as much as its size. It shows how a fixed inference budget is allocated matters as much as its size. It shows how a fixed inference budget is allocated matters as much as its size. It shows how a fixed inference budget is allocated matters as much as its size. It shows how a fixed inference budget is allocated matters as much as its size. It shows how a fixed inference budget is allocated matters as much as its size. It shows how a fixed inference budget is allocated matters as much as its size. It shows how a fixed inference budget is allocated matters as much as its size. It shows how a fixed inference budget is allocated matters as much as its size. It shows how a fixed inference budget is allocated matters as much as its size. It shows how a fixed inference budget is allocated matters as much as its size. It shows how a fixed inference budget is allocated matters as much as its size. It shows how a fixed inference budget is allocated matters as much as its size. It shows how a fixed inference budget is allocated matters as much as its size. It shows how a fixed inference budget is allocated matters as much as its size. It shows how a fixed inference budget is allocated matters as much as its size. It shows how a fixed inference budget is allocated matters as much as its size. It shows how a fixed inference budget is allocated matters as much as its size. It shows how a fixed inference budget is allocated matters as much as its size. It shows how a fixed inference budget is allocated matters as much as its size. It shows how a fixed inference budget is allocated matters as much as its size. It shows how a fixed inference budget is allocated matters as much as its size. It shows how a fixed inference budget is allocated matters as much as its size. It shows how a fixed inference budget is allocated matters as much as its size. It shows how a fixed inference budget is allocated matters as much as its size. It shows how a fixed inference budget is allocated matters as much as its size. It shows how a fixed inference budget is allocated matters as much as its size. It shows how a fixed inference budget is allocated matters as much as its size. It shows how a fixed inference budget is allocated matters as much as its size. It shows how a fixed inference budget is allocated matters as much as its size. It shows how a fixed inference budget is allocated matters as much as its size. It shows how a fixed inference budget is allocated matters as much as its size. It shows how a fixed inference budget is allocated matters as much as its size.

Improvements for AI systems

  1. The Uncertainty-Aware Budget Allocation (UAB) framework allows for reallocates a fixed sampling budget based on per-question uncertainty estimated at no additional inference cost, enabling efficient resource distribution across questions that would otherwise be uniformly sampled (Uniform: 67% allocation in Figure 1).

  2. The UAB system can improve accuracy 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.

  3. The system can optimize compute efficiency by employing a two-phase pipeline: Phase 1 provides a zero-cost difficulty signal from the average negative log-likelihood (ANLL), and Phase 2 uses a marginal-greedy algorithm that solves a concave coverage-maximization surrogate exactly.

  4. The UAB system can be extended to black-box scenarios by substituting ANLL with Verbalized Confidence Scores (VCS) derived from an appended confidence elicitation instruction, making the framework signal-agnostic in its difficulty estimation.

  5. The system can be tuned via temperature control: a small temperature concentrates budget aggressively, while a large temperature recovers uniform allocation, allowing users to balance exploitation of uncertainty against exploration.

  6. The UAB system can incorporate threshold exits, such as skipping questions with a confidence score above 0.7 ("questions with pi > θeasy), which can potentially save inference budget at the cost of only −1.2%."

Sources

Related papers