Where Should a Reasoning Model Spend Its Thinking? A Proposed Uncertainty-Adaptive Test-Time Compute System

Abstract

Large language models can now trade extra inference-time computation for higher answer accuracy — sampling several reasoning chains, searching over intermediate steps, or scoring candidates with a learned verifier. Published evidence shows this "test-time compute" is often more effective per FLOP than growing model size [1,4], and that step-level verifiers substantially outperform outcome-only supervision at catching reasoning errors [3]. What remains open is *how* a system should decide, step by step, where to spend that extra compute rather than spreading it evenly or fixing it in advance. We propose the Uncertainty-Adaptive Reasoning Compute System (UARCS), an inference-time controller that estimates calibrated uncertainty at each reasoning step and allocates a fixed compute budget toward the steps most likely to determine correctness, formulated as a constrained-optimization problem solved with a greedy marginal-utility rule. We describe the architecture, its mathematical formulation, and a planned evaluation protocol on established reasoning benchmarks. No experiments under this protocol have been run; the outcomes discussed are hypotheses, not findings.

Keywords: large language models, test-time compute, reasoning systems, process reward models, uncertainty estimation, inference optimization.

1. Introduction

Chain-of-thought prompting showed that asking a large language model to write out intermediate steps before an answer measurably improves performance on arithmetic, symbolic, and commonsense reasoning tasks [1]. Self-consistency built on this by sampling many independent reasoning chains for the same question and taking the majority final answer, trading extra sampling compute for a further accuracy gain without any change to the underlying model [2]. Separately, step-level process reward models (PRMs) — verifiers trained to judge whether each intermediate step, not just the final answer, is correct — were shown to catch errors that outcome-only reward signals miss, and to guide search toward better solutions [3]. Most recently, both industrial reasoning systems and open replications built around reinforcement-learned chains of thought have made "spend more compute at inference time" a first-class design axis of large language model systems, alongside model size and training compute [4,5].

A recent systematic study of this axis found that the value of extra test-time compute is highly non-uniform: on easy questions a model already answers correctly with a single pass, so extra samples mostly waste compute, while on harder questions additional search or verification meaningfully changes the outcome — and an adaptive policy that spends compute where it helps can match or beat a much larger fixed compute budget spent uniformly [4]. That finding establishes the value of adaptivity at the level of a whole question. It does not, on its own, say how a system should decide moment-to-moment, *within* one reasoning trace, whether the step it just produced is trustworthy enough to move on, or uncertain enough to justify a verifier call or a branch into alternative continuations.

Research gap. Self-consistency samples a fixed number of full chains regardless of how confident any individual step is [2]. Process reward models are typically used to re-rank or filter already-generated candidate chains after the fact, not to decide, step by step and before generation continues, whether to keep spending compute on the current trace [3]. Question-level adaptive compute studies confirm that non-uniform allocation helps in aggregate, but do not specify a step-level mechanism for making that allocation decision online, under an explicit compute budget, using a single calibrated uncertainty signal [4]. We are not aware of a system that couples a step-level uncertainty estimate directly to a compute-budgeted branching decision inside one inference pass — the gap this proposal addresses.

Research question. Can a step-level, calibrated uncertainty signal be used to allocate a fixed test-time compute budget across a reasoning trajectory more efficiently than uniform sampling or a static verifier threshold, measured as accuracy achieved per unit of inference compute?

Hypothesis. We hypothesize that an inference-time controller which spends its budget on steps flagged as uncertain by a combined entropy-and-verifier-disagreement signal, and skips verification on low-uncertainty steps, will reach a given accuracy level at lower total compute than self-consistency sampling or best-of-N re-ranking with a fixed budget, because it concentrates search exactly where the trajectory is most likely to still be corrigible.

2. AI System & Proposed Approach

UARCS sits between a base reasoning model and the user, wrapping ordinary autoregressive decoding with four cooperating components: a step segmenter, an uncertainty estimator, a branching controller, and a lightweight process verifier ensemble. The base model itself is not modified; UARCS is an inference-time control layer, in the same spirit as self-consistency and best-of-N re-ranking [2,3], but operating at step granularity and under an explicit, shared compute budget rather than a fixed sample count.

Step segmenter. The generated chain is split into discrete reasoning steps s_1, ..., s_n using the model's own step-boundary markers (line breaks, numbered steps, or a lightweight learned boundary classifier when the base model does not emit explicit markers) — the same segmentation unit used to train step-level verifiers [3].

Uncertainty estimator. For step s_t, token-level predictive entropy is averaged over the step:

H̄_t = (1 / |s_t|) · Σ_{k=1}^{|s_t|} ( − Σ_w p(w | context, x_{<k}) · log p(w | context, x_{<k}) )

where p(w | context, x_{<k}) is the model's next-token distribution at position k of step t and the inner sum runs over the vocabulary. A small ensemble of m lightweight verifier heads independently scores the step's correctness probability v_1(s_t), ..., v_m(s_t); their disagreement is measured as the sample variance Var(v(s_t)). The two signals are combined into one step uncertainty score:

U_t = α · H̄_t + (1 − α) · Var(v(s_t))

with α ∈ [0,1] a tunable weight fit on a held-out set. U_t is intentionally cheap: entropy is read off the decoding pass at no extra cost, and the verifier ensemble is a small model relative to the base reasoner, following the same asymmetry between generator and verifier used in process-reward-model work [3].

Figure 1 shows the resulting inference-time control loop.

Figure 1. Proposed UARCS control loop — a proposed system diagram, not a deployed pipeline.

Branching controller and budget scheduler. Given U_t and the remaining compute budget B_remaining, the controller chooses one of three actions a ∈ {continue, verify, branch}: continue proceeds to the next step at no extra cost; verify spends a small fixed cost c_v to query the full verifier ensemble and, if confidence is still low, revise the step; branch spends a larger cost c_b to sample k alternative continuations of the step and keep the one the verifier ensemble scores highest. Framing this as a compute-constrained optimization over the whole trajectory,

maximize Σ_{t=1}^{n} E[correct_t | a_t]   subject to   Σ_{t=1}^{n} cost(a_t) ≤ C

where C is the total inference compute budget for the trajectory, the Lagrangian relaxation of this constraint yields a greedy per-step decision rule: take the action with the highest marginal utility per unit cost,

a_t* = argmax_a   [ ΔP_correct(a, U_t) / cost(a) ] ,   taken only if   ΔP_correct(a, U_t) / cost(a) > λ

where λ is the Lagrange multiplier (shadow price of compute) tuned so that expected total spend matches C, and ΔP_correct(a, U_t) is the estimated gain in trajectory correctness probability from taking action a at uncertainty level U_t, itself modeled by a small calibration function fit on held-out data relating U_t to observed error rates. This is the same knapsack-style greedy logic used for constrained resource allocation problems, applied here to reasoning steps instead of items: cheap, high-uncertainty steps are verified or branched first, and the budget is exhausted on the steps most likely to still be wrong.

3. Methodology

Architecture. UARCS wraps an existing instruction-tuned reasoning model (frozen, no fine-tuning of the base weights is proposed) with the four components above. The verifier ensemble is proposed as m = 3–5 small classifier heads sharing a distilled encoder, kept intentionally light relative to the base model, following the generator/verifier asymmetry that made process-reward-model verification computationally practical in prior work [3].

Training. The verifier ensemble is trained with weak step-level labels derived the same way process reward models have been trained without exhaustive human step annotation: a step is labeled correct if trajectories continuing from it agree with the majority self-consistency answer at a higher rate than trajectories that diverge from it [2,3]. The calibration function mapping U_t to an error-rate estimate is fit by isotonic regression on a held-out validation split, and the Lagrange multiplier λ is tuned by binary search so that mean realized compute spend across the validation set matches the target budget C.

Datasets. We propose training and evaluating on GSM8K and MATH for quantitative reasoning, HumanEval and MBPP for code reasoning, and a held-out split of ARC-style multi-step logic problems, mirroring the benchmark families used to establish self-consistency, process-reward verification, and test-time compute scaling results [2,3,4].

Baselines. Single-pass greedy decoding (no extra compute); self-consistency at matched total sample budget [2]; best-of-N with the same verifier ensemble applied only after full generation, at matched total compute [3]; and a fixed, non-adaptive test-time compute allocation that spends the same average budget uniformly across all questions, representing the non-adaptive alternative to the question-level adaptivity result reported in prior work [4].

Evaluation. The primary proposed metric is accuracy achieved as a function of total inference compute (FLOPs or generated-token count), reported as an accuracy-versus-compute curve per benchmark, allowing UARCS to be compared against each baseline at matched budget rather than at a single operating point. Secondary proposed metrics are the expected calibration error of U_t against realized step error rates, and tokens spent per solved problem. Figure 2 summarizes the proposed data flow from raw problem sets through weak-label generation to the trained system and its evaluation.

Figure 2. Proposed data flow from benchmark problems to a trained, evaluated system.

4. Results & Discussion

No experiments have been conducted under the protocol above, and this section makes no accuracy, compute-savings, or calibration claims for UARCS itself. What can be reported is the published evidence that motivates the design. Self-consistency's gains from multi-chain sampling over single-pass decoding are well established across arithmetic and commonsense benchmarks [2]. Process reward models trained with step-level supervision have been shown to substantially outperform outcome-only reward models at locating the first incorrect step in a flawed solution, which is the capability UARCS's verifier ensemble is proposed to reuse for step-level uncertainty rather than post-hoc filtering [3]. Most directly, the finding that adaptive, question-difficulty-aware test-time compute allocation can match a several-times-larger fixed compute budget spent uniformly is the empirical basis for expecting a step-level adaptive policy to help, since it demonstrates that non-uniform allocation is where the available gain lives, not merely that more compute helps [4]. Large-scale reinforcement-learning-trained reasoning systems have further shown that letting a model's own generation length and structure vary with task difficulty, rather than fixing it, is compatible with strong benchmark performance, reinforcing that variable-depth reasoning is a viable system design rather than a workaround [5].

Taken together, this evidence supports the plausibility of the hypothesis without validating it: question-level adaptivity works [4], step-level verification signal is informative [3], and multi-sample aggregation helps when spent [2] — but no published result measures whether combining these three signals into one step-level, budget-constrained controller outperforms spending the same total budget via self-consistency or post-hoc best-of-N re-ranking. That comparison is the proposed experimental design in Section 3; until it is run, any specific percentage improvement would be speculation. If the hypothesis holds, the expected impact is lower inference cost for a given accuracy target in latency- and cost-constrained deployments such as on-device assistants or high-volume API serving, recovering further savings on the "easy majority, hard minority" step distribution typical of multi-step reasoning traces. If it does not hold — for instance, if entropy and verifier disagreement correlate poorly with actual step error, or the calibration function fails to transfer across benchmark families — that negative result would itself clarify whether step-level uncertainty is a reliable currency for compute allocation at all, which is useful independently of UARCS.

5. Conclusion & Future Research

Findings. Published work establishes that test-time compute is a distinct, tunable axis of large language model system performance, that step-level process verification outperforms outcome-only supervision, and that non-uniform allocation of test-time compute across questions of varying difficulty can outperform a uniform fixed budget [1–5]. No published result, to our knowledge, specifies a step-level, budget-constrained controller that ties a calibrated uncertainty signal directly to an online branching decision within a single reasoning trace.

Proposed solution. UARCS is offered as a concrete, mathematically specified answer to that gap: a step segmenter, an entropy-plus-verifier-disagreement uncertainty score, and a Lagrangian-derived greedy budget-allocation rule that spends compute where the marginal probability of correcting an error is highest per unit cost.

Limitations. The proposal assumes the weak self-consistency-derived step labels used to train the verifier ensemble are informative enough to teach step-level correctness, which is itself an assumption inherited from, not proven beyond, prior process-reward-model work [3]. The calibration function mapping uncertainty to error probability is likely benchmark- and model-specific and may need re-fitting per deployment. The greedy Lagrangian rule is a relaxation of the true combinatorial budget-allocation problem and is not guaranteed optimal. Verifier ensemble cost, while designed to be small relative to the base model, is not zero, and on very short reasoning traces the fixed overhead of segmentation and scoring could erode any compute savings.

Next steps. The immediate next step is implementing the four components against an open instruction-tuned reasoning model and running the accuracy-versus-compute comparison against self-consistency and best-of-N baselines described in Section 3, before any broader claim is made. Beyond that initial validation, natural extensions include testing whether the same uncertainty signal transfers to tool-using agentic settings where a "step" is a tool call rather than a reasoning sentence, studying the robustness of the uncertainty estimator against adversarially confident-sounding but incorrect steps, and examining whether a single calibration function can serve multiple base models or whether per-model recalibration is required in practice.

References

1. Wei, J. et al. (2022). Chain-of-Thought Prompting Elicits Reasoning in Large Language Models. NeurIPS.
2. Wang, X. et al. (2023). Self-Consistency Improves Chain of Thought Reasoning in Language Models. ICLR.
3. Lightman, H. et al. (2024). Let's Verify Step by Step. ICLR.
4. Snell, C. et al. (2024). Scaling LLM Test-Time Compute Optimally can be More Effective than Scaling Model Parameters. arXiv preprint.
5. DeepSeek-AI (2025). DeepSeek-R1: Incentivizing Reasoning Capability in LLMs via Reinforcement Learning. arXiv preprint.

This is a ScienceTrace research proposal. It presents a candidate system, its mathematical formulation, and a planned evaluation protocol. It does not report trained models, measured accuracy, or empirical validation of any kind.