---
title: Two-Stage Inference-Time Budget Control
url: https://www.emergentmind.com/topics/two-stage-inference-time-budget-control
type: topic
---

# Two-Stage Inference-Time Budget Control

Two-stage inference-time budget control denotes a family of methods that separates budgeted inference into an upstream assessment, screening, or planning step and a downstream execution step that spends a constrained resource more selectively. In the recent literature, the controlled resource may be reasoning tokens, total generation length, sampled reasoning paths, wall-clock latency, executed Transformer layers, model capacity by phase, or the number of expensive response-model calls. The topic is heterogeneous: some systems implement a literal two-part runtime pipeline, such as predict-then-actuate latency control or cheap-answer-then-escalate serving, whereas others are “two-stage” only in training and should not be misread as two-phase inference algorithms [2512.21859][2602.01237][2604.27536][2606.09514][2508.17196].

## 1. Resource models and formal objectives

A first unifying feature is that the literature makes the controlled budget explicit. In token-budgeted reasoning, the controlled quantity is often the length of the intermediate reasoning process. BudgetThinker formulates budget-aware reasoning for chain-of-thought generation with a target budget \(B\), where the model should adapt its reasoning so that it completes within that budget while preserving as much task accuracy as possible; the budget is enforced over the generated reasoning segment before the final answer [2508.17196]. SelfBudgeter uses an explicit textual interface,
```text
<budget>an integer, token budget</budget><solution>response</solution>
```
so that the model first emits a token budget and then a solution, treating output length as a proxy for inference cost and user wait time [2505.11274]. BET uses a related but continuous interface,
```text
<predict> Solvability: [0,1] Budget: [0,1] </predict>
<think> ... reasoning trace ... </think>
\boxed{final answer}
```
where the budget is measured in reasoning tokens inside the `<think>...</think>` block, and the maximum completion length is \(L_{\max}=16{,}384\) [2605.11625].

Other papers formalize different resources. TimeBill treats the resource as end-to-end wall-clock latency per request and writes the time-budgeted inference problem as maximizing a response-performance metric \(\mathcal M(\hat{\mathbf y}(\theta), \mathbf y)\) subject to \(t_{\mathrm{e2e}}(x,\theta)\le T\) and \(N\le N_{\max}\), where \(T\) is the request deadline and \(\theta\) is instantiated as a KV-cache eviction ratio \(\alpha\) [2512.21859]. Adaptive Test-Time Compute Allocation instead assumes a discrete action set \(\mathcal B=\{b_1,\ldots,b_K\}\) and optimizes expected accuracy under an average compute budget,
\[
V(B)=\max_{\pi\in\Pi} E_{x\sim D}[Acc(x,\pi(x))]\quad \text{s.t.}\quad E_{x\sim D}[C(\pi(x))]\le B,
\]
so the budget is global and population-level rather than per request [2604.14853]. ROI-Reasoning moves one step further and allocates a hard total token budget across an ordered set of multiple problems,
\[
\max_{\pi_\theta}\mathbb E\!\left[\sum_{i=1}^N r(\pi_\theta(s_i,x_i))\right]
\quad \text{s.t.}\quad
\sum_{i=1}^N c(\pi_\theta(s_i,x_i))\le B,
\]
making early over-spending directly detrimental to later tasks [2601.03822].

Several systems budget internal compute rather than output length. BUDDY uses a normalized depth budget \(b\in(0,1]\), converted to an integer number of executed middle layers by \(k=\operatorname{round}(b\cdot L)-2\), and always executes the first and last Transformer blocks [2606.09514]. DDC jointly budgets sampling width and reasoning depth, minimizing total generated tokens subject to a target reliability constraint on the final consensus answer [2605.15100]. Star Elastic controls compute by switching among nested submodels between the explicit thinking phase and answering phase, so budget becomes phase-specific model capacity rather than a single monolithic scale factor [2605.07182].

This diversity suggests that “budget control” is not tied to one metric. The common structure is selective resource allocation under explicit constraints, but the resource may be tokens, latency, paths, layers, experts, or model scale.

## 2. Canonical runtime architectures

Across the literature, several recurrent two-stage architectures appear. One pattern is **predict-then-actuate**. TimeBill first predicts latent workload through a Response Length Predictor and an Execution Time Estimator, then chooses the smallest KV-cache eviction ratio \(\alpha\) whose predicted worst-case latency satisfies the time budget [2512.21859]. Predictive Scheduling first runs a lightweight predictor, either a hidden-state MLP or a LoRA-based classifier, and then allocates a fixed total token budget across queries by a greedy allocator or by exhaustive search over difficulty-tier budgets [2602.01237]. Adaptive Test-Time Compute Allocation first solves the constrained allocation problem offline by Lagrangian relaxation to obtain oracle budget labels, then trains a lightweight classifier to predict those oracle actions from cheap features at deployment [2604.14853]. Veroic first produces a cheap default response \(y_t^{(0)}\), then uses verifiable observations from \((x_t,y_t^{(0)})\) to decide whether to accept that output or trigger a higher-cost inference pathway \(y_t^{(1)}\) [2604.27536].

A second pattern is **internal pre-reasoning planning** within one autoregressive pass. SelfBudgeter makes the model emit `<budget>...</budget>` before `<solution>...</solution>`, so the model first predicts the reasoning budget and then spends it [2505.11274]. BET similarly emits a `<predict>` block with solvability and budget before the `<think>` block, and trains explicit “short solve,” “nice fold,” and “hero call” behaviors [2605.11625]. ROI-Reasoning forces pre-computation meta-cognition via `<predicted_level>Level-k</predicted_level>` tags, where Level-3 means “too difficult, skip reasoning and answer \(\boxed{\text{NA}}\)” [2601.03822]. These methods do not separate planning and execution into different models, but they still implement a genuine two-part runtime behavior: estimate first, reason second.

A third pattern is **phase-specific execution**. Star Elastic preserves an explicit thinking phase followed by an answering phase and allows different nested submodels in each phase, including \(\mathcal M_S\to \mathcal M_L\), \(\mathcal M_L\to \mathcal M_L\), \(\mathcal M_S\to \mathcal M_S\), and \(\mathcal M_L\to \mathcal M_S\), with the paper reporting that \(\mathcal M_S\to \mathcal M_L\) is the best default [2605.07182].

Important exceptions clarify the terminology. BudgetThinker is frequently grouped with two-stage budget control, but its “two-stage” label refers to supervised fine-tuning followed by reinforcement learning; inference itself is a single decoding process with periodic control-token insertion [2508.17196]. ORBIT is likewise multi-stage in training and exposes explicit Low, Mid, High, and Xhigh reasoning modes at inference time, but it does not learn the stage-1 mode selector from input [2601.08310].

## 3. Token-budgeted reasoning and controllable chain-of-thought

BudgetThinker is a direct attempt to make a reasoning model obey a user-specified token budget during inference. Its central mechanism is a fixed set of control tokens \(\mathcal C=\{c_1,\dots,c_K\}\) that are injected deterministically during decoding at budget-fraction milestones, with \(K=8\) in the reported setup. The backbone model conditions on these inserted tokens exactly as on normal prior context. Training uses a two-stage pipeline: supervised fine-tuning on a 41k-example reasoning dataset, followed by GRPO with a composite reward combining correctness, format, and a length-aware term that penalizes overshooting much more strongly than undershooting, using \(\gamma=16\) when \(|y|>B\). The RL curriculum decreases budgets through \(6000\rightarrow 4000\rightarrow 3000\rightarrow 2000\), then mixes budgets sampled from \(\{6000,4000,3000,2000\}\). In evaluation on MATH-500, AMC 2023, and AIME 2024, the paper reports that BudgetThinker improves accuracy by an average of 4.9% across tested budgets in the abstract, and by 4.2% over original models and 5.7% over ThinkPrune on average for MATH-500 and AMC 2023 in the main text, while also improving budget following ratio and budget utilization ratio [2508.17196].

SelfBudgeter uses a more explicit budget-first protocol. Stage 1 teaches the model to output a numeric budget before the solution; Stage 2 applies budget-guided GRPO with a budget penalty and a Precise Budget Control Reward. In autonomous mode, the model predicts its own budget; in controlled mode, the user can pre-fill the `<budget>` field directly. The paper reports that SelfBudgeter achieves up to 74.47% response length compression on MATH while maintaining a 2.16-point accuracy drop relative to the baseline, and on GSM8K a best variant improves accuracy from 78.32% to 81.50% while reducing average tokens from 1737.92 to 662.08 [2505.11274].

BET generalizes token budgeting beyond length compression by tying budget allocation to policy-dependent solvability. It estimates current-policy solvability by Monte Carlo rollouts, defines an efficient solution cost from the shortest correct trajectories, and learns three named behaviors: short solve, nice fold, and hero call. The fold action is explicit, via `\boxed{<Unsolvable>}`, and is rewarded only when rollout-derived solvability is effectively zero. Across seven benchmarks and three base models, BET reduces reasoning tokens by about 55% on average while improving overall performance, and it transfers zero-shot from mathematical reasoning to scientific QA and logical reasoning [2605.11625].

This line of work shows two distinct conceptions of inference-time budget control. One conception prioritizes **budget adherence**—BudgetThinker and SelfBudgeter are explicit about matching or approximately matching a target budget. The other prioritizes **computational return on investment**—BET and ROI-Reasoning emphasize deciding whether a problem is worth solving at all under a shared token cap [2601.03822]. Budget-Aware Anytime Reasoning is adjacent but importantly different: it does not learn an online controller, but it defines the Anytime Index,
\[
\text{AnytimeIndex}=
\frac{\sum_{t=1}^{T-1} \frac{(Q^*_t+Q^*_{t+1})}{2}(b_{t+1}-b_t)}
{(b_T-b_1)Q_{\max}},
\]
and uses LLM-synthesized preference pairs at fixed token budgets to improve the quality of partial solutions under interruption [2601.11038].

## 4. Latency, depth, routing, and phase-specific model selection

TimeBill exemplifies the most explicit two-stage latency controller in the literature. Stage 1 consists of an SLM-based Response Length Predictor and a workload-guided analytical Execution Time Estimator. The predictor maps a prompt to a bucketed response-length estimate \(\hat N=\min(N_{\max},\mathrm{Predict}(x)B)\), and the ETE models prefill and decode cost analytically, then inflates the prediction to a worst-case estimate \(\hat N_W=\min(k\hat N,N_{\max})\). Stage 2 solves for the smallest eviction ratio \(\alpha\) that meets the request deadline \(T\). The paper reports that the RLP with 512 buckets achieves MAE 42.71, RMSE 78.13, and \(R^2=0.723\); the ETE reaches Mean Absolute Percentage Error of 1.22% for prefill and 1.69% for decoding-step estimation; and TimeBill achieves the highest average response performance scores among tested approaches while attaining a similar task completion rate as \(\alpha=95\%\) [2512.21859].

BUDDY brings the same predict-then-actuate idea inside the Transformer. Its optional Budget Predictor chooses an input-dependent compute level when the user does not specify one, and the Decision Module then scores middle layers and deterministically executes the top-\(k\) subset consistent with the budget. The routing rule is
\[
P(s',b)=\operatorname{Top}\text{-}K(\operatorname{Softmax}(s'),\,k=\operatorname{round}(bL)-2),
\]
and the selected mask can be recomputed at every decoding step using first-layer KV cache information as a lightweight global context source. On Llama3-8B, the paper reports that a single multi-budget model retains approximately 99.9% of original accuracy at 12.5% sparsity, 90.8% at 25%, 80.7% at 37.5%, and 71.3% at 50%, while also supporting strict budget control and decode-time rerouting [2606.09514].

Star Elastic addresses a different axis: model capacity by reasoning phase. It turns one parent model into a nested family of submodels and preserves the standard reasoning protocol of a token-bounded thinking phase followed by an answering phase. The system can therefore run, for example, a smaller model during the long token-heavy thinking phase and a larger model during the short accuracy-critical answering phase. The paper reports that the best default pairing is \(23\text{B}\to 30\text{B}\), and that dynamic per-phase model selection yields up to 16% higher accuracy and 1.9× lower latency. The reported experiments still recompute cache states when switching models, so the measured timings already include switching overhead [2605.07182].

Budget-Aware Value Tree extends budget control from pure decoding to tool-augmented search. It is not an explicit hard two-stage scheduler, but it naturally decomposes into early broad exploration and late focused exploitation through the remaining-budget ratio
\[
r_t=\min\!\left(\frac{b_{\mathrm{tool},t}}{B_{\mathrm{tool}}},\frac{b_{\mathrm{token},t}}{B_{\mathrm{token}}}\right),
\qquad
\alpha_t=\frac{1}{r_t},
\qquad
\mathbb P(n_i)\propto V(n_i)^{\alpha_t}.
\]
As \(r_t\) falls, the node-selection distribution sharpens from diffuse sampling toward near-greedy exploitation, and a budget backstop forces answer generation when resources are nearly exhausted [2603.12634].

## 5. Shared-budget allocation, escalation, and deferral

Some of the clearest two-stage controllers allocate a shared budget across many units rather than controlling one request in isolation. Predictive Scheduling assumes a fixed total token budget over a batch of GSM8K queries. Stage 1 predicts either a full early-stopping probability vector or a coarse difficulty class before any full reasoning trace is generated. Stage 2 then allocates 16-token windows greedily across queries, or solves a discrete class-level allocation problem over easy, medium, and hard budgets. The paper reports up to 7.9 percentage points of absolute accuracy gain over uniform budgeting at identical token cost, and finds that middle transformer layers 12–17 are most informative for hidden-state prediction, with layer 16 reaching Pearson correlation 0.742 [2602.01237].

Adaptive Test-Time Compute Allocation formalizes the same idea more abstractly. In the Solve stage, Lagrangian relaxation yields the oracle action
\[
b^\star(x;\lambda)=\arg\max_{b\in\mathcal B}[Acc(x,b)-\lambda C(b)],
\]
and a binary search over \(\lambda\) targets the desired average budget through the monotone oracle-induced cost curve. In the Learn stage, a lightweight classifier predicts those oracle actions from 16 cheap features. On MATH and GSM8K with \(\mathcal B=\{1,2,4,8,16\}\) in self-consistency experiments, the method achieves up to 12.8% relative accuracy improvement on MATH under matched budget constraints and over 91% imitation accuracy while closely tracking the oracle upper bound [2604.14853].

ROI-Reasoning applies shared-budget control to ordered multi-problem reasoning. It treats a three-problem “test paper” under a strict global token cap as an Ordered Stochastic Multiple-Choice Knapsack Problem, uses Meta-Cognitive Fine-Tuning to emit `<predicted_level>Level-k</predicted_level>` tags before reasoning, and then uses Rationality-Aware Reinforcement Learning to optimize solve-or-skip decisions and reasoning length under the hard cap. Under the Hard/512 setting, the paper reports Score 0.93 for MFT+RARL versus 0.81 for MFT and 0.15 for the base model, with corresponding Regret values of 0.16, 0.11, and 2.73 [2601.03822].

Veroic treats repeated cheap-answer-then-escalate decisions as a long-horizon partially observable control problem. The latent state is whether the default response is reliable, the belief state is updated from verifiable hard and soft signals, and the policy decides between accepting the default output and triggering a stronger inference pathway. The discounted budget constraint is
\[
\mathbb E_\pi\!\left[\sum_{t\ge0}\gamma^t C(a_t)\right]\le B.
\]
In CMIS settings such as LLaMA-3.1-8B \(\rightarrow\) LLaMA-3.1-70B, Veroic improves task quality, calibration, and long-horizon robustness relative to thresholding baselines [2604.27536].

HIA is similar in spirit but works in black-box prompt optimization. It generates \(N\) candidate prompt modifications, scores them with heuristic reward models, keeps only the top \(K\), and spends expensive response-model calls only on that filtered subset. In the low-budget regime \(N=128, K=1\), the paper reports HelpSteer single-objective goal completion improving from 24.00 to 31.00 for BoN+H relative to BoN+Random, a 29% improvement [2508.05165]. An adjacent but importantly different line is Budgeted Multiple-Expert Deferral, which uses selective expert-cost querying to reduce **training-time** cost when learning a two-stage router; its budget mechanism is not inference-time deployment control in the strict sense [2510.26706].

## 6. Evaluation patterns, misconceptions, and limitations

The literature evaluates two-stage budget control with several recurring metrics. Token-budgeted reasoners report accuracy under a maximum budget, budget following ratio, budget utilization ratio, and length-compression figures [2508.17196][2505.11274]. Latency-oriented systems report completion rate, average response score, and deadline-sensitive performance under overrun strategies such as Kill and Skip-Next [2512.21859]. Shared-budget allocators use score, regret, and oracle imitation accuracy [2601.03822][2604.14853]. Belief-based escalators add Brier score, NLL, ECE, low-quality occurrence, CVaR, and recovery delay [2604.27536]. Anytime reasoning adds the Anytime Index to quantify how quickly quality improves with additional reasoning tokens [2601.11038].

Several misconceptions recur. First, “two-stage” does not always mean a two-phase inference controller. BudgetThinker’s two stages are supervised fine-tuning and reinforcement learning; inference remains a single autoregressive process with periodic control-token insertion [2508.17196]. ORBIT is multi-stage in training and exposes explicit budget modes at inference, but it does not learn a stage-1 mode selector from input [2601.08310]. BAVT is often interpretable as early exploration followed by late exploitation, but the paper’s mechanism is continuous budget annealing rather than a hard stage boundary [2603.12634]. By contrast, TimeBill, Predictive Scheduling, Veroic, SelfBudgeter, BET, and Star Elastic do implement unmistakable runtime decompositions [2512.21859][2602.01237][2604.27536][2505.11274][2605.11625][2605.07182].

Second, most budgets are proxies rather than exact deployment costs. BudgetThinker explicitly notes that token budget is only a proxy for actual latency and that inserted control tokens add small overhead [2508.17196]. BUDDY gives strict executed-layer-count control, but the paper reports that realized speedups are not perfectly proportional because routing and gather/scatter overhead can offset savings at light pruning [2606.09514]. TimeBill’s ETE coefficients are hardware- and implementation-dependent and must be recalibrated on new systems [2512.21859]. Star Elastic currently recomputes cache states on model switches, so the present latency gains are conservative with respect to future cache-reuse-capable runtimes [2605.07182].

Third, more inference-time compute is not uniformly beneficial. The robustness study based on budget forcing shows that in hidden chain-of-thought settings, larger reasoning budgets can improve robustness for prompt injection and prompt extraction, but when intermediate reasoning is exposed or tool-integrated, increased inference-time computation can reduce robustness, with an inverse scaling law formalized through the monotone growth of malicious-token exposure risk with chain length [2507.15974]. This is a genuine controversy for deployment: a second budget stage can improve performance under one threat model and worsen security under another.

Finally, two-stage control inherits the failure modes of both stages. Predictors can be miscalibrated, as shown by the crossing point where fine-grained size-based scheduling can underperform uniform allocation at larger budgets [2602.01237]. Belief-based escalation depends on the informativeness of black-box verifiable signals [2604.27536]. Preference-data prompting improves partial-solution quality but still lacks an explicit online stopping rule [2601.11038]. These limitations do not negate the framework; they delimit its current operational regimes.

Taken together, the literature defines two-stage inference-time budget control less as one algorithm than as a design principle: estimate, screen, or plan before spending scarce inference resources, then execute under an explicit budget with mechanisms that preserve accuracy, calibration, or utility as much as possible. The principal open distinction is no longer whether to control inference budgets, but which budget is being controlled, where the control signal is computed, and whether the two-stage decomposition occurs in training, in inference, or in both.

Source: https://www.emergentmind.com/topics/two-stage-inference-time-budget-control