ThinkPrune:通过强化学习对大型语言模型的长思维链进行剪枝
ThinkPrune: Pruning Long Chain-of-Thought of LLMs via Reinforcement Learning
AI总结:
提出ThinkPrune方法,通过强化学习结合递减token限制的迭代剪枝策略,有效缩短长思考LLM的推理链长度,在AIME24上实现长度减半且性能仅降2%。
摘要:
我们提出了ThinkPrune,一种简单而有效的方法,用于对长思考LLM的思考长度进行剪枝,该方法已被发现经常产生低效且冗余的思考过程。现有减少思考长度的初步探索主要集中在强制思考过程提前退出,而不是使LLM适应并优化和巩固思考过程,因此目前观察到的长度-性能权衡是次优的。为了填补这一空白,ThinkPrune提供了一种简单的解决方案,通过强化学习(RL)持续训练长思考LLM,并增加了一个token限制,超出此限制后任何未完成的思考和答案都将被丢弃,导致零奖励。为了进一步保持模型性能,我们引入了一种迭代长度剪枝方法,其中进行多轮RL,每轮都有越来越严格的token限制。我们观察到ThinkPrune产生了显著的性能-长度权衡——在AIME24数据集上,DeepSeek-R1-Distill-Qwen-1.5B的推理长度可以减少一半,而性能仅下降2%。我们还观察到,剪枝后,LLM可以绕过不必要的步骤,同时保持核心推理过程的完整性。代码可在https://github.com/UCSB-NLP-Chang/ThinkPrune获取。
英文摘要:
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.
机构、总结、中文摘要均由AI分析而来,可能存在错误,请以论文正文为准。