Plan and Budget: Effective and Efficient Test-Time Scaling on Reasoning Large Language Models

• Junhong Lin, Xinyue Zeng, Jie Zhu, and Song Wang appear in the first author line.Source: SONING · p. 1
• Junhong Lin and Xinyue Zeng carry asterisks in the author listing.Source: SONING · p. 1
• Julian Shun, Jun Wu, and Dawei Zhou appear in the second author line.Source: SONING · p. 1

1 INTRODUCTION

• LLMs’ strong generalization and advances in test-time computation have expanded their use in complex reasoning, including high-stakes domains.Source: §1 · p. 1
• Fixed token constraints can reduce performance on complex queries by inducing underthinking, while overthinking wastes computation without improving accuracy.The introduction identifies limited work systematically addressing both failures in one framework.Source: §1 · p. 2
• Reasoning miscalibration is framed as a mismatch between computational effort and uncertainty across a query’s sub-questions.The paper proposes BAM to characterize heterogeneous uncertainty reduction and motivate structured, adaptive allocation.Source: §1 · pp. 2–3
• Plan-and-Budget decomposes queries into sub-questions, then adaptively assigns token budgets according to their uncertainty patterns.The paper also introduces E3 to capture the trade-off between reasoning accuracy and computational cost.Source: §1 · p. 3
• Across four LLMs and three domains, PLAN-AND-BUDGET improves accuracy by up to 70%, reduces token use by up to 39%, and raises E3 by up to 193.8% over strong baselines.On agentic planning, DS-Qwen-32B raises E3 from 0.16 to 0.47 with PLAN-AND-BUDGET, nearing DS-LLaMA-70B's 0.50 without planning.Source: §1 · p. 3

2.1 RELATED WORK

• Prior test-time scaling work reports that more inference compute can reduce failures but often has diminishing returns; this paper instead studies structured, uncertainty-aware allocation.Source: §2.1 · p. 3
• Uncertainty research distinguishes epistemic and aleatoric components, and LLM work has extended estimation through consistency checks and parameter-efficient ensembles.The paper positions its use of uncertainty decomposition for inference-time token allocation as a different application.Source: §2.1 · p. 3

2.2 PRELIMINARIES

• Reasoning miscalibration is a mismatch between query complexity and test-time reasoning depth, manifesting as overthinking or prematurely terminated reasoning.Source: §2.2 · p. 4
• Miscalibration is framed as unregulated effort across sub-questions: epistemic uncertainty may warrant more inference, while aleatoric uncertainty may favor concise solutions.The stated objective is to improve efficiency while mitigating miscalibration.Source: §2.2 · p. 4
• E3 jointly captures reasoning quality and computational cost, using average accuracy and decoding-token count while weighting accuracy more heavily than accuracy-per-token.The preliminary problem formulation seeks a computation strategy maximizing E3 under a query-level token budget.Source: §2.2 · p. 4

3 BUDGET ALLOCATION MODEL (BAM)

• The model assumes epistemic uncertainty decreases with tokens according to an inverse power law, with initial uncertainty and reduction complexity varying by sub-question.Higher β_ij means uncertainty is easier to reduce.Source: §3 · p. 5
• BAM represents total uncertainty as epistemic plus aleatoric uncertainty, treating the aleatoric component as irreducible through additional tokens.Source: §3 · p. 5
• The optimal allocation maximizes total utility under a token-budget constraint, with sub-question utility decreasing as uncertainty rises and α scaling by model or task.Source: §3 · p. 5
• BAM’s allocation rule is unimodal in complexity: budgets rise up to a peak, then fall as further effort yields diminishing returns.Moderately difficult sub-questions receive more tokens to avoid underthinking, while the most difficult receive fewer to prevent overthinking.Source: §3 · p. 5

4 REASONING CALIBRATION FRAMEWORK: PLAN-AND-BUDGET

• Exact marginal uncertainty gains and parameters c_ij and β_ij are unavailable in black-box LLMs, while estimating them directly would require costly sampling.Decay-based scheduling approximates uncertainty-aware allocation while keeping the total allocation within the global budget.Source: §4 · p. 6

4.1 PLANNING STEP: QUESTION DECOMPOSITION AS GUIDED SCAFFOLD

• The planning step decomposes a query into an ordered sequence of sub-questions and estimates each sub-question’s complexity as a soft, non-unique reasoning scaffold.Complexity scores may draw on LLM confidence, problem structure, or other heuristics and reflect epistemic uncertainty.Source: §4.1 · p. 6
• The main LLM answers each planned sub-question within its allocated token budget, then a synthesis function combines the responses into an answer to the original query.Source: §4.1 · p. 6

4.2 BUDGET STEP: DECAY-BASED BUDGET ALLOCATION

• Budget scheduling combines normalized complexity weights with a positional decay prior, allocating more tokens to early stages where uncertainty is typically higher.Later steps are generally narrower or more deterministic, and extra computation there risks wasting inference effort.Source: §4.2 · p. 7
• Figure 2 shows distinct allocation patterns under uniform complexity: polynomial and cosine favor early steps, linear declines gradually, and exponential is more balanced.The paper describes decay strategies as practical surrogates for Bayesian-optimal allocation, with different prioritization schemas.Source: §4.2 · p. 7, Figure 2 · p. 7

5.2 COMPARATIVE RESULTS

• Across datasets and model scales, PLAN-AND-BUDGET improves E3 by up to 193.8% while maintaining comparable or higher accuracy than Vanilla and Global Budget.Tables 3–5 cover MATH-500, NaturalInstructions, and TravelPlanner.Source: Table 3 · p. 8, §5.2 · pp. 8–9, Table 4 · p. 9

5.1 EXPERIMENT SETUP

• Experiments use MATH-500, NaturalInstructions, and TravelPlanner; four reasoning models; and baselines ranging from Vanilla to fixed and local budgets.Performance is measured by dataset score, average billed tokens, and E3; results are averaged over five runs.Source: §5.1 · p. 8

5.2 COMPARATIVE RESULTS

• On NaturalInstructions, E3 improves by 19.3–36.0%; on TravelPlanner, it reaches 0.16 → 0.47 (+193.8%) for DS-Qwen-32B.The reported gains are greatest on the more complex TravelPlanner benchmark.Source: Table 4 · p. 9, Table 5 · p. 9, §5.2 · p. 9
• On MATH-500, E3 rises from 4.55 →5.89 (+29.4%) for DS-LLaMA-70B and from 13.25 →15.95 (+20.3%) for o4-mini without compromising accuracy.The passage attributes limited Global Budget gains to its lack of uncertainty-awareness and reports 2–13% efficiency gains from planning alone.Source: §5.2 · p. 9
• Adding local budgeting improves E3 beyond structured planning alone, with reported gains of 14.1% on MATH-500, 24.2% on NaturalInstructions, and 33.3% on TravelPlanner.The results support adapting budgets to sub-questions rather than using a global allocation.Source: §5.2 · p. 9
• Polynomial decay and cosine annealing yield the highest E3 on mathematical and long-form planning tasks, while weighted or uniform schedules usually perform well on NaturalInstructions.The former strategies front-load computation toward early, uncertain steps.Source: §5.2 · p. 10
• On TravelPlanner, DS-Qwen-32B reaches E3 = 0.47 after Plan-and-Budget, near the unplanned DS-LLaMA-70B’s E3 = 0.50.The framework is model-agnostic and requires no retraining or fine-tuning.Source: §5.2 · p. 10

5.3 ROBUSTNESS AND SENSITIVITY ANALYSIS

• Robustness tests find better efficiency–accuracy trade-offs than early termination, 38–44% lower latency despite 3–7% planning overhead, stable rankings across E3 and A/T, and effective budget prompts and planner replacement.The latency comparison is against vanilla decoding; replacing the default planner with LLaMA 3.3-70B yields consistent improvements.Source: §5.3 · p. 10

6 CONCLUSION

• PLAN-AND-BUDGET improves LLM reasoning efficiency through structured planning and uncertainty-aware token budgeting, but currently requires an additional LLM call to generate its decomposition plan.The authors aim to develop a dedicated planner to support end-to-end reasoning within a single model.Source: §6 · p. 10

REPRODUCIBILITY STATEMENT

• The authors provide proofs, experimental details, dataset and model licenses, code, and scripts to reproduce every table and figure.Prompt templates and robustness and sensitivity analyses are also included.Source: REPRODUCIBILITY STATEMENT · p. 11

ETHICS STATEMENT

• The study uses public datasets and open or API-accessible models; its framework does not modify base-model internals and inherits their limitations and biases.The authors recommend task-specific evaluation before real-world deployment, particularly where efficiency could be prioritized over accuracy in high-stakes applications.Source: ETHICS STATEMENT · p. 11

A BROADER IMPACTS

• The authors suggest that reduced computational costs and energy use, along with a narrower small-to-large-model performance gap, may promote access to LLM capabilities.These are described as potential societal impacts.Source: A BROADER IMPACTS · p. 14

B LLM USAGE

• LLMs assisted with linguistic refinement and minor editing, but not with research design, experiments, results interpretation, or scientific contributions.The authors state that the scientific contributions are their responsibility.Source: B LLM USAGE · p. 14

C PROOF OF UNCERTAINTY DECOMPOSITION FOR LLMS

• The proof defines predictive uncertainty as the entropy of the marginal predictive distribution and applies the law of total entropy to decompose it.The marginal distribution is obtained by averaging over the parameter posterior.Source: C PROOF OF UNCERTAINTY DECOMPOSITION… · p. 14
• Aleatoric uncertainty is the expected conditional entropy, while epistemic uncertainty is mutual information reflecting disagreement across posterior samples.The proof states that total uncertainty is the sum of these components.Source: C PROOF OF UNCERTAINTY DECOMPOSITION… · p. 15
• The proof describes Monte Carlo estimation by drawing parameter samples from the posterior to approximate the uncertainty decomposition.It cites Hüllermeier & Waegeman (2021) for this approximation.Source: C PROOF OF UNCERTAINTY DECOMPOSITION… · p. 15
• For Plan-and-Budget, the decomposition is a theoretical lens rather than an inference-time computation; inference instead uses the greedy or most-likely generation path as a deterministic proxy.The authors state that they do not perform the computationally expensive Monte Carlo sampling during inference.Source: C PROOF OF UNCERTAINTY DECOMPOSITION… · p. 15

D PROOF OF LAGRANGE OPTIMALITY

• Because the aleatoric uncertainty term and α are constant with respect to each budget, maximizing total utility reduces to minimizing the budget-dependent expression.The derivation forms a Lagrangian, differentiates with respect to each sub-question budget, and solves for the allocation.Source: D PROOF OF LAGRANGE OPTIMALITY · pp. 15–16
• The allocation must satisfy the fixed total budget constraint, which is substituted into the derived per-sub-question budget expression.The proof then normalizes the allocations to satisfy the constraint.Source: D PROOF OF LAGRANGE OPTIMALITY · p. 16
• Homogeneous difficulty permits a closed-form solution for λ, whereas heterogeneous difficulty requires numerical determination of λ to satisfy the budget constraint.The proof gives the normalized general form and states that it satisfies the constraint.Source: D PROOF OF LAGRANGE OPTIMALITY · p. 16

E ANALYSIS OF THE RELATIONSHIP BETWEEN bij AND βij

• The sign of g′(β) is non-monotonic, so g(β) first increases and then decreases, making both g(β) and f(β) unimodal.This analysis concerns the relationship between the allocation function and βij.Source: E ANALYSIS OF THE RELATIONSHIP BETWEEN… · pp. 16–17

F.1 ADDITIONAL BASELINES: EARLY TERMINATION

• PLAN-AND-BUDGET consistently achieves superior efficiency–accuracy trade-offs across three LLMs compared with Certaindex, which often reduces tokens at an accuracy cost.The comparison is on MATH-500; Plan-and-Budget preserves reasoning depth in high-uncertainty stages while removing redundant later computation.Source: §F.1 · p. 17, Table 7 · p. 17

F.2 LATENCY BREAKDOWN

• Planning adds modest overhead, but shorter reasoning chains reduce total latency.The authors state that the latency gains represent practical runtime improvements, not only token savings.Source: §F.2 · p. 17

F.3 METRIC ROBUSTNESS: E3 VS. A/T

• Plan-and-Budget’s ranking over baselines remains consistent under A/T across methods and datasets; E3 distinguishes cases where A/T is comparable, including 6.86 vs. 6.87 under A/T and 2.96 vs. 2.88 under E3.E3 penalizes accuracy drops more heavily, while the reported rankings remain stable across the two metrics.Source: §F.3 · p. 18, Table 8 · p. 18, Table 9 · p. 18

F.4 BUDGET SENSITIVITY ANALYSIS

• Performance scales smoothly as initial and per-level budgets increase, without abrupt degradation or instability, supporting controllable behavior across resource regimes.The sensitivity analysis varies Binit and Bper-level; Table 11 relates total budget to question difficulty.Source: §F.4 · p. 19, Table 11 · p. 19

F.5 PLANNER SENSITIVITY

• Improvements persist when LLaMA 3.3-70B replaces the default planner, and the analysis attributes efficiency gains to the allocation mechanism rather than planner-specific behavior.The comparison covers MATH-500, NaturalInstructions, and TravelPlanner.Source: §F.5 · p. 19, Table 12 · p. 19, Table 13 · p. 19

F.6 TOKEN USAGE DISTRIBUTION ON TRAVELPLANNER

• On TravelPlanner, global-budget methods use more tokens and have lower pass rates than Plan-and-Budget at every difficulty level, while local budgeting adapts computation to sub-question difficulty.Figure 4 shows token usage distributions and pass rates by difficulty level; the text associates this calibration with higher overall E3.Source: Figure 4 · p. 20, §F.6 · p. 20

G.1 DATASET DESCRIPTIONS AND EVALUATION METRICS

• MATH-500 uses exact-match accuracy; NaturalInstructions evaluates 500 sampled queries with ROUGE, while TravelPlanner uses validation-set hard-constraint pass rate in a no-tool, sole-planning setting.TravelPlanner omits the stricter full success rate to isolate planning competence; the datasets span symbolic math, instruction following, and long-horizon planning.Source: §G.1 · pp. 20–21

G.2 LICENSES FOR EXISTING ASSETS

• The appendix lists licenses for the DeepSeek models, QwQ-32B, o4-mini, LLaMA models, and the three evaluation datasets.It notes that o4-mini is accessed through the API and that its weights are not released or modified.Source: §G.2 · p. 21

G.3 COMPUTE RESOURCES

• Experiments use API-accessible models and cloud-based preprocessing and evaluation; no model training or fine-tuning is involved, and reported token counts can inform runtime estimates.The stated generation-rate range is typically 50–80 tokens/sec, depending on provider.Source: §G.3 · p. 21

H PROMPT TEMPLATES

• The local-budget prompt requests 2–5 linked sub-questions, fewer for easier problems, and assigns each sub-question a solving budget.Each sub-question is paired with a brief hint explaining its purpose; the template specifies solving each with up to a given number of words.Source: H PROMPT TEMPLATES · pp. 22, 25