Steering LLM Thinking with Budget Guidance

TL;DR

The paper tackles the problem of prohibitive inference costs from deep-thinking LLMs by enabling budget-aware reasoning without fine-tuning. It introduces budget guidance, a test-time framework that uses a lightweight auxiliary predictor to model a Gamma distribution over the remaining thinking length and softly modulates the LLM's token generation to meet a target thinking budget. Framed via Bayes-like budget conditioning, the approach achieves strong token efficiency, achieving up to 26% accuracy gains on MATH-500 under tight budgets while using only about 63% of the thinking tokens of full-thinking baselines. The predictor trained on math reasoning generalizes across domains, offering cross-domain applicability and enabling real-time, cost-conscious reasoning across diverse tasks without modifying the base LLM.

Abstract

Recent deep-thinking large language models often reason extensively to improve performance, but such lengthy reasoning is not always desirable, as it incurs excessive inference costs with disproportionate performance gains. Controlling reasoning length without sacrificing performance is therefore important, but remains challenging, especially under tight thinking budgets. We propose budget guidance, a simple yet effective method for steering the reasoning process of LLMs toward a target budget without requiring any LLM fine-tuning. Our approach introduces a lightweight predictor that models a Gamma distribution over the remaining thinking length during next-token generation. This signal is then used to guide generation in a soft, token-level manner, ensuring that the overall reasoning trace adheres to the specified thinking budget. Budget guidance enables natural control of the thinking length, along with significant token efficiency improvements over baseline methods on challenging math benchmarks. For instance, it achieves up to a 26% accuracy gain on the MATH-500 benchmark under tight budgets compared to baseline methods, while maintaining competitive accuracy with only 63% of the thinking tokens used by the full-thinking model. Budget guidance also generalizes to broader task domains and exhibits emergent capabilities, such as estimating question difficulty. The source code is available at: https://github.com/UMass-Embodied-AGI/BudgetGuidance.

Figures (11)

• Figure 1: Deep-thinking models often produce excessively long reasoning traces, leading to high latency and unnecessary computation. Existing inference-time methods like budget forcing rely on simplistic heuristics such as abruptly stopping, which can result in incomplete reasoning and degraded answer quality. In contrast, our method, budget guidance, steers the reasoning process toward the target budget in a smoother and more natural way, without any LLM fine-tuning.
• Figure 2: An overview of budget guidance. A lightweight predictor uses the LLM’s hidden states to predict a Gamma distribution over the remaining reasoning length for each candidate token. We then use the CDF of Gamma distribution to compute a predictor score, which is combined with the LLM’s output score to guide generation. The result is soft, token-level steering toward budget-conditioned reasoning without any LLM fine-tuning.
• Figure 3: Accuracy vs. thinking length on math benchmarks.
• Figure 4: Thinking length controllability measured on MATH-500 benchmark.
• Figure 5: Correlation between question difficulties and estimated thinking lengths.