2SSP: A Two-Stage Framework for Structured Pruning of LLMs

📄 arXiv: 2501.17771v2 📥 PDF

作者: Fabrizio Sandri, Elia Cunegatti, Giovanni Iacca

分类: cs.CL, cs.AI, cs.LG

发布日期: 2025-01-29 (更新: 2025-08-17)

备注: Published in Transactions on Machine Learning Research (TMLR)

🔗 代码/项目: GITHUB

💡 一句话要点

提出2SSP框架，结合宽度和深度剪枝，高效压缩大型语言模型。

🎯 匹配领域: 支柱九：具身大模型 (Embodied Foundation Models)

关键词: 大型语言模型 结构化剪枝 宽度剪枝 深度剪枝 模型压缩 稀疏性 Transformer

📋 核心要点

  1. 大型语言模型参数众多，计算开销大，现有剪枝方法难以在效率和性能间取得平衡。
  2. 2SSP框架结合宽度和深度剪枝，先剪枝神经元，再剪枝注意力模块，并平衡两阶段稀疏率。
  3. 实验表明，2SSP在多个LLM和数据集上优于现有方法，剪枝速度提升显著。

📝 摘要（中文）

本文提出了一种用于大型语言模型（LLMs）结构化剪枝的新型两阶段框架（ extsc{2SSP}），它结合了两种不同的剪枝策略，即宽度剪枝和深度剪枝。第一阶段（宽度剪枝）移除整个神经元，从而移除它们对应的行和列，旨在保持每个Transformer块中前馈网络中间状态下剪枝结构之间的连接。这是基于一个重要性分数完成的，该分数衡量每个神经元对输出幅度的影响。第二阶段（深度剪枝）则移除整个注意力子模块。这是通过一个迭代过程完成的，该过程移除对给定感兴趣指标（在我们的例子中是困惑度）影响最小的注意力模块。我们还提出了一种新的机制来平衡两个阶段的稀疏率，使其与所需的全局稀疏度相关。我们在四个LLM家族和三个稀疏率（25％，37.5％和50％）上测试了 extsc{2SSP}，测量了三个语言建模数据集上的困惑度以及六个下游任务的性能。我们的方法在三个语言建模和六个下游任务上始终优于五个最先进的竞争对手，在剪枝时间方面获得了高达两个数量级的提升。代码可在https://github.com/FabrizioSandri/2SSP获得。

🔬 方法详解

问题定义：论文旨在解决大型语言模型（LLMs）的结构化剪枝问题。现有方法在剪枝过程中，要么效率低下，需要大量的计算资源和时间，要么剪枝后的模型性能下降明显，无法满足实际应用的需求。因此，如何在保证模型性能的前提下，高效地进行LLM的结构化剪枝是一个关键挑战。

核心思路：论文的核心思路是将剪枝过程分解为两个阶段：宽度剪枝和深度剪枝。宽度剪枝通过移除不重要的神经元来减少模型的参数量，深度剪枝则通过移除整个注意力子模块来进一步压缩模型。这种两阶段的设计允许更精细的控制剪枝过程，从而在性能和效率之间取得更好的平衡。

技术框架：2SSP框架包含两个主要阶段：1) 宽度剪枝：计算每个神经元的重要性得分，并移除得分较低的神经元，保持剪枝后结构的连接性。2) 深度剪枝：迭代地移除对性能影响最小的注意力子模块。此外，该框架还包含一个稀疏率平衡机制，用于调整两个阶段的剪枝比例，以达到期望的全局稀疏度。

关键创新：2SSP的关键创新在于其两阶段剪枝策略和稀疏率平衡机制。通过结合宽度和深度剪枝，该框架能够更有效地压缩LLM，同时保持较高的性能。稀疏率平衡机制允许用户根据实际需求调整两个阶段的剪枝比例，从而实现更灵活的剪枝策略。

关键设计：宽度剪枝阶段，神经元的重要性得分基于其对输出幅度的影响进行计算。深度剪枝阶段，使用困惑度作为评估指标，迭代地移除对困惑度影响最小的注意力模块。稀疏率平衡机制通过调整两个阶段的剪枝比例，以达到预设的全局稀疏度。具体的参数设置和损失函数细节在论文中未明确说明，可能使用了默认或常用的设置。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

实验结果表明，2SSP框架在多个LLM家族（具体名称未知）和数据集上，相比于五个最先进的竞争方法，在语言建模和下游任务上均取得了更好的性能。更重要的是，2SSP在剪枝时间上获得了高达两个数量级的提升，这使得大规模LLM的剪枝成为可能。具体的性能提升数据和对比基线可以在论文中找到。

🎯 应用场景

该研究成果可应用于各种需要部署大型语言模型的场景，例如移动设备、边缘计算设备等资源受限的环境。通过高效的剪枝，可以在不显著降低模型性能的前提下，减少模型的存储空间和计算开销，从而实现LLM在更多场景下的应用。此外，该方法还可以用于加速LLM的训练和推理过程。

📄 摘要（原文）

We propose a novel Two-Stage framework for Structured Pruning (\textsc{2SSP}) for pruning Large Language Models (LLMs), which combines two different strategies of pruning, namely Width and Depth Pruning. The first stage (Width Pruning) removes entire neurons, hence their corresponding rows and columns, aiming to preserve the connectivity among the pruned structures in the intermediate state of the Feed-Forward Networks in each Transformer block. This is done based on an importance score measuring the impact of each neuron on the output magnitude. The second stage (Depth Pruning), instead, removes entire Attention submodules. This is done by applying an iterative process that removes the Attention with the minimum impact on a given metric of interest (in our case, perplexity). We also propose a novel mechanism to balance the sparsity rate of the two stages w.r.t. to the desired global sparsity. We test \textsc{2SSP} on four LLM families and three sparsity rates (25\%, 37.5\%, and 50\%), measuring the resulting perplexity over three language modeling datasets as well as the performance over six downstream tasks. Our method consistently outperforms five state-of-the-art competitors over three language modeling and six downstream tasks, with an up to two-order-of-magnitude gain in terms of pruning time. The code is available at https://github.com/FabrizioSandri/2SSP.