Steering LLM Thinking with Budget Guidance

1 Introduction

• Long reasoning traces can impose unnecessary computation and latency, particularly in real-time customer-facing interactions.The paper notes that excessive latency can degrade user experience and responsiveness.Source: §1 · p. 1
• Fine-tuning can restructure reasoning but costs substantial resources and may unexpectedly affect safety, while heuristic inference-time interventions can disrupt reasoning quality.The paper frames the central challenge as combining inference-time flexibility with principled restructuring under budget.Source: §1 · pp. 1–2
• Budget guidance predicts a distribution over remaining reasoning length and uses it to modulate generation toward a target budget without fine-tuning the LLM.The approach uses a lightweight auxiliary module for test-time control.Source: §1 · pp. 2–3
• 26% higher accuracy than budget forcing under tight budgets on MATH-500, while reducing full thinking length by 37% with minimal accuracy degradation.Both comparisons are reported for MATH-500.Source: §1 · p. 2
• The predictor considers task difficulty and instruction type, and the approach is reported to generalize across datasets and domains.The contributions also describe strong trade-offs between thinking length and accuracy across multiple benchmarks.Source: §1 · pp. 2–3

2.1 Efficient LLM Reasoning

• Prior efficiency methods such as ThinkPrune, Z1, and COCONUT aim to shorten reasoning, typically rely on expensive LLM fine-tuning, and do not primarily target length control.More recent approaches explore length control through heuristic rules or fine-tuning; budget guidance is presented as a fine-tuning-free alternative.Source: §2.1 · p. 3

2.2 Guidance and Guided Generation

• Diffusion-model guidance steers generation at sampling time, with classifier guidance modifying the diffusion score using an auxiliary classifier’s log-likelihood gradient.The passage describes this as biasing sampling toward desired content.Source: §2.2 · p. 3
• LLM guided generation typically constrains outputs to structural requirements such as regular expressions or context-free grammars for downstream syntactic correctness.Source: §2.2 · p. 3
• Budget guidance extends guided generation to budget-conditioned generation, softly steering LLM output toward a specified thinking budget.The authors describe this as enabling efficient, controlled reasoning without compromising output quality.Source: §2.2 · p. 3

3.1 The Budget Guidance Framework

• At each token, budget guidance combines the LLM’s unconditional distribution with the predicted probability that a candidate token’s remaining reasoning fits the budget, then renormalizes.This produces a budget-conditional generation distribution without fine-tuning the LLM.Source: §3.1 · p. 4

3.2 An Auxiliary Thinking Length Predictor

• The predictor models log remaining length with a Gamma distribution, using shape and rate vectors whose CDF yields within-budget probabilities for candidate tokens.This replaces predicting a separate full distribution for each vocabulary token with two n-dimensional parameter vectors.Source: §3.2 · p. 5

3.3 Training the Predictor

• Predictor training maximizes likelihood on reasoning chains truncated at varied positions; the traces come from OpenR1-Math-220k, distinct from the inference dataset and task.The predictor is trained to estimate remaining length from partial reasoning chains.Source: §3.3 · p. 5

3.4 Architecture of the Predictor

• BERT-base consumes concatenated last-token hidden states across target-LLM layers; a [CLS] summary projects to per-token Gamma parameters.A softplus activation ensures the shape and rate parameters are non-negative.Source: §3.4 · p. 5

3.5 Skipping Modulation

• 0.6% total latency increase for a 7B-parameter LLM; modulation runs only at reasoning-paragraph starts, where uncertainty is typically highest.Paragraph starts are identified by newline delimiters.Source: §3.5 · p. 6

4.1 Settings

• The predictor is trained on math reasoning traces, while evaluation includes four math benchmarks and four broader-domain reasoning benchmarks in a zero-shot setting.The evaluated models are DeepSeek-R1-Distill-Qwen-7B, DeepSeek-R1-Distill-Qwen-32B, and Qwen3-8B; decoding is greedy.Source: §4.1 · p. 6

4.2.1 Evaluation on Math Reasoning Benchmarks

• Across three models and four math datasets, budget guidance consistently outperforms budget forcing at comparable average thinking lengths, with substantially higher performance than NoThinking.The evaluations use approximately half of each original model’s full thinking length.Source: §4.2.1 · p. 6, Table 1 · p. 7

4.2.2 Accuracy–Thinking Length Tradeoff Analysis

• Across most math benchmarks, budget guidance achieves higher accuracy than budget forcing over a range of thinking lengths, especially as budgets tighten.The paper attributes this pattern to adapting reasoning under strict budgets rather than abruptly truncating traces.Source: Figure 3 · p. 7, §4.2.2 · p. 8

4.2.3 Fine-Grained Control of Thinking Length

• At each tested MATH-500 budget, at least 75% of budget-guided samples stay within budget, and median thinking length closely matches the target.Unlike budget forcing, guidance softly steers generation rather than enforcing a hard cutoff.Source: §4.2.3 · p. 8, Figure 4 · p. 8

4.2.4 Out-of-Domain Transferability

• Despite training exclusively on math data, the predictor outperforms budget forcing on all four evaluated non-math benchmarks, though gains are less pronounced than in-domain.The authors identify training on reasoning traces from broader domains as future work.Source: §4.2.4 · pp. 8–9, Table 2 · p. 8

4.3 Insights into What the Predictor Learns

• On MATH-500 (in-domain) and LiveCodeBench (out-of-domain), estimated thinking length increases with question difficulty, suggesting the predictor captures a general understanding of difficulty for effective difficulty estimation.Source: §4.3 · p. 9
• On MATH-500, a long-reasoning prompt yields longer estimated thinking lengths than a concise prompt, with a statistically significant difference (p = 0.0028).The authors interpret this as evidence that the predictor is prompt-aware.Source: §4.3 · p. 9, Figure 6 · p. 9

4.4 Case Study

• In the MATH-500 case study, stricter budgets produce concise answers without reflection, while larger budgets elicit fuller analysis and reflective phrases.Both traces end appropriately rather than being cut off.Source: Figure 7 · p. 9, §4.4 · p. 10

5 Conclusion

• Budget guidance steers LLM reasoning under a thinking budget without requiring LLM fine-tuning, and the authors identify budget-conditioned generation as a promising direction.The conclusion reports natural control of reasoning and improved token efficiency on challenging benchmarks.Source: §5 · p. 10

A Quantitative Reasoning Behavior Analysis

• Smaller budgets substantially reduce reasoning-related keyword frequency, while sufficiently high budgets produce behavior close to the full-thinking baseline.The comparison counts terms such as “wait” and “alternatively” for 500-, 2000-, and 4000-token budgets on DS-7B using MATH-500.Source: A Quantitative Reasoning Behavior… · p. 12, Figure 8 · p. 12

B Additional Reasoning Samples

• Under identical budget constraints, budget guidance solves illustrative MATH-500 problems where budget forcing fails.The authors describe its thinking-length control as more precise and rational than budget forcing.Source: B Additional Reasoning Samples · p. 12

• In the 1000-token examples, budget guidance controls reasoning to just under the limit and also produces a drastically shortened reasoning process that still solves the question correctly.Source: §B.1 · p. 12, Figure 9 · p. 13
• The examples compare budget forcing and budget guidance under the same thinking budget.Source: §B.1 · p. 12
• Budget forcing and budget guidance use a similar number of thinking tokens under the specified budget, but only budget guidance reaches the correct answer in Figure 9.Source: Figure 9 · p. 13

• Across diverse budget settings, budget guidance flexibly steers reasoning while ensuring an accurate reasoning process.Source: §B.2 · p. 13
• The method controls thinking length under varying budget constraints.Source: §B.2 · p. 13
• The reported reasoning remains accurate across the budget settings examined.Source: §B.2 · p. 13

• MATH-500 is a 500-problem subset of MATH with difficulty levels 1–5, while AIME-2024 contains 30 problems and AMC includes 83 problems from two AMC12 years.Source: C Dataset Description · p. 13
• OlympiadBench contributes 675 problems from its math subset, while GPQA Diamond contains 198 difficult questions across biology, physics, and chemistry.Source: C Dataset Description · p. 13
• The non-math evaluations include FOLIO’s 203-problem test set, TableBench’s 493 numerical-reasoning problems, and 323 LiveCodeBench problems from August 2024 to January 2025.Source: C Dataset Description · p. 14

D Training Data Augmentation

• One new sample is added for each original sample by using the answer message as the predictor’s THINK_MESSAGE, doubling the training set.The predictor operates only on THINK_MESSAGE, so the original ANSWER_MESSAGE is otherwise unused during training.Source: D Training Data Augmentation · p. 14

E Prompt Description

• The prompt-awareness experiment uses system prompts asking for thorough step-by-step reasoning or concise reasoning with minimal steps.Source: E Prompt Description · p. 15