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

📄 arXiv: 2504.01296v1 📥 PDF

作者: Bairu Hou, Yang Zhang, Jiabao Ji, Yujian Liu, Kaizhi Qian, Jacob Andreas, Shiyu Chang

分类: cs.CL

发布日期: 2025-04-02

备注: 15 pages, 7 figures

🔗 代码/项目: GITHUB

💡 一句话要点

ThinkPrune：通过强化学习剪枝LLM的长链式思考过程

🎯 匹配领域: 支柱二：RL算法与架构 (RL & Architecture) 支柱九：具身大模型 (Embodied Foundation Models)

关键词: 长链式思考 大语言模型 强化学习 模型剪枝 推理优化

📋 核心要点

  1. 现有方法在减少LLM思考长度时，侧重于强制提前退出，忽略了优化和巩固思考过程。
  2. ThinkPrune通过强化学习训练LLM，限制token数量，超出限制则给予零奖励，促使模型优化思考过程。
  3. 实验表明，ThinkPrune能显著减少推理长度，同时保持性能，例如在AIME24数据集上减少一半长度，性能仅下降2%。

📝 摘要（中文）

本文提出ThinkPrune，一种简单而有效的方法，用于剪枝长思考LLM的思考长度，因为研究发现它们经常产生低效和冗余的思考过程。现有的减少思考长度的初步探索主要集中在强制思考过程提前退出，而不是调整LLM来优化和巩固思考过程，因此目前观察到的长度-性能权衡并不理想。为了填补这一空白，ThinkPrune提供了一个简单的解决方案，通过强化学习（RL）持续训练长思考LLM，并增加一个token限制，超过该限制，任何未完成的想法和答案都将被丢弃，导致零奖励。为了进一步保持模型性能，我们引入了一种迭代长度剪枝方法，其中进行多轮RL，每轮都具有越来越严格的token限制。我们观察到，ThinkPrune在性能-长度权衡方面取得了显著的效果——在AIME24数据集上，DeepSeek-R1-Distill-Qwen-1.5B的推理长度可以减少一半，而性能仅下降2%。我们还观察到，剪枝后，LLM可以绕过不必要的步骤，同时保持核心推理过程的完整。

🔬 方法详解

问题定义：论文旨在解决长链式思考（Chain-of-Thought, CoT）的大语言模型（LLM）在推理过程中产生的冗余和低效问题。现有方法主要通过强制提前终止推理过程来缩短思考链，但这种方式缺乏对模型推理过程的优化，导致性能下降明显。因此，如何在保证性能的前提下，有效地缩短LLM的思考链是本文要解决的核心问题。

核心思路：ThinkPrune的核心思路是通过强化学习（RL）来训练LLM，使其学会如何在有限的token数量内完成推理任务。通过设置token限制，并对超出限制的推理过程给予零奖励，引导模型优化其思考过程，避免冗余步骤，从而在保证性能的同时，缩短推理长度。

技术框架：ThinkPrune的整体框架包含以下几个主要步骤：1) 使用预训练的LLM作为初始模型；2) 定义强化学习环境，包括状态、动作、奖励等；3) 使用强化学习算法（如策略梯度）训练LLM，使其学习在token限制下完成推理任务；4) 采用迭代长度剪枝方法，逐步降低token限制，进一步优化模型。

关键创新：ThinkPrune的关键创新在于其利用强化学习来优化LLM的思考过程，而不是简单地强制提前终止。这种方法能够使模型学会如何在有限的资源下进行高效推理，从而在性能和推理长度之间取得更好的平衡。与现有方法相比，ThinkPrune能够更好地保留模型的推理能力，并减少冗余步骤。

关键设计：ThinkPrune的关键设计包括：1) 奖励函数的设计，超出token限制给予零奖励，引导模型优化推理过程；2) 迭代长度剪枝策略，逐步降低token限制，避免模型性能大幅下降；3) 强化学习算法的选择，可以使用策略梯度等算法来训练LLM。具体的参数设置和网络结构取决于所使用的LLM和强化学习算法。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

ThinkPrune在AIME24数据集上取得了显著的性能-长度权衡。实验结果表明，使用ThinkPrune后，DeepSeek-R1-Distill-Qwen-1.5B模型的推理长度可以减少一半，而性能仅下降2%。这表明ThinkPrune能够有效地缩短LLM的思考链，同时保持较高的推理精度。此外，实验还观察到，剪枝后的LLM能够绕过不必要的步骤，保留核心推理过程。

🎯 应用场景

ThinkPrune具有广泛的应用前景，可以应用于各种需要长链式思考的LLM任务，例如数学推理、知识问答、代码生成等。通过缩短推理长度，可以降低计算成本，提高推理效率，并使LLM更易于部署在资源受限的设备上。此外，ThinkPrune还可以用于提高LLM的可解释性，通过分析剪枝后的推理过程，可以更好地理解模型的推理逻辑。

📄 摘要（原文）

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.