TLDR: This research introduces a framework that combines reinforcement learning (RL) with supervised fine-tuning (SFT) to improve the accuracy and token efficiency of “budget forcing” in large language models (LLMs) for mathematical reasoning. Budget forcing is a test-time scaling method that allocates extra compute for thinking. The study demonstrates that their SFT+RL model significantly boosts accuracy and reduces token usage by over 40% on the GSM8K dataset, especially for smaller models, by mitigating issues like verbosity and incomplete generations caused by SFT alone.
Large Language Models (LLMs) have made remarkable strides in solving complex mathematical problems. A key area of research focuses on improving their reasoning capabilities, especially through methods that enhance performance during the inference phase, known as test-time scaling. One such method, called budget forcing, is designed to allocate additional computational resources for the model to “think” before providing a final answer, thereby encouraging its inherent self-correction abilities.
However, the traditional implementation of budget forcing often relies on supervised fine-tuning (SFT) using extensive reasoning traces. While effective, this approach can lead to performance degradation in smaller LLMs, primarily because it encourages verbose responses. Models might “overthink” a problem, backtrack unnecessarily, or even produce incomplete answers by exhausting their token budget before reaching a solution.
To address these challenges, a new research paper titled Boosting Accuracy and Efficiency of Budget Forcing in LLMs via Reinforcement Learning for Mathematical Reasoning by Ravindra Aribowo Tarunokusumoa and Rafael Fernandes Cunhaa introduces an innovative framework. This framework integrates reinforcement learning (RL) with supervised fine-tuning to significantly improve token efficiency and boost the mathematical reasoning performance of smaller LLMs, specifically demonstrated with a 1.5B model.
The Integrated Approach: SFT + RL
The proposed framework involves a three-stage pipeline:
First, the model undergoes **Supervised Fine-Tuning (SFT)**. This phase initializes the model’s reasoning behavior using a carefully curated dataset. This dataset contains reasoning steps for math problems, including examples of self-correction. Importantly, the study highlights that a modestly sized, specialized dataset with sufficiently long sample lengths is enough to instill this initial reasoning capability.
Following SFT, **Reinforcement Learning (RL)** is employed to refine the model’s reasoning strategies. The goal here is to teach the model to utilize its token budget more efficiently. The researchers adopted a novel algorithm called Group-Relative Policy Optimization (GRPO), which is more sample-efficient and less prone to “reward hacking” compared to traditional RL methods. GRPO estimates advantages by comparing each response within a group to the mean performance of that group, reducing the need for a separate value network.
Crucially, the RL phase incorporates specialized reward models:
- **Accuracy Reward Model:** This scores the correctness of the final answer. To combat verbosity, this model includes a length-dependent accuracy reward, penalizing longer correct responses and encouraging conciseness.
- **Completion Penalty:** An additional penalty is applied for incomplete generations, ensuring the model is explicitly discouraged from truncating its responses.
Budget Forcing During Inference
During the evaluation phase, budget forcing is applied as a test-time strategy. This mechanism dynamically controls the number of “thinking tokens” generated by the model. If the model is about to stop its Chain-of-Thought (CoT) reasoning too early, a special “Wait” token is injected to prompt further reasoning. Conversely, if the model exceeds a predefined token budget, an end-of-thinking delimiter is enforced to truncate the reasoning process and elicit the final answer.
Also Read:
- Bridging the Gap: How Symbolic AI Enhances Transparency and Reasoning in Large Language Models
- DeepAgent: Advancing AI with Autonomous Reasoning and Dynamic Tool Use
Key Findings and Impact
The experiments, conducted on the GSM8K dataset, yielded significant improvements:
- The SFT+RL model achieved an average accuracy gain of 14.0% over the default model and 9.2% over the SFT-only model across various thinking steps.
- It drastically reduced the average token length by 41.8% compared to the SFT model and 14.1% compared to the default model, demonstrating superior token efficiency.
- While SFT alone showed performance gains with increased inference token budget, it was notably inefficient and verbose, with a high percentage of responses being truncated. The RL component effectively mitigated these issues.
The research highlights that RL can effectively recover performance losses introduced by long-context training data in the SFT stage, leading to overall improved mathematical reasoning. Although the 1.5B model showed a performance cap, suggesting potential overfitting due to its size, the framework successfully demonstrated the power of combining SFT and RL for more accurate and efficient budget forcing.
This work paves the way for future research to explore scaling these methods to larger models and extending evaluations to broader reasoning domains, promising more capable and efficient LLMs for complex problem-solving.
