Boosting Accuracy and Efficiency of Budget Forcing in LLMs via Reinforcement Learning for Mathematical Reasoning

TL;DR

The paper tackles inefficiencies in budget forcing for math reasoning when applied to small LLMs trained with long-context SFT. It proposes a three-stage framework that combines supervised fine-tuning with reinforcement learning using Group-Relative Policy Optimization (GRPO) and reward shaping, plus test-time budget forcing to control thinking tokens. On a 1.5B model with only about 1.5K training samples, the approach achieves higher GSM8K accuracy and substantially reduces token usage (about $41.9\%$ fewer tokens) compared to SFT alone, demonstrating that RL can recover losses from long-context training and enable effective test-time scaling. These results suggest practical gains for deploying resource-efficient math reasoning systems without extensive post-training data.

Abstract

Test-time scaling methods have seen a rapid increase in popularity for its computational efficiency and parameter-independent training to improve reasoning performance on Large Language Models. One such method is called budget forcing, a decoding intervention strategy which allocates extra compute budget for thinking and elicits the inherent self-correcting behavior of the model. However, this relies on supervised fine-tuning (SFT) on long-context reasoning traces which causes performance degradation on smaller models due to verbose responses. For this reason, we offer a framework integrating reinforcement learning (RL) to improve token efficiency and boost the performance of a 1.5B model for mathematical reasoning. We demonstrate this using only 1.5K training samples and found that our SFT+RL model performed better on the GSM8K dataset with varying compute budgets. Our main findings showed an overall higher accuracy while significantly reducing its token usage by over 40% compared to the SFT model, revealing how RL can recover the losses due to long-context training and altogether improving performance in mathematical reasoning.

Figures (11)

• Figure 1: Example generation using budget forcing. At first, the model outputs an incorrect answer due to its faulty reasoning step. Then, the appended "Wait" token prompted the model to revise its reasoning and corrects itself.
• Figure 2: Pipeline of the framework. We start with the default model and apply supervised fine-tuning on a reasoning dataset to get our SFT model, and finally we use reinforcement learning to get our SFT+RL model. All three models will be tested using the same budget forcing inference scheme.
• Figure 3: Accuracy rewards of 16 randomly generated (correct) dummy samples, following Equation \ref{['eq:r_acc']} Samples with token length exceeding the sample mean receives a reward of $< 1.0$ and vice versa.
• Figure 4: Average completion lengths during RL training (smoothed over 50 steps). During the first 2000 steps, the model consistently hits the token budget limit, but then it decreases and settles at around 350 tokens.
• Figure 5: Zero-shot accuracy score (Left) and average tokens per sample (Right) of all three models under the Step-Level condition on the GSM8K test set. The SFT+RL model scores the highest in accuracy in all five thinking steps, while the SFT model performs similarly with the Default model in 1-2 thinking steps but gradually surpasses it. However, the SFT model outputs the most tokens while the SFT+RL model shows similar token consumption as the Default model.