NeurIPS 2024 ｜ SparseLLM：突破性全局剪枝技术，大语言模型稀疏化革命

1.摘要

大型语言模型（LLMs）如 GPT 和 LLaMA 已经彻底改变了自然语言处理领域，但其计算成本极高。我们在论文中提出了 SparseLLM ，一种创新的全局剪枝框架，使得大规模模型压缩更加高效。 SparseLLM 通过将全局剪枝问题分解为更小的子问题，在保持卓越性能的同时实现了内存和计算效率，尤其在高稀疏率情况下表现尤为突出。 我们的方法在困惑度（perplexity）降低和灵活性上超越了当前的最先进方法 。

2.作者

Guangji Bai¹，Yijiang Li²，Chen Ling¹，Kibaek Kim²，Liang Zhao¹*

• ¹Emory University
• ²Argonne National Lab
• *Corresponding Author

论文链接：https://arxiv.org/abs/2402.17946

代码： https://github.com/BaiTheBest/SparseLLM

3.情景导入

随着大型语言模型（LLMs）如 GPT 和 LLaMA 在自然语言处理领域的突破，现如今的模型能够在各种复杂的语言任务中表现优异。然而，这些模型往往包含数十亿参数，导致计算资源的需求极为庞大。为了让LLMs在更多的实际应用中变得可行，研究人员进行了大量的模型压缩工作，其中包括剪枝、量化、知识蒸馏和低秩分解等方法。

剪枝作为一种重要的压缩策略，通过引入稀疏性提升了内存和计算效率。尽管剪枝的有效性已在视觉任务和较小规模的语言模型中得到验证， 但全局剪枝由于需要将整个模型加载到同一个GPU中 ，对于如今的大规模LLMs而言是不现实的。因此，近年来出现了局部剪枝方法比如SparseGPT和Wanda，尽管这些方法简单粗暴地局部化每个layer的剪枝，从而在效率上有了提升，但局部剪枝带来的效果往往是次优 (suboptimal) 的。

我们提出的 SparseLLM 框架，通过将全局剪枝问题分解为更易管理的子问题，从而在高稀疏度下也能实现高效的优化和优异的性能。SparseLLM 在实现了内存和计算效率的同时，超越了当前最先进的剪枝方法，为LLMs的压缩和应用树立了新的标杆。

4.现有方法及其缺陷

在模型剪枝中，剪枝方法主要分为两大类：全局剪枝（Global Pruning）和局部剪枝（Local Pruning）。

全局剪枝

全局剪枝旨在对整个模型应用统一的稀疏化掩码（sparsity mask），以最小化未压缩模型与压缩模型之间的全局损失。虽然理论上全局剪枝可以为模型提供最优的性能，尤其是在高稀疏度情况下，但它的主要缺点在于其内存开销巨大。为了实现全局剪枝，整个模型必须被加载到同一个GPU中，这对于如今规模巨大的LLMs（如GPT和LLaMA）来说是不现实的 。

局部剪枝

为了规避全局剪枝的内存瓶颈，局部剪枝通过将模型压缩分解为每一层的子问题来减少内存消耗。局部剪枝的方法通常会对每一层的输出进行独立的稀疏化，并构造局部损失来衡量未压缩和压缩层之间的差异。虽然局部剪枝方法在资源使用方面更加高效，但由于它只关注每一层的局部最优解，导致了全局上次优的模型性能，尤其是在高稀疏度下表现尤为显著 。

方法缺陷

局部剪枝的局限性在于它无法很好地捕捉模型层之间的相互依赖关系，只针对各层的稀疏化进行优化。这种过度约束每层激活值的行为，容易导致全局性能的下降。因此，尽管局部剪枝在某些情况下能带来效率提升，但它往往无法达到全局剪枝所能带来的最优性能。

5.技术方法

在SparseLLM框架中，我们的目标是实现全局剪枝。S parseLLM通过将全局剪枝目标分解为多个子问题，每个子问题可以使用较少的资源解决，并且可以协同实现全局剪枝目标 。SparseLLM的优势在于它能够在内存消耗较低的情况下实现全局剪枝。

动机

SparseLLM的开发基于以下观察：LLMs可以被表述为一个复合函数，后一个模块的输出是下一个模块的输入。这使得可以通过辅助变量将全局剪枝目标重新表述为等价形式，从而将其分解为多个子问题。然后，开发了一种高效的算法，通过交替优化每个子问题来实现全局最优解。

剪枝问题的重新表述

我们通过对模型的稠密线性部分和非线性部分进行解耦，重新表述了剪枝问题。将每一层的输出存储为一个新变量 ，非线性层的输出则表示为激活值 。接着，我们优化以下目标函数：

其中约束条件为：

这里 表示模型的总层数， 表示需要剪枝的层集合， 和 分别代表原始预训练模型的中间变量值。 是模型的最终输出结果 。

Remark

这个公式具有高度的通用性和灵活性。当 时，解决该问题等同于全局剪枝。而当 时，问题简化为局部剪枝，独立地考虑每一层。通过调节 的大小，我们可以在全局和局部剪枝之间无缝切换。

SparseLLM优化过程

原始的全局剪枝问题包含全局约束条件，为了避免直接求解带约束的目标函数，我们引入了辅助变量和惩罚函数。通过引入 惩罚项，原始的全局目标函数被转化为无约束问题，具体表达如下：

其中， 和 是超参数。这样，我们将每一层的剪枝问题解耦为局部子问题，同时保留了层之间的依赖性。

SparseLLM算法的关键在于上面公式中的 的灵活性 ，目的是在完全全局剪枝（会导致内存瓶颈）和完全局部剪枝（会导致次优性能）之间找到更好的折中。粗暴地对所有层进行全局剪枝是不现实的。最近的研究表明，前馈网络 (FFN) 模块在每个解码器层中的参数占LLM总参数的三分之二以上。因此，SparseLLM优先对FFN模块进行全局剪枝，同时仍然遵循对多头注意力模块 (MHA) 的局部剪枝策略（见Figure 2）。这种策略在剪枝大规模模型的计算可行性与剪枝过程的有效性之间取得了平衡，并遵循了现有LLM剪枝框架的限制和实践。

具体而言，SparseLLM将 设置为 ，其中 表示每个解码器层FFN模块中线性层的索引集（具体数学细节详见论文）。也就是说， 同一个FFN模块中的线性层被全局剪枝 ，而剩下的线性层仍按照局部剪枝。对于每一层的FFN模块，我们的目标是最小化以下无约束的目标函数：

此时，每一层的剪枝问题被表述为线性投影模块（上投影和下投影）的优化问题，并且可以针对每一个子问题进行交替优化。这一过程通过引入辅助变量和惩罚函数，使得我们能够在不消耗过多内存的前提下，求解全局最优解。

Figure 2: 该图展示了 SparseLLM 框架在不同LLM架构中的工作原理。左侧描绘了在 OPT模型中的剪枝流程。SparseLLM通过引入辅助变量，将全局剪枝问题分解为可管理的子问题，利用线性层的上下投影模块（Up Proj 和 Down Proj）以及ReLU激活函数来实现模型压缩。右侧描绘了在 LLaMA模型 中的剪枝流程。此处，SparseLLM除了使用线性层的上下投影外，还利用了SiLU激活函数与Gate Proj模块来实现更复杂的剪枝操作。图中标注了剪枝层（蓝色）、辅助变量（绿色）和预训练模型的固定输入/输出（粉色），以帮助理解各模块在剪枝过程中的相互关系与作用。

OPT 模型中的子问题求解

在 OPT 模型中，我们的目标是对每一层的FFN模块进行全局剪枝。根据公式5，我们首先考虑交替优化子问题，对于每一个子问题的闭式解，我们分以下步骤进行：

1. 剪枝权重的优化 ：首先，我们固定其他变量，优化权重 和 。具体来说，优化目标是最小化以下表达式：

   为了求解该问题，我们首先对 进行分解： ，其中 （ 表示伪逆）。通过将分解后的 带入原始损失函数中，可以得到权重剪枝问题的闭式解，类似于 SparseGPT 等剪枝方法。

2. 激活值的更新 ：对于激活值 的优化，这是一个类似于权重剪枝的简单最小二乘问题。激活值 的更新可以通过如下公式求解：

3. 输出 的更新 ：对于输出 ，我们最小化以下非凸的损失函数：

   由于ReLU函数的逐元素操作，问题可以通过简单的if-then逻辑分为两种情况： ，以及 。

LLaMA 模型中的子问题求解

在LLaMA模型中，我们的剪枝过程与OPT模型类似，唯一不同的是LLaMA模型中包含额外的门控投影层，并使用了SiLU激活函数。我们针对LLaMA的各层子问题进行如下求解：

1. 剪枝权重的优化 ：对于LLaMA模型中的FFN模块，权重剪枝的优化过程与OPT模型类似。我们通过最小化以下表达式来优化权重：