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

🔗 Project Code: https://github.com/junhongmit/P-and-B

🚀 In this work, we introduce Plan-And-Budget, a test-time framework that improves reasoning efficiency in LLMs by combining structured planning and uncertainty-aware compute allocation. No retraining needed—just smarter token usage.

😆 Takeaways:

1. Reasoning Miscalibration in LLMs
   LLMs often overthink easy queries (verbose, wasteful) or underthink hard ones (premature, incorrect). We identify this mismatch as a core inefficiency in current inference.

2. Bayesian Budget Allocation Model (BBAM)
   We formalize reasoning as a sequence of sub-questions with varying uncertainty. BBAM allocates more tokens to epistemically uncertain steps, and fewer to those dominated by aleatoric noise.

3. Plan-And-Budget Framework

   • Plan step: Decompose queries into sub-questions using lightweight planning.
   • Budget step: Allocate tokens using decay-based heuristics (linear, polynomial, cosine, etc.).
     It’s model-agnostic, inference-only, and compatible with any LLM.

4. ℰ³ Score: A New Efficiency-Aware Effectiveness Evaluation Metric
   Defined as ℰ³ = Accuracy² / Tokens, this metric rewards high-accuracy, low-token reasoning—penalizing both wasteful and careless inferences.

5. Strong Empirical Results
   We extensively evaluate Plan-And-Budget on 3 types of reasoning tasks and 4 state-of-the-art reasoning LLMs, and shown improvements:

   • 📈 Accuracy: up to +70%
   • 🔻 Token use: up to –39%
   • 🔥 ℰ³: up to +187.5%

On agentic planning task, our method lifts a small 32B model to match the efficiency of a 70B model—without fine-tuning.

6. Early Steps Matter Most
   Polynomial and cosine decay schedulers work best, confirming that early reasoning steps—where uncertainty is highest—deserve more compute.

🧠 Plan before you think. Budget as you go. Smarter reasoning starts now.