BudgetThinker: Empowering Budget-aware LLM Reasoning with Control Tokens

TL;DR

BudgetThinker introduces a control token-based framework for LLMs that enables precise management of reasoning length, balancing accuracy and resource constraints through a two-stage training process involving supervised fine-tuning and reinforcement learning.

Contribution

The paper presents a novel method for budget-aware reasoning in LLMs using control tokens and a two-stage training pipeline, improving efficiency and performance under resource constraints.

Findings

Outperforms baseline models in reasoning accuracy across various budgets.

Effectively maintains reasoning quality while adhering to token limits.

Scalable approach suitable for real-time, resource-constrained environments.

Abstract

Recent advancements in Large Language Models (LLMs) have leveraged increased test-time computation to enhance reasoning capabilities, a strategy that, while effective, incurs significant latency and resource costs, limiting their applicability in real-world time-constrained or cost-sensitive scenarios. This paper introduces BudgetThinker, a novel framework designed to empower LLMs with budget-aware reasoning, enabling precise control over the length of their thought processes. We propose a methodology that periodically inserts special control tokens during inference to continuously inform the model of its remaining token budget. This approach is coupled with a comprehensive two-stage training pipeline, beginning with Supervised Fine-Tuning (SFT) to familiarize the model with budget constraints, followed by a curriculum-based Reinforcement Learning (RL) phase that utilizes a length-aware reward function to optimize for both accuracy and budget adherence. We demonstrate that BudgetThinker significantly surpasses strong baselines in maintaining performance across a variety of reasoning budgets on challenging mathematical benchmarks. Our method provides a scalable and effective solution for developing efficient and controllable LLM reasoning, making advanced models more practical for deployment in resource-constrained and real-time environments.