SelfBudgeter: Adaptive Token Allocation for Efficient LLM Reasoning

1. Problem Statement

Recent Large Reasoning Models (LRMs), such as O1 and DeepSeek-R1, have demonstrated exceptional performance on complex tasks by utilizing long Chain-of-Thought (CoT) processes. However, a critical inefficiency has emerged: these models tend to "overthink" even simple queries.

• Resource Waste: Models generate excessive tokens for trivial questions (e.g., "2+3"), leading to unnecessary computational costs.
• User Latency: The excessive generation time results in prolonged waiting times for users.
• Limitations of Existing Solutions:
  • Prompt-based methods lack precise control over output length.
  • Integrated training (SFT/RL) often fails to strictly adhere to user-specified length constraints or requires additional overhead.
  • Router-based methods introduce extra computational costs via auxiliary classifiers.
  • Existing budget-aware methods (e.g., TALE, TOPS) fail to simultaneously enable autonomous budget estimation and strict adherence to those budgets.

2. Methodology: SelfBudgeter

The authors propose SelfBudgeter, a self-adaptive reasoning strategy that enables models to autonomously estimate the required token budget and generate responses strictly within that budget. The framework operates in two main stages:

A. Output Format

The model is trained to output a specific structure before the solution:
<budget> [integer] </budget><solution> [response] </solution>

• No Constraint: The model estimates the minimal tokens needed based on problem complexity.
• User Constraint: The <budget> tag is pre-filled by the user, and the model must generate the solution within that limit.

B. Training Framework

The training consists of a Cold-Start phase followed by Reinforcement Learning (RL).

1. Cold-Start (Supervised Fine-Tuning):

   • The model is fine-tuned on datasets (GSM8K, MATH, s1k-1.1) to learn the new output format (<budget> + <solution>).
   • The <budget> tag is pre-populated with the token count of the ground-truth or high-quality reasoning chains to teach the model the correlation between problem difficulty and token usage.

2. Reinforcement Learning (GRPO):

   • The model is optimized using Group Relative Policy Optimization (GRPO).
   • Reward Function: A composite reward function guides the model to balance accuracy, length compression, and budget adherence. It consists of three components:
     • Format Penalty ($r_f$): Penalizes deviations from the required output format.
     • Budget Penalty ($PB$): Penalizes the model if the estimated budget ($b$) exceeds the maximum acceptable budget ($b_{max}$). $b_{max}$ is set to the length of a correct answer if the base model solves it correctly; otherwise, it is $\infty$.
     • Precise Budget Control Reward (PreB): This is the core innovation. It uses a tightness coefficient ($\alpha$) to align the actual response length ($\ell$) with the predicted budget ($b$).
       • If $|\ell - b| / b > \alpha$, the reward drops to a minimum.
       • Inside the acceptable range, a cosine reward is applied.
       • Dual Mechanism: For incorrect answers, the reward encourages longer chains (to find the solution). For correct answers, the reward peaks at the minimal sufficient length to prevent overthinking.

3. Dynamic $\alpha$ Scheduling:

   • To prevent "reward hacking" (where the model inflates the budget to easily meet the length constraint), the tightness coefficient $\alpha$ is dynamically decreased linearly during training (from 6.0 to 0.1). This gradually forces the model to tighten its budget estimation and output length.

3. Key Contributions

1. SelfBudgeter Framework: A novel training framework enabling LLMs to explicitly predict and control output length via self-estimated token budgets.
2. Efficiency without Accuracy Loss: Demonstrates that strict length control does not degrade reasoning quality; in many cases, it improves accuracy by forcing concise, essential reasoning.
3. Scalability: The method is validated across different model sizes (1.5B and 7B), showing consistent benefits.
4. User Control: Allows users to either let the model self-estimate the budget or pre-fill a specific token limit, providing transparency on expected latency.

4. Experimental Results

Experiments were conducted on GSM8K, MATH500, and AIME2025 benchmarks using DeepSeek-R1-Distill-Qwen models (1.5B and 7B).

• Performance vs. Baselines:

  • SelfBudgeter-1.5B achieved 84.10% accuracy on GSM8K (vs. 73.09% baseline) while compressing response length by 57% (from ~2865 to ~1231 tokens).
  • On MATH500, it improved accuracy by 3.54 points while reducing length by 56%.
  • On AIME2025, it compressed length to 30% of the original while maintaining comparable accuracy.
  • Compared to baselines like L1-Max (which suffers accuracy drops on hard tasks) and E1-Math (which relies on hard truncation), SelfBudgeter achieved the best or second-best accuracy across all datasets with superior compression.

• Generalization:

  • Tested on general knowledge benchmarks (GPQA and SCoRE), SelfBudgeter maintained competitive accuracy while significantly reducing reasoning length compared to long-reasoning baselines.

• Ablation Studies:

  • Removing any component of the reward function (Format, Correctness, PreB, or Budget Penalty) led to either accuracy collapse or poor length control. The combination of all components was necessary for robust performance.
  • Linear $\alpha$ scheduling proved superior to fixed or cosine scheduling, offering the best trade-off between compression and accuracy.

5. Significance and Impact

• Resource Efficiency: SelfBudgeter offers a practical solution to the "overthinking" problem, significantly reducing inference costs and latency for both simple and complex queries.
• User Experience: By predicting token budgets upfront, the system allows users to anticipate waiting times and decide whether to wait for a full solution or terminate early, enhancing interaction control.
• Scalable Reasoning: It proves that reasoning models do not need to generate excessively long chains of thought to be accurate; adaptive budgeting allows for "thinking just enough."
• Deployment Potential: The method is highly suitable for resource-constrained environments and can be integrated into existing LLM pipelines to optimize inference costs without sacrificing model intelligence.

In conclusion, SelfBudgeter represents a significant step toward controllable, efficient, and user-centric reasoning models, bridging the gap between autonomous budget estimation and strict adherence to computational constraints.