Table of Contents
Fetching ...

Beyond Token Length: Step Pruner for Efficient and Accurate Reasoning in Large Language Models

Canhui Wu, Qiong Cao, Chang Li, Zhenfang Wang, Chao Xue, Yuwei Fan, Wei Xi, Xiaodong He

TL;DR

The paper tackles overthinking and inefficiency in large reasoning models by shifting the optimization objective from token count to the number of reasoning steps. It introduces Step Pruner (SP), a step-based RL framework that uses a paragraph-based segmentation proxy and a step-aware reward, paired with a dynamic stopping rule to prevent reward hacking. SP leverages GRPO for policy optimization and demonstrates state-of-the-art accuracy with substantially shorter outputs across four benchmarks (AIME24, MATH500, GSM8K, GPQA), including a ~69.7% reduction in token usage on AIME24. The work also provides ablations and semantic analyses showing that SP reduces redundant reasoning and yields more focused, pivotal reasoning paths, with robust performance across model scales and domains.

Abstract

Large Reasoning Models (LRMs) demonstrate strong performance on complex tasks but often suffer from excessive verbosity, known as "overthinking." Existing solutions via reinforcement learning (RL) typically penalize generated tokens to promote conciseness. However, these methods encounter two challenges: responses with fewer tokens do not always correspond to fewer reasoning steps, and models may develop hacking behavior in later stages of training by discarding reasoning steps to minimize token usage. In this work, we introduce \textbf{Step Pruner (SP)}, an RL framework that steers LRMs toward more efficient reasoning by favoring compact reasoning steps. Our step-aware reward function prioritizes correctness while imposing penalties for redundant steps, and withholds rewards for incorrect responses to prevent the reinforcement of erroneous reasoning. Moreover, we propose a dynamic stopping mechanism: when the model's output no longer shortens, training is halted to prevent hacking behavior caused by the merging of steps. Extensive experiments across four reasoning benchmarks demonstrate that SP achieves state-of-the-art accuracy while significantly reducing response length. For instance, on AIME24, SP reduces token usage by \textbf{69.7\%}.

Beyond Token Length: Step Pruner for Efficient and Accurate Reasoning in Large Language Models

TL;DR

The paper tackles overthinking and inefficiency in large reasoning models by shifting the optimization objective from token count to the number of reasoning steps. It introduces Step Pruner (SP), a step-based RL framework that uses a paragraph-based segmentation proxy and a step-aware reward, paired with a dynamic stopping rule to prevent reward hacking. SP leverages GRPO for policy optimization and demonstrates state-of-the-art accuracy with substantially shorter outputs across four benchmarks (AIME24, MATH500, GSM8K, GPQA), including a ~69.7% reduction in token usage on AIME24. The work also provides ablations and semantic analyses showing that SP reduces redundant reasoning and yields more focused, pivotal reasoning paths, with robust performance across model scales and domains.

Abstract

Large Reasoning Models (LRMs) demonstrate strong performance on complex tasks but often suffer from excessive verbosity, known as "overthinking." Existing solutions via reinforcement learning (RL) typically penalize generated tokens to promote conciseness. However, these methods encounter two challenges: responses with fewer tokens do not always correspond to fewer reasoning steps, and models may develop hacking behavior in later stages of training by discarding reasoning steps to minimize token usage. In this work, we introduce \textbf{Step Pruner (SP)}, an RL framework that steers LRMs toward more efficient reasoning by favoring compact reasoning steps. Our step-aware reward function prioritizes correctness while imposing penalties for redundant steps, and withholds rewards for incorrect responses to prevent the reinforcement of erroneous reasoning. Moreover, we propose a dynamic stopping mechanism: when the model's output no longer shortens, training is halted to prevent hacking behavior caused by the merging of steps. Extensive experiments across four reasoning benchmarks demonstrate that SP achieves state-of-the-art accuracy while significantly reducing response length. For instance, on AIME24, SP reduces token usage by \textbf{69.7\%}.

Paper Structure

This paper contains 37 sections, 39 equations, 9 figures, 6 tables.

Figures (9)

  • Figure 1: Comparison of Step Pruner with seven baselines on the MATH500 dataset. SP uses only 33% of the tokens compared to DeepSeek-R1-Distilled-Qwen-7B, and maintains the same level of accuracy.
  • Figure 2: Pipeline of Step Pruner: (1) Prompt LRMs to generate multiple responses; (2) Split each response into logical steps; (3) Calculate a reward score based on both correctness and the number of steps; (4) Update LRMs using GRPO.
  • Figure 3: Training curves showing paragraph length, paragraph count, and total response length over training steps.
  • Figure 4: Comparison of different segmentation methods during training: output length (left) and accuracy change (right) on MATH500. The paragraph-based segmentation achieved the best trade-off between accuracy and output length.
  • Figure 5: The Semantic Analysis of Reasoning in R1-Qwen-7B and SP-7B.
  • ...and 4 more figures