PaperlayerView exact revision on arXiv

Source-linked AI summary

Token-Budget-Aware LLM Reasoning

Tingxu Han, Zhenting Wang, Chunrong Fang, Shiyu Zhao, Shiqing Ma, Zhenyu Chen

arXiv:2412.18547v5cs.CLcs.AIcs.LG

TL;DR

CoT improves LLM reasoning but adds substantial token overhead and cost, motivating the question of whether current reasoning is unnecessarily lengthy. The paper proposes TALE, which dynamically adjusts reasoning-token budgets by problem complexity, and reports major token reductions with limited accuracy loss. The framework is currently evaluated only for text-input and text-output LLMs.

  • Problem

    CoT improves reasoning but produces substantial additional tokens, increasing inference resource use, running time, monetary cost, and energy cost.

  • Method

    TALE dynamically adjusts reasoning-token budgets by problem complexity through budget estimation and prompting or post-training internalization.

  • Results

    TALE-EP reduces token usage by 67% with less than a 3% accuracy decrease, while TALE-PT cuts usage by around 50% versus Vanilla CoT with competitive performance.

  • Takeaways & Limitations

    TALE reduces output token usage and expense while balancing token efficiency and answer correctness across various LLMs.

  • Takeaways & Limitations

    The experiments focus on text-input and text-output LLMs and do not account for models producing multimodal outputs such as interleaved images and text.

Abstract

from arXiv · show

Reasoning is critical for large language models (LLMs) to excel in a wide range of tasks. While methods like Chain-of-Thought (CoT) reasoning and enhance LLM performance by decomposing problems into intermediate steps, they also incur significant overhead in token usage, leading to increased costs. We find that the reasoning process of current LLMs is unnecessarily lengthy and it can be compressed by including a reasonable token budget in the prompt, but the choice of token budget plays a crucial role in the actual compression effectiveness. We then propose a token-budget-aware LLM reasoning framework that dynamically adjusts the number of reasoning tokens based on the reasoning complexity of each problem. Experiments show that our method effectively reduces token costs in CoT reasoning with only a slight performance reduction, offering a practical solution to balance efficiency and accuracy in LLM reasoning. Code: https://github.com/GeniusHTX/TALE

1 Introduction

CoT improves LLM reasoning by decomposing problems into intermediate steps, but its detailed outputs increase token, computational, monetary, and energy costs. The paper finds that token budgets can compress reasoning and proposes dynamically adapting budgets to problem complexity.

  • Motivation: CoT guides LLMs through intermediate steps, improving answer reliability across tasks.The paper describes CoT as prompting models to decompose questions and solve each step sequentially.
  • Motivation: Detailed CoT reasoning substantially increases token usage, inference time, computational resources, monetary costs, and energy costs.The added cost comes from intermediate reasoning steps included in the output.
  • Token budgeting: A reasonable token budget reduces CoT output costs, but compression effectiveness depends critically on the chosen budget.The paper motivates searching for the smallest budget that preserves correctness and minimizes actual token usage.
  • Token budgeting: 258 output tokens fell to 86 with a 50-token budget, while a 10-token budget still produced 157 tokens.The smaller budget was less effective because models often exceed very small specified budgets, a phenomenon termed Token Elasticity.
  • Contribution: TALE dynamically adjusts reasoning-token counts to problem complexity through estimation-and-prompting or post-training implementations.TALE-EP uses zero-shot budget estimation, whereas TALE-PT internalizes budget awareness during post-training.
  • Contribution: TALE-EP reduces token usage by 67% with less than a 3% accuracy decrease, while TALE-PT cuts usage by around 50% versus Vanilla CoT.TALE-PT achieves competitive performance according to the introduction’s reported results.

2 Related Work

Prior work develops LLM reasoning techniques, especially CoT, while token-budget-aware prompting is illustrated as a related approach to controlling reasoning length.

  • LLM Reasoning: LLM reasoning techniques generate intermediate steps to improve performance across diverse domains.Related work includes variational reasoning and knowledge distillation approaches.
  • LLM Reasoning: CoT prompting improves reasoning by breaking problems into intermediate steps, including on benchmarks such as GSM8K.The related-work discussion positions CoT as a key technique among broader reasoning methods.
  • Token Budget: Table 1 illustrates the difference between vanilla CoT prompts and prompts that include an explicit token budget.This comparison motivates the paper’s focus on controlling reasoning length through prompting.

3 Token Redundancy in LLM Reasoning

The paper identifies token redundancy in CoT reasoning and studies explicit token budgets as a way to control output length while retaining correct answers.

  • Token Budget: A 50-token budget produces 86 output tokens, showing that LLMs can partially follow explicit length constraints.The example contrasts vanilla CoT with CoT prompted using a token budget.
  • Token Redundancy Phenomenon: Token-budget prompting can reduce reasoning length, but the budget must be chosen carefully for effective compression.The paper identifies this dependence on budget choice as part of the token redundancy phenomenon.

4 Searching Optimal Token Budget

The paper searches for token budgets that preserve answer correctness while reducing reasoning cost, then accounts for token elasticity when budgets become too tight. It uses binary and greedy feasibility searches to identify cost-effective budgets and motivates TALE’s estimation-based workflow.

  • Vanilla Method for Optimal Budget Search: Binary search selects the shortest budget that preserves correctness under an implicit monotonicity assumption.A budget is feasible when the model remains correct; the search uses vanilla CoT token cost as the right boundary.
  • Vanilla Method for Optimal Budget Search: The implicit monotonicity assumption may fail in some instances but remains a reasonable and effective approximation for guiding budget search.
  • Observation of Token Elasticity: When budgets fall below a reasonable range, token costs increase because the model abandons the tight constraint and returns to longer reasoning.This token elasticity produces non-monotonic costs: usage first decreases as budgets tighten, then rebounds when budgets become too small.
  • Token Elasticity based Optimal Budget Search: Greedy feasibility adds a lower-token-cost requirement to correctness, seeking budgets that reduce cost relative to the previously searched budget.The updated feasibility condition compares actual token costs under the current and previous budgets.
  • TALE-EP Workflow: TALE-EP estimates a reasonable budget with zero-shot prompting, constructs a budget-aware prompt, and sends it to the LLM without iterative search at inference time.Budget search is used offline for TALE-PT target generation, while TALE-EP uses a lightweight estimator directly at inference.

5 Methodology

TALE provides two token-budget-aware implementations: estimating and prompting a budget for each problem, or internalizing budget awareness through post-training. TALE-PT searches optimal-budget outputs and trains with SFT or DPO, while TALE-EP constructs prompts using estimated budgets.

  • TALE offers estimation-and-prompting (TALE-EP) and post-training (TALE-PT) solutions for token-budget-aware reasoning.
  • The method searches for a budget in a reasonable range because only a well-chosen budget effectively minimizes token costs while preserving performance.The searched budget is used to create concise target outputs that retain correctness.
  • TALE-EP: TALE-EP estimates a reasonable budget for each question and incorporates it into a token-budget-aware reasoning prompt.Its workflow uses budget estimation followed by prompt construction to align with an optimal searched budget.
  • TALE-PT: TALE-PT searches target outputs under optimal budgets, builds a specialized dataset, and post-trains the LLM to internalize token-budget awareness.The workflow has target-output generation followed by LLM post-training.
  • TALE-PT: SFT trains on searched target outputs by minimizing cross-entropy, while DPO ranks those outputs above vanilla-CoT outputs.The target outputs serve as positive samples for DPO, with vanilla-CoT outputs as negative samples.

6 Evaluation

Across mathematical benchmarks and multiple LLMs, TALE-EP substantially reduces output-token and expense costs while retaining competitive accuracy. TALE-PT likewise reduces tokens, with SFT achieving the strongest reported accuracy and DPO emphasizing efficiency on GSM8K-Zero.

  • TALE-EP: TALE-EP achieves 80.22% average accuracy, 138.53 average output tokens, and 118.46 average expense on GPT-4o-mini.
  • TALE-EP: TALE-EP reduces token costs by 68.64% on average while achieving 81.03% accuracy, compared with Vanilla CoT’s 83.75%.On GSM8K, TALE-EP reaches 84.46% accuracy and surpasses Vanilla CoT.
  • TALE-EP: TALE-EP reduces output tokens by 64.63% and expenses by 45.30% on average across Yi-lightning, GPT-4o-mini, GPT-4o, and o3-mini.Accuracy remains comparable to Vanilla CoT for the reported models, with the largest drop observed for GPT-4o-mini.
  • TALE-PT: TALE-PT-SFT achieves 78.57% and 78.43% accuracy, while TALE-PT-DPO cuts GSM8K-Zero token consumption by over 50% versus Vanilla CoT.The reported accuracy pairs correspond to GSM8K and GSM8K-Zero, respectively.

7 Conclusion

TALE reduces redundancy in Chain-of-Thought reasoning by dynamically adjusting reasoning tokens to problem complexity. Experiments report lower token usage and expense with acceptable accuracy loss and generalization across LLMs.

  • TALE dynamically adjusts reasoning-token counts according to each problem’s reasoning complexity.
  • Experiments show that TALE significantly reduces output-token usage and expense with acceptable accuracy loss.
  • TALE outperforms Vanilla CoT in cost-effectiveness and generalizes across various LLMs.

8 Limitations

The experiments focus on text-only input and output, excluding models that generate multimodal content such as interleaved images and text.

  • The experiments evaluate LLMs with text-only input and output.
  • Multimodal output content, including interleaved images and text, is not covered.Future work proposes modality-specific budget constraints and adaptive strategies for images and videos.

A.1 Definition of Ideal Budget Range

The ideal budget range is the region of feasible budgets where actual token consumption is minimized. TALE seeks this range without searching every budget, while allowing an acceptable sacrifice.

  • The ideal budget range is the region where actual token costs approach their lowest bound.Token costs increase before or after this bottom range.
  • Among budgets that maintain answer correctness, the ideal range minimizes actual token consumption.
  • The workflow illustrates TALE-EP using an example question and a 68-output-token TALE response.The supplied example contrasts direct answering, CoT prompting, and TALE.
  • TALE aims to estimate a budget in the ideal range without a search process.The estimated budget may involve an acceptable sacrifice.

A.2 Effectiveness of Budget Estimation.

Budget estimation is evaluated by whether predicted budgets fall within the ideal range and by their distance from it. The estimator achieves 60.61% in-range accuracy with an average out-of-range distance of 109.64 tokens.

  • The evaluation defines in-range accuracy by whether the predicted budget falls within the ideal budget range.
  • Out-of-range distance measures the distance between a predicted budget and the ideal budget range when the prediction falls outside it.
  • 60.61% in-range accuracy indicates that predicted budgets fall within the ideal range in the evaluation.
  • 109.64 tokens is the average offset for estimated budgets outside the ideal range.Higher in-range accuracy and lower out-of-range distance indicate better estimation.
  • The estimation prompt asks the model to estimate the minimum number of reasoning tokens required for a question.It evaluates reasoning depth, structure, completeness, and surface-level difficulty.

A.3 Details of TALE’s Implementation

TALE-EP estimates budgets with an additional query, while TALE-PT uses post-training; experiments examine efficiency, accuracy, latency, generalization, larger budgets, and the search assumption.

  • Implementation: TALE-EP uses zero-shot budget estimation followed by prompting, whereas TALE-PT uses post-training implementations including SFT and DPO.
  • TALE-EP versus TALE-PT: TALE-EP uses fewer tokens but slightly lower accuracy than TALE-PT, which achieves higher accuracy with more tokens on GSM8K.The comparison frames a trade-off between token efficiency and strict correctness preservation.
  • Generalization: TALE-EP achieves comparable or better BLEU scores than Vanilla CoT while using around 40% of output tokens on three open-ended generative tasks.
  • Latency: 2.3 seconds per sample versus 10.2 seconds for Vanilla CoT shows lower end-to-end latency for TALE-EP on GSM8K-Zero.The measurement includes TALE-EP’s additional budget-estimation query.
  • Larger token budgets: 67.33% to 72.66% accuracy and 210.97 to 279.78 tokens show the effect of scaling the estimated budget from α=1 to α=2.The larger budget increases both accuracy and token usage.
  • Search assumption: The search procedure is empirically evaluated against an implicit monotonicity assumption across scaled budget values.
Loading 2412.18547v5…