PaperlayerView exact revision on arXiv

Source-linked AI summary

SelfBudgeter: Adaptive Token Allocation for Efficient LLM Reasoning

Zheng Li, Qingxiu Dong, Jingyuan Ma, Di Zhang, Kai Jia, Zhifang Sui

arXiv:2505.11274v6cs.AIcs.CL

TL;DR

Reasoning models can spend excessive tokens on simple queries, wasting resources and prolonging user waits. SelfBudgeter predicts or accepts a token budget and trains with GRPO to produce budget-consistent responses. The paper reports shorter responses while preserving performance, and users can use the budget to anticipate duration or constrain output length.

  • Problem

    Reasoning models may use overly long thought processes even for simple questions, wasting computational resources and increasing user waiting times.

    Source: §1 · p. 1

  • Method

    SelfBudgeter trains models to estimate a token budget from the query and generate a solution that follows the estimated or user-defined budget.

    Source: §1 · p. 1, §3.1 · p. 3

  • Results

    SelfBudgeter reduces response length while preserving model performance, sometimes improving accuracy, and shows consistent benefits across larger models.

    Source: §1 · p. 2

  • Takeaways & Limitations

    The predicted budget lets users anticipate inference duration and decide whether to continue, while a user-set budget constrains response length.

    Source: §1 · p. 1

    The methodology may face practical challenges in more complex or diverse real-world scenarios, and external factors beyond the study’s scope could affect generalizability.

    Source: Limitations · p. 9

Abstract

from arXiv · show

Recently, large reasoning models demonstrate exceptional performance on various tasks. However, reasoning models always consume excessive tokens even for simple queries, leading to resource waste and prolonged user latency. To address this challenge, we propose SelfBudgeter - a self-adaptive reasoning strategy for efficient and controllable reasoning. Specifically, we first train the model to self-estimate the required reasoning budget based on the query. We then introduce budget-guided GRPO for reinforcement learning, which effectively maintains accuracy while reducing output length. Experimental results demonstrate that SelfBudgeter dynamically allocates budgets according to problem complexity, achieving an average response length compression of 61% on math reasoning tasks while maintaining accuracy. Furthermore, SelfBudgeter allows users to see how long generation will take and decide whether to continue or stop. Additionally, users can directly control the reasoning length by setting token budgets upfront.

1 Introduction

Reasoning models can overthink simple questions, wasting compute and extending user waits. SelfBudgeter predicts or accepts a token budget and generates within it; the paper reports shorter responses while preserving performance, with benefits across model sizes.

1 Introduction

  • Reasoning models may overthink simple questions, wasting computational resources and increasing user waiting times.The paper illustrates the issue with a simple arithmetic query that prompts many solutions and substantially more tokens.Source: §1 · p. 1
  • SelfBudgeter estimates the minimal token budget for unrestricted queries or follows a user-defined budget, then generates a response of the corresponding length.The displayed budget also helps users anticipate generation time and decide whether to continue or stop.Source: §1 · p. 1, Figure 1 · p. 2
  • The training framework first teaches budget prediction with a cold-start stage, then uses RL to produce concise, budget-consistent outputs.The introduction identifies explicit length prediction and control as the framework’s central contribution.Source: §1 · pp. 1–2
  • SelfBudgeter reduces response length while preserving performance, sometimes improving accuracy, with consistent benefits across larger model sizes.The introduction presents both outcomes as contributions.Source: §1 · p. 2

2 Related Work

Prior approaches address overthinking through prompting, training, or routing, but the paper identifies gaps in precise control, overhead, and autonomous budget estimation. SelfBudgeter aims to estimate an appropriate budget and generate text that adheres to it.

2 Related Work

  • Prompt-based, training-based, and router-based approaches address overthinking, but are described as lacking precise control, requiring extra computation, or failing to bridge autonomous estimation and strict adherence.The limitations vary by approach: prompts lack precision, routers add overhead, and training methods do not reliably satisfy user-specified length constraints.Source: §2 · p. 2
  • Prompted word, character, token, and step budgets have been explored, but precise control over model output remains challenging.The passage also notes that TALE and TOPS do not explicitly guide an optimal token budget or effectively enforce a given budget.Source: §2 · p. 2
  • SelfBudgeter is presented as combining autonomous token-budget estimation with generation that strictly follows the estimated budget.This contrasts with methods described as controlling length under a given budget without autonomously estimating an appropriate response length.Source: §2 · pp. 2–3

3 Method

SelfBudgeter uses rewards to train a model to predict a budget based on question difficulty and generate an accurate response that conforms to it. Its reward design penalizes excessive budgets, aligns output length with the target, and prioritizes correctness.

3 Method

  • SelfBudgeter uses a reward-guided format to predict a token budget before answering, allocate budgets by question difficulty, and generate accurate solutions of appropriate length.The method is motivated by the observation that response lengths generally increase with problem difficulty, though models may overthink simple problems.Source: §3 · p. 3

3.1 SelfBudgeter

  • GRPO trains the model to predict difficulty-based budgets and produce outputs whose lengths conform to them.The reward design includes format, budget, length, and correctness considerations.Source: §3.1 · p. 3

3.2 Budget Penalty

  • The Budget Penalty penalizes an estimated budget above bmax; bmax is the maximum acceptable budget for a question and is based on a correct base-model response length, or set to infinity otherwise.No penalty applies when the estimated budget is within bmax.Source: §3.2 · pp. 3–4

3.3 Precise Budget Control Reward

  • PreB uses a tightness coefficient α to align response length with a specified budget, allowing approximate adherence within α · b.Responses outside that range receive the minimum reward; the reward design favors minimal sufficient length for correct answers and longer reasoning for incorrect ones.Source: §3.3 · p. 4

3.4 Accuracy Reward

  • The Accuracy Reward is configured so that the minimum reward for a correct response exceeds the maximum reward for an incorrect response, balancing correctness with length control.Together, Budget Penalty, PreB, and Accuracy Reward form the method’s three core modules.Source: §3.4 · p. 4

4 Experiment

The training template places a budget before the solution, with the budget predicted when no user limit is given and prefilled when one is specified. Data preprocessing retains correct cold-start responses and derives an acceptable budget for RL questions.

4.1 Training Template

  • The training format requires a budget estimate before the answer; without a user constraint the model predicts it, while a user-specified limit is prefilled for solution generation.The output uses separate budget and solution fields.Source: §4.1 · p. 4

4.2 Data Preprocessing

  • Cold-start preprocessing retains correct responses and their lengths, while RL preprocessing calculates bmax as the maximum acceptable token budget for each question.Incorrect cold-start answers are discarded to avoid reinforcing incorrect responses.Source: §4.2 · p. 4

＋Budget Penalty

SelfBudgeter combines cold-start format adaptation with RL-based token budgeting and a dynamic α schedule. On reported benchmarks, it compresses responses while retaining strong accuracy at both 1.5B and 7B scales.

＋Budget Penalty

  • For correct responses, the budget penalty encourages further length compression and sets bmax to the current response length; for incorrect answers, any budget is acceptable.Source: ＋Budget Penalty · p. 5

4.3 Cold Start

  • Jointly learning output format, budgets, budget-consistent solutions, and accuracy proved difficult, so cold-start first teaches the output format before RL training.Cold-start examples use prepopulated budget counts and model-generated solutions.Source: §4.3 · p. 5

4.4 Experiment Settings

  • Experiments use R1-1.5B and R1-7B, cold-start on 3,630 samples from GSM8K, MATH, and s1k-1.1, then train three epochs on 30K STILL-3-Preview-RL-Data.The RL dataset contains samples based on MATH, NuminaMathCoT, and AIME 1983–2023.Source: §4.4 · pp. 5–6

4.5 Main Results

  • On GSM8K, SelfBudgeter gains 11.01 accuracy points and compresses length to 43% of baseline; on MATH500, it gains 3.54 points with length at 44%.On AIME2025, response length is 30% of baseline while accuracy remains comparable.Source: Table 1 · p. 6, §4.5 · p. 6
  • SelfBudgeter consistently balances accuracy and response length better than existing baselines, achieving the best or second-best accuracy across all datasets while substantially reducing response length.L1 collapses on AIME2025, while E1’s hard truncation severely harms accuracy.Source: §4.5 · pp. 6–7
  • At 7B, SelfBudgeter achieves the highest accuracy on MATH500 and AIME2025, second-best on GSM8K by 0.68 points, and a 48% average compression ratio.The passage reports robust gains across datasets of varying difficulty.Source: §4.5 · p. 7

4.6 Dynamic Alpha Schedule

  • A linearly decreasing α tightens the acceptable-length range and encourages predicted budgets to converge toward actual output length, avoiding loose- and tight-fixed-α failures.The schedule uses αstart = 6.0 and αend = 0.1.Source: §4.6 · p. 7

5 Discussion

The discussion tests how reward components and α scheduling affect the accuracy–compression trade-off, then evaluates transfer beyond math benchmarks. Its appendices further examine budget adherence, complexity-sensitive allocation, and response quality under compression.

5 Discussion

  • Appendices C–F examine pre-filled-budget adherence, difficulty-sensitive allocation, constraint compliance with response quality, and case studies of compression’s effects on response quality and structure.The discussion also describes reward-component and α-scheduling analyses and evaluates general-domain knowledge reasoning.Source: §5 · pp. 7–8

5.1 Ablation Study on Reward Components

  • SelfBudgeter combines format, correctness, PreB, and budget-penalty rewards to achieve the strongest overall accuracy–compression trade-off across the ablation.The individual components address structure, reasoning quality, predicted-budget stability, and budget compliance; only their combination yields robust length control without degrading accuracy.Source: Table 2 · p. 8, §5.1 · p. 8

5.2 α Scheduling

  • Linear α scheduling maintains accuracy with over 61% compression, while Fixed and Cosine compress more but tend to over-compress difficult problems and reduce accuracy.Fixed achieves the strongest length compression, with Cosine slightly weaker.Source: Table 3 · p. 8, §5.2 · p. 8

5.3 Generalization Ability

  • On GPQA and SCoRE, SelfBudgeter-1.5B reduces reasoning length while retaining competitive accuracy; it leads the compared methods on SCoRE accuracy with shorter responses.Its GPQA accuracy is slightly below the strongest baseline but remains comparable despite using far fewer reasoning tokens.Source: §5.3 · pp. 8–9, Table 4 · p. 9

6 Conclusion

SelfBudgeter predicts and follows token budgets for reasoning, balancing accuracy and response length. Its budget predictions can help users anticipate inference duration, and output length can also be dynamically regulated.

6 Conclusion

  • SelfBudgeter autonomously predicts required reasoning token budgets and adheres to self-imposed constraints to optimize the accuracy–response-length trade-off.Output length can also be dynamically regulated through transformation functions when required.Source: §6 · p. 9
  • Users can use SelfBudgeter’s token-budget predictions to anticipate total inference duration.Source: §6 · p. 9
  • SelfBudgeter demonstrates performance comparable to several existing methods in resource-efficient reasoning, with potential for resource-constrained environments.Source: §6 · p. 9

Limitations

The authors flag practical challenges in applying the methodology to more complex or diverse real-world scenarios and note that external factors beyond the study’s scope could affect generalizability.

Limitations

  • The methodology may face practical challenges in more complex or diverse real-world scenarios, and external factors beyond the study’s scope could influence generalizability.Source: Limitations · p. 9

Ethical Considerations

The paper reports no new data collection, human participants, or direct ethical risks, and identifies no specific ethical concerns beyond general LLM-deployment risks. It also documents its compute environment and fine-tuning and GRPO configurations.

Ethical Considerations

  • The work identifies no specific ethical concerns beyond general risks associated with LLM deployment and involves no new data collection or human participants.Source: Ethical Considerations · p. 9

A.1 Experimental Environments

  • Fine-tuning and inference ran on two A40 GPUs, while GRPO training ran on two A100 GPUs.The server had two 80GB A100 GPUs and two 45GB A40 GPUs.Source: §A.1 · p. 12

A.2 Parameter Settings

  • Cold-start fine-tuning used a 16,384-token sequence cap, batch size 1 per device, two-step gradient accumulation, a cosine learning-rate schedule, and one training epoch.The base learning rate was 5e-5 with 10% warm-up; 10% of the training set was used for validation.Source: §A.2 · p. 12
  • GRPO training used three epochs, with checkpoints selected at specified training steps for the GSM-initialized, s1k-initialized, and larger-model variants.The appendix also points to Table 5 for reward-function hyperparameters.Source: §A.2 · p. 12, Table 5 · p. 12

B Cold-start Data Selection

Cold-start initialization changes the accuracy–length trade-off; responses also grow with user-set budgets and allocated budgets rise with problem difficulty. The analyses report budget-following across configurations and increasing budgets across difficulty tiers.

B Cold-start Data Selection

  • Cold-start prompts ask the model to estimate the required token total before reasoning and to enclose the budget and solution in separate tags.Source: B Cold-start Data Selection · p. 12, Figure 4 · p. 12
  • s1k initialization yields 74.18% vs. 63.46% on MATH at α = 0.2 and 80.44% vs. 74.68% on GSM8K at α = 0.5, while GSM initialization shortens responses.GSM-initialized variants reduce response length by approximately 15–24% on MATH and 21–28% on GSM8K.Source: B Cold-start Data Selection · p. 12
  • s1k-fine-tuned models gain 10.54% on GSM8K and 11.90% on MATH over GSM-fine-tuned models, but generate outputs 97.58% and 31.48% longer, respectively.The s1k responses average 7,677.43 tokens versus GSM8K’s 837.14 tokens; the passage attributes the difference partly to s1k’s DeepSeek-R1-generated responses.Source: B Cold-start Data Selection · pp. 12–13

C Prefilled Token Budget Following

  • User-defined budgets have a stable positive correlation with response lengths, enabling deterministic length control through derived transformation functions.The evaluation tested budgets from 50 to 2,000 tokens using GSM initialization and α = 0.2.Source: C Prefilled Token Budget Following · p. 13

D Adaptive Budget Allocation

  • Across all model variants, allocated token budgets show a consistent positive correlation with problem complexity; allocation is near-linear across difficulty tiers, with minimal variance across configurations.Evaluations cover GSM8K, MATH, and AIME 2024 across four architectural variants combining GSM8K-versus-s1k cold-start initialization and α values 0.2 versus 0.5.Source: D Adaptive Budget Allocation · p. 13

E Budget-Constrained Generation

SelfBudgeter’s generated lengths track token budgets, whether budgets are predicted or supplied by users. Most responses remain within the reported target-budget deviation thresholds.

E Budget-Constrained Generation

  • Generated lengths follow a linear functional relationship with user-defined budgets, indicating alignment under explicit external constraints.Source: Figure 5 · p. 14, E Budget-Constrained Generation · p. 15
  • 1.025 (95% CI [0.9466, 1.1042]) is the MATH500 predicted-budget-to-response-length slope, while 0.793 (95% CI [0.7512, 0.8354]) is the GSM8K slope.The passage interprets the near-unity MATH500 slope as an average output increase of about one token for each one-token increase in predicted budget.Source: Figure 7 · p. 15, E Budget-Constrained Generation · p. 15
  • 96% of responses deviate by no more than 50% from target budgets, and 65.40% deviate by no more than 20%.On full benchmarks, 97.65% of GSM8K and 95.82% of MATH samples satisfy the ≤50% deviation constraint; GSM8K initialization with α = 0.2 is reported as the best-performing variant.Source: E Budget-Constrained Generation · p. 15

F Case Study

The case study examines how SelfBudgeter’s shorter GSM8K responses relate to correctness, including both jointly correct answers and cases where restrictive budgets omit crucial steps. It also presents a payment example whose total matches across the methods considered.

F Case Study

  • Among 381 jointly correct samples with outputs under 70% of baseline length, 96.06% of reductions removed conversational filler and 3.94% reflected more efficient problem-solving.The analysis suggests that shorter chains usually preserve the underlying logic while removing non-essential verbalization.Source: F Case Study · p. 15
  • In cases where correctness diverged, 13.12% (173 examples) changed from baseline-wrong to compressed-correct, while 1.52% (20 examples) changed from baseline-correct to compressed-wrong.The latter errors occurred when the tighter budget omitted crucial intermediate steps.Source: F Case Study · pp. 15–16
  • The representative examples cover baseline-wrong to compressed-correct, baseline-correct to compressed-wrong, and jointly correct outcomes.Together, the findings characterize compression as generally benign or beneficial, while noting that excessively restrictive budgets can harm solution completeness.Source: F Case Study · pp. 15–16
  • The arithmetic 40 + 24 gives a total of 64.Source: F Case Study · p. 22
  • The discount calculation considers 8 discounted glasses.Source: F Case Study · p. 22
  • Kylar needs to pay 64 dollars for 16 glasses.The passage also states that the total is the same regardless of the method.Source: F Case Study · p. 22
Loading 2505.11274v6…