PaperlayerView exact revision on arXiv

Source-linked AI summary

Training Language Models to Reason Efficiently

Daman Arora, Andrea Zanette

arXiv:2502.04463v4cs.LGcs.CL

TL;DR

Long chain-of-thought reasoning raises inference costs, motivating methods that reduce token usage without sacrificing accuracy. The paper uses reinforcement learning with a tunable length penalty to derive more efficient reasoning models. Experiments show substantial token reductions with small accuracy losses, while the approach remains subject to imperfect length targeting and generally small performance degradation.

  • Problem

    Long chain-of-thought reasoning increases inference costs, creating a need for computationally efficient reasoning models.

  • Method

    Policy-gradient reinforcement learning rewards correct solutions with shorter chain-of-thoughts, with α controlling the efficiency level.

  • Results

    Across two open-weight reasoning models, the method substantially reduces reasoning cost while approximately maintaining accuracy; the 7B model cuts tokens by 50% with less than 5% accuracy reduction.

  • Takeaways & Limitations

    A single model can adapt its inference budget to problem difficulty, offering controllable deployment-cost reductions without separately training models for each trade-off.

  • Takeaways & Limitations

    The penalty coefficient does not precisely target a required generation length, and length regularization generally brings a small performance loss.

Abstract

from arXiv · show

Scaling model size and training data has led to great advances in the performance of Large Language Models (LLMs). However, the diminishing returns of this approach necessitate alternative methods to improve model capabilities, particularly in tasks requiring advanced reasoning. Large reasoning models, which leverage long chain-of-thoughts, bring unprecedented breakthroughs in problem-solving capabilities but at a substantial deployment cost associated to longer generations. Reducing inference costs is crucial for the economic feasibility, user experience, and environmental sustainability of these models. In this work, we propose to train large reasoning models to reason efficiently. More precisely, we use reinforcement learning (RL) to train reasoning models to dynamically allocate inference-time compute based on task complexity. Our method incentivizes models to minimize unnecessary computational overhead while maintaining accuracy, thereby achieving substantial efficiency gains. It enables the derivation of a family of reasoning models with varying efficiency levels, controlled via a single hyperparameter. Experiments on two open-weight large reasoning models demonstrate significant reductions in inference cost while preserving most of the accuracy.

1 Introduction

The paper trains reasoning models to reduce inference-time token usage while preserving accuracy, producing controllable efficiency–performance trade-offs across model scales and tasks.

  • Motivation: Long chain-of-thought reasoning increases inference costs through quadratic attention costs and linear KV-cache growth.These costs affect profitability, latency, responsiveness, and deployment sustainability.
  • Approach: Policy-gradient reinforcement learning rewards correct solutions with fewer tokens, with α controlling the reduction in inference compute.Varying α derives a family of models with increasingly shorter chain-of-thoughts.
  • Results: 50% fewer tokens for the 7B model incurs less than a 5% accuracy reduction.This result is aggregated across the evaluation setting described in the paper.
  • Results: 16% fewer tokens on AIME2024 accompanies a 3.3% points accuracy drop, while MATH500 and GSM8K reach 37% and 65% reductions with 2.2% and 1.7% accuracy drops.The achievable reduction depends on problem hardness.
  • Practicality: The procedure uses only 100 RL steps, approximately 200 gradient updates, making efficient-reasoning training viable with modest academic resources.The method requires only a couple of lines of changes to a standard reinforcement-learning implementation.

2 Related Work

Prior work improves reasoning by spending more inference compute or shortening chain-of-thoughts through heuristics, prompting, or alternative training, while this method targets tunable RL-based efficiency.

  • More inference compute: Chain-of-thought, self-consistency, search, process rewards, and self-correction improve performance by spending more inference-time compute.These methods represent the compute-increasing direction of related work.
  • Large reasoning models: Frontier reasoning models rely on long monolithic chain-of-thoughts and large-scale reinforcement learning to develop advanced reasoning abilities.The proposed approach seeks to make such models more efficient.
  • Efficient serving: Speculative decoding and batch engines accelerate inference at the system level and can be combined directly with this model-based method.Existing efficient-serving techniques are described as orthogonal to the paper's approach.
  • Concurrent work: Concurrent methods mitigate overthinking with preference data or shortened chain-of-thoughts, but do not easily tune models to a user's compute budget.The paper positions controllable compute allocation as a distinction from these approaches.
  • Efficient chain-of-thought: Earlier methods use conditional training or prompt engineering for shorter chain-of-thoughts, without RL training of long-CoT models for reduced verbosity.The paper also notes that these methods rely on an assumption about models following token limits.

3 Setup

The setup models autoregressive responses, evaluates mathematical correctness with a verifier, and frames reasoning as reinforcement learning over sampled complete outputs.

  • Response generation: A language model generates a response y=(y_1,...,y_t) autoregressively, conditioning each next token on the prompt and previously generated tokens.Generation ends when the model emits the EOS token.
  • Reinforcement-learning pipeline: The method samples multiple solutions for each prompt, rewards correctness and response length, and updates the language model with policy gradients.The shortest correct answers receive the highest rewards.
  • Chain-of-thought: A chain-of-thought consists of intermediate tokens produced before the final answer, with y denoting the model's full output.The final answer is typically not formally separated from the chain-of-thought.
  • Reward and correctness: For math problems, the scoring function tests whether the generated final answer matches the gold answer rather than whether complete output strings are identical.Because y includes the chain-of-thought, correctness concerns the final answer within the full response.

4 Method

The method adds a length penalty to reinforcement learning, normalizes and soft-clips response lengths, and theoretically preserves accuracy while favoring shortest correct solutions under simplified assumptions.

  • Method: The objective encourages models to solve simple problems with few tokens while retaining the ability to spend thousands of tokens on difficult AIME problems.This supports adaptive inference-time compute based on problem difficulty.
  • Length regularization: Increasing α strengthens regularization toward shorter correct responses, while α=0 recovers the standard reinforcement-learning objective.The coefficient controls the compute–accuracy trade-off rather than specifying an exact generation length.
  • Length normalization: Per-prompt length normalization avoids disproportionately penalizing long chains on hard problems, and sigmoid soft-clipping bounds the objective between 0 and 1.Correct responses remain preferred to incorrect ones even when they are long.
  • Optimization: The non-differentiable objective is optimized with reinforcement learning, using PPO with an RLOO Monte Carlo advantage estimator.RLOO avoids maintaining a separate value network for variance reduction.
  • Theoretical guarantees: Under the simplified population-level setup, the objective yields the shortest correct response for every prompt while preserving accuracy.The guarantees rely on assumptions including tabular representation and coverage of a correct response.
  • Theoretical guarantees: The idealized analysis preserves accuracy and reduces average length, but the practical length function depends on sampling-derived statistics such as MEAN(x) and STD(x).Finite-sample guarantees are deferred to standard statistical techniques.

5 Experiments

Experiments evaluate the method on two open-weight reasoning models across mathematical benchmarks, comparing compute–accuracy trade-offs, baselines, and changes in chain-of-thought behavior. The results show controllable token reduction, stronger savings on easier problems, and a compression–faithfulness trade-off.

  • Experimental setup: Experiments evaluate DeepSeek-R1-Distill-Qwen-1.5B and DeepSeek-R1-Distill-Qwen-7B on GSM8K, MATH500, and AIME2024, alongside simple inference-cost baselines.The study also uses 3.2k numerical-answer prompts for post-training and reports average pass rate@k.
  • Performance on the test sets: 36% lower MATH chain-of-thought length accompanies a 2.2% accuracy loss at α = 0.1 for the 7B model.The chain-of-thought decreases from approximately 4,000 to 2,600 tokens.
  • Performance on the test sets: 27% lower AIME token usage accompanies a 4% accuracy drop at α = 0.2 for the 7B model.Token usage falls from approximately 13,000 to 9,000 compared with DeepSeek-R1-Distill-Qwen-7B.
  • Performance on the test sets: The method is more performant than SFT and DPO at the same token usage, while α provides controllable response-length reductions.Increasing α produces greater response-length decreases, but lower generation length is also accompanied by lower performance.
  • Performance on the test sets: 83% token savings on GSM8K versus 27% on AIME2024 at α = 0.2 show greater compression on easier problems.The authors hypothesize that reasoning models are more wasteful on easier problems because those tasks have a larger relative spread in response lengths.
  • Qualitative and quantitative changes to the Chain-of-Thought: Greater compression generally reduces verification, backtracking, and exploration, while faithfulness also declines as compression increases.The reported faithfulness scores remain higher for compressed reasoning models than for nonreasoning models, at 0.480 versus 0.301.

6 Limitations and Conclusions

The method substantially reduces reasoning-model inference cost while minimally affecting accuracy, but its reinforcement-learning setup is more involved and its length control is imprecise.

  • The reinforcement-learning optimization is more involved than SFT- or DPO-derived techniques.
  • The penalty coefficient α affects generation cost but does not precisely target a desired generation length.
  • Length-penalty regularization is generally accompanied by a small performance loss, leaving better performance at reduced compute as an open question.
  • The approach reduces inference cost while minimally affecting accuracy and lets one model adapt its inference budget to problem difficulty.

A Implementation Details

The implementation uses OpenRLHF-based training and demonstrates how the trained model produces shorter reasoning on a simple arithmetic prompt.

  • Implementation Details: Training uses OpenRLHF with 4 GH200 GPUs for the 1.5B model and 8 GH200 GPUs across two nodes for the 7B model.
  • Implementation Details: The shared prompt asks the model to reason step by step and place its final answer within \boxed{}.
  • Example: On the 1+1 prompt, the distilled model generates extended reasoning through multiple explanatory perspectives before answering 2.
  • Example: After training with α = 0.4, the model solves 1 + 1 with a short three-step rationale and outputs 2.

B.2 Example 2

The currency-conversion example contrasts verbose and concise reasoning: both derive that Kathy has 22 more euros than pounds, while the trained model presents the answer more directly.

  • Example 2: The problem divides Kathy’s $300 equally between pounds and euros using exchange rates of 1.64 USD and 1.32 USD, respectively.
  • Example 2: The detailed calculation converts $150 to approximately 91.46 pounds and 113.64 euros.
  • Example 2: Subtracting the converted amounts gives approximately 22.18, which rounds to 22.
  • Example 2: The concise response states that Kathy will have 22 more euros than pounds.
  • Baseline comparison: The vLLM-cutoff baseline performs poorly relative to models achieving higher performance with lower compute requirements.

C Extended Results

Extended results cover five benchmarks, compare multiple baselines, and examine how α changes normalized accuracy and token usage for 1.5B and 7B models.

  • Baseline comparison: Compared with O1-Pruner, the method produces higher-accuracy models at lower token cost, specifically for α = 0.05 and α = 0.1.
  • Sensitivity analysis: Figure 6 aggregates five datasets and uses three seeds to show normalized accuracy and token usage as functions of α for both model sizes.
  • Evaluation setup: The evaluation uses a shared step-by-step prompt template that places each final answer within \boxed{}.

E.1 Proof of Theorem 4.3

The proof establishes properties of the efficiency-regularized objectives by showing that population maximizers place probability only on correct answers under the stated conditions.

  • Proof of Theorem 4.3: Lemma E.1 states that if a prompt has a correct answer, the population maximizer places no mass on incorrect answers.The proof argues that shifting mass from an incorrect answer to a correct one increases the objective.
  • Proof of Theorem 4.3: The contradiction proof constructs a modified distribution that concentrates more mass on positive, correct-answer terms.This distribution would improve the objective, contradicting optimality of the original maximizer.
  • Proof of Theorem 4.3: Theorem E.1 is applied to conclude that the effective policy has support only on correct answers.The conclusion follows once the theorem's premise—that a correct answer exists—is satisfied.
  • Proof of Theorem 4.3: Theorem E.1 assumes a monotonically increasing length function bounded in [0, 1] and α ∈ [0, 1).These conditions are used in the contradiction argument for the efficiency-regularized objective.
  • Length-constrained prompting: Prompting distilled models with explicit token limits does not reliably control their generated length on MATH500.The experiment varies requested limits from 256 to 4096 tokens for 1.5B and 7B models.
  • Additional analysis: The appendix reports normalized token and accuracy changes for experiments applying fixes to the 7B distilled model.Table 16 defines ΔNT as normalized-token change and ΔNA as normalized-accuracy change relative to baseline.

H Training dynamics

Training with a length penalty reduces token usage while preserving initial training accuracy, with the effect varying across α values and dataset difficulty.

  • Training dynamics: Up to 50% fewer tokens are generated as α increases, while training accuracy remains at its initial level.The α = 0 condition corresponds to RL without a length penalty.
  • Training dynamics: Figure 12 tracks smoothed training accuracy and generation length for the 7B model across α values.Both curves use running averages over 25 training iterations.
  • Training dynamics: Table 17 reports relative spread across GSM8K, MATH500, and AIME2024 for the distilled 7B model.The aggregation includes prompts with at least one correct response among eight generations.

I Difficulty based analysis

Length reductions are larger on easier datasets, which exhibit greater relative variation in response lengths and therefore more token waste.

  • Difficulty-based analysis: Relative spread is larger for GSM8K than for MATH500 and AIME2024 in the distilled 7B model.The analysis computes normalized standard deviation and normalized mean-minus-minimum response length.

J Length penalty when training a base model.

In a base-model Countdown experiment, length regularization slows convergence and lowers peak performance, but produces shorter responses at matched performance and more token-efficient solutions.

  • Figure guide: Figure 13 compares training accuracy and generation lengths for vanilla and length-regularized RL on a base Countdown model.The figure provides the training-dynamics view underlying the reported convergence and length comparison.
  • Motivation and setup: Base models generally lack long chain-of-thoughts before RLVR training, motivating the preliminary Countdown study.The experiment compares RLOO with and without a length penalty under a 2048-token context limit.
  • Training dynamics: Length-regularized RLOO reaches peak performance at about 400 RL steps, versus about 220 steps for vanilla RLOO.Vanilla RLOO later saturates because of the 2048-token context-window limit.
  • Training dynamics: Vanilla RLOO reaches 67.48% validation performance, whereas length-regularized RLOO reaches 59.68%.The comparison uses Qwen2.5-3B with α = 0.1.
  • Efficiency at matched performance: At matched training accuracy, length-regularized RLOO uses substantially shorter responses than vanilla RLOO.Vanilla response length grows significantly as performance improves; only one hyperparameter combination was tested.
  • Interpretation: The toy setup finds that length regularization slows convergence and reduces peak performance but discovers more token-efficient solutions at matched performance.The authors therefore expect the method to be more effective as post-training than for RLVR from scratch.
Loading 2502.04463v4…