Source-linked AI summary
ThinkPrune: Pruning Long Chain-of-Thought of LLMs via Reinforcement Learning
Bairu Hou, Yang Zhang, Jiabao Ji, Yujian Liu, Kaizhi Qian, Jacob Andreas, Shiyu Chang
TL;DR
Long-thinking LLMs often produce redundant reasoning, while early-exit budget forcing may yield a suboptimal length-performance tradeoff. ThinkPrune trains these models with reinforcement learning under token limits and progressively tightens those limits. It reduces reasoning length substantially while preserving performance and retaining core reasoning processes.
Problem
Long-thinking LLMs can produce inefficient, redundant reasoning, and existing approaches mainly force early exits rather than adapting models to consolidate their reasoning.
Method
ThinkPrune applies reinforcement learning with a token limit that discards over-limit unfinished outputs, followed by iterative training with increasingly stringent limits.
Results
ThinkPrune achieves a strong performance-length tradeoff, including a 50% length reduction for DeepSeek-R1-Distill-Qwen-1.5B with average performance maintained or slightly improved.
Takeaways & Limitations
After pruning, models avoid unnecessary reasoning steps while preserving key reasoning processes.
Takeaways & Limitations
ThinkPrune’s success depends on choosing a length limit that is not too stringent relative to the base model’s original output length.
Abstract
from arXiv · showhide
We present ThinkPrune, a simple yet effective method for pruning the thinking length for long-thinking LLMs, which has been found to often produce inefficient and redundant thinking processes. Existing preliminary explorations of reducing thinking length primarily focus on forcing the thinking process to early exit, rather than adapting the LLM to optimize and consolidate the thinking process, and therefore the length-performance tradeoff observed so far is sub-optimal. To fill this gap, ThinkPrune offers a simple solution that continuously trains the long-thinking LLMs via reinforcement learning (RL) with an added token limit, beyond which any unfinished thoughts and answers will be discarded, resulting in a zero reward. To further preserve model performance, we introduce an iterative length pruning approach, where multiple rounds of RL are conducted, each with an increasingly more stringent token limit. We observed that ThinkPrune results in a remarkable performance-length tradeoff -- on the AIME24 dataset, the reasoning length of DeepSeek-R1-Distill-Qwen-1.5B can be reduced by half with only 2% drop in performance. We also observed that after pruning, the LLMs can bypass unnecessary steps while keeping the core reasoning process complete. Code is available at https://github.com/UCSB-NLP-Chang/ThinkPrune.
1 Introduction
Long-thinking LLMs can generate far more tokens than necessary, and existing budget-forcing methods may waste limited budgets on redundant steps. ThinkPrune instead trains models to shorten reasoning while preserving performance.
- DeepSeek-R1-Distill-Qwen-1.5B generates over 15,000 tokens on average on MATH500, although many questions could be solved in under 1,000 tokens by regular LLMs.
- Budget-forcing can force an original model to exit before completing reasoning, whereas ThinkPrune removes unnecessary steps during training.In the illustrated 2000-token example, the budget-forced model exhausts its budget before identifying the mistake, while ThinkPrune solves the problem correctly with fewer tokens.
- The introduction frames three open questions: whether long-CoT models can be fine-tuned for shorter reasoning, the resulting length-performance tradeoff, and which reasoning steps are pruned.
- ThinkPrune reduces generation length while maintaining a strong performance tradeoff across long-thinking LLMs.
2 Related Work
Prior work improves long-CoT efficiency through training penalties, latent-vector reasoning, or short-to-long-CoT training, but reducing reasoning length in already trained long-CoT models remains less explored.
- Reinforcement learning with verifiable rewards has produced long-CoT models with deep thinking behaviors across math, coding, and browser-surfing tasks.
- Several methods reduce reasoning length by adding a length penalty during reinforcement learning.
- Other approaches represent reasoning with latent vectors instead of text tokens to shorten the reasoning process.
- Few prior works target reasoning-length reduction for models that have already been trained with long chain-of-thought.
3 Method
ThinkPrune adds length clipping to reinforcement learning so outputs exceeding a limit receive no reward, then tightens that limit iteratively to preserve performance.
- 3.1 Overview: ThinkPrune fine-tunes a long-CoT LLM to reduce output length while desirably maintaining overall performance.
- 3.2 Reinforcement Learning with Length Clipping: ThinkPrune uses the GRPO algorithm and otherwise follows a reward scheme similar to DeepSeek-R1, adding length clipping as the key difference.
- 3.2 Reinforcement Learning with Length Clipping: The method clips each sampled output to length L before reward evaluation, making over-limit outputs unable to produce a valid answer and therefore yielding zero reward.
- 3.2 Reinforcement Learning with Length Clipping: During training, a system prompt explicitly tells the model to keep its output within the length limit.
- 3.3 Iterative Length Pruning Strategy: Performance can be seriously compromised when the length limit is set too stringently relative to the base model’s original output length.
- 3.3 Iterative Length Pruning Strategy: The iterative schedule uses progressively smaller limits L1 > L2 > ··· > L* and fine-tunes each model from the preceding iteration.This gradual compaction is intended to help the LLM recover performance while shortening its reasoning chain.
- 3.3 Iterative Length Pruning Strategy: Each iteration selects a checkpoint using AIME22 and AIME23 validation data, allowing up to a relative 10% pass@1 accuracy drop and choosing the shortest qualifying output.
4 Experiment
The experiments evaluate ThinkPrune across long-reasoning models and math benchmarks, comparing one-shot and iterative length pruning with unpruned and budget-forced baselines. Results show substantial token reduction with generally modest performance loss, while iterative pruning and moderate budgets can better preserve performance.
- Experiment Setup: The evaluation uses three long-reasoning backbones, four math benchmarks, and accuracy averaged over multiple sampled responses.The models are Distill-Qwen-1.5B, DeepScaleR-1.5B-Preview, and QwQ-32B; the benchmarks are MATH-500, AIME24, AMC23, and OlympiadBench.
- One-shot Length Pruning: One-shot pruning reduces token usage by 40–50% for saturated models and up to 50% for DeepSeek-R1-Distill-Qwen-1.5B, with moderate or minimal performance degradation.Lowering the token limit from 4k to 2k consistently reduces length while slightly lowering accuracy.
- Iterative Length Pruning: Iterative pruning improves accuracy over one-shot pruning for the two 1.5B models at similar or shorter lengths, but does not recover QwQ-32B performance.On DeepScaleR, the 4k→3k→2k schedule exceeds one-shot 2k accuracy by 1.5% while using 59 fewer tokens on average; QwQ-32B loses 2.7% on average.
- Performance-Length Trade-off: Moderate pruning can reduce QwQ-32B generation length from 14K to 8K while keeping performance close to the original model, whereas further reduction beyond a critical threshold sharply lowers performance.The 4k limit performs worse than the 3k limit for QwQ-32B, and further optimization is left for future work because of computation cost.
- Inference-time Trade-off: Under inference-time budget-forcing, ThinkPrune outperforms the original model at matched thinking-token budgets and uses about 50% of the original tokens for Qwen1.5B-Distill-R1 on Math-500 at similar accuracy.ThinkPrune removes more tokens on easier Math-500 problems than on AIME-24, indicating difficulty-dependent redundancy.
- Reasoning Behavior Analysis: After pruning, hesitation and self-verification keywords decline while core computation and reasoning keywords slightly increase in frequency despite shorter responses.The analysis interprets this pattern as preserving core reasoning while removing peripheral steps.
5 Conclusion
THINKPRUNE reduces reasoning length by imposing length constraints during RL and tightening them iteratively, while aiming to preserve performance and key reasoning processes.
- THINKPRUNE imposes a length constraint during RL training, discarding unfinished thoughts and answers when sampling responses.
- Iterative pruning gradually tightens the length limit over multiple training rounds to maintain model performance.
- THINKPRUNE achieves a strong trade-off between reasoning length and performance while removing redundant steps and preserving key reasoning processes.
A.1 Implementation Details
The implementation uses training prompts with explicit token limits and modifies budget forcing to terminate thinking with a final-answer marker.
- For QwQ-32B, the training prompt instructs the model to keep its output within {N} tokens.
- Training prompts replace {N} with the length limit for DeepSeek-R1-Distill-Qwen-1.5B and DeepScaleR, such as 2,000 or 4,000 tokens.
- Budget forcing appends an end-of-thinking and final-answer sequence to stop the model's thinking process.
A.2 Analyze the Reasoning Behavior Change
The analysis segments long model-generated solutions into high-level problem-solving phases and estimates the reasoning steps devoted to each phase.
- GPT-4o identifies high-level problem-solving phases by returning the first and last reasoning steps of each chunk.
- String matching aligns phase chunks with raw model outputs, while double-newline delimiters estimate the number of reasoning steps per phase.
A.3 Additional Examples
Additional Math-500 traces show that THINKPRUNE removes repeated self-reflection steps while preserving readability and keeping the model focused on problem solving.
- Original Distill-R1-1.5B traces repeatedly check previous reasoning on simple math questions, creating redundant self-reflection steps.
- THINKPRUNE removes repeated steps and helps the model focus on problem solving while keeping perfect readability.