arXiv:2504.01296cs.CL2025-04被引 146

用强化学习精简大模型思维链,省一半步骤仅损失2%性能。

ThinkPrune: Pruning Long Chain-of-Thought of LLMs via Reinforcement Learning

  • 通过强化学习训练模型在限定字数内完成思考,超限则得零分。
  • 在AIME24上将推理长度减半,性能仅降2%。
  • 能自动跳过冗余步骤,保留核心推理逻辑,适合高效推理场景。

我们提出ThinkPrune,一种简化长思维链大模型的方法,该模型常产生低效冗余的思考过程。现有方法多强制提前退出,未能优化和整合思考过程,导致长度与性能权衡不理想。ThinkPrune通过强化学习持续训练长思考模型,在设定的词元上限外未完成的思考和答案将被丢弃,获得零奖励。为更好保持性能,采用迭代式长度剪枝策略,逐步收紧词元限制进行多轮强化学习。实验显示,ThinkPrune实现显著的性能-长度权衡:在AIME24数据集上,DeepSeek-R1-Distill-Qwen-1.5B的推理长度可减少50%,性能仅下降2%。剪枝后模型能跳过不必要的步骤,同时保持核心推理完整。代码已公开于https://github.com/UCSB-NLP-Chang/ThinkPrune。

原文摘要 · Abstract (English)

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.

思维链强化学习模型剪枝推理优化

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。