AIxiv专栏是机器之心发布学术、技术内容的栏目。过去数年，机器之心AIxiv专栏接收报道了2000多篇内容，覆盖全球各大高校与企业的顶级实验室，有效促进了学术交流与传播。如果您有优秀的工作想要分享，欢迎投稿或者联系报道。投稿邮箱：liyazhou@jiqizhixin.com；zhaoyunfeng@jiqizhixin.com
- 论文标题：SparseLLM: Towards Global Pruning of Pre-trained Language Models
- 论文链接：https://arxiv.org/abs/2402.17946
- 代码：https://github.com/BaiTheBest/SparseLLM
情景导入
随着大型语言模型（LLMs）如 GPT 和 LLaMA 在自然语言处理领域的突破，现如今的模型能够在各种复杂的语言任务中表现优异。然而，这些模型往往包含数十亿参数，导致计算资源的需求极为庞大。为了让LLMs在更多的实际应用中变得可行，研究人员进行了大量的模型压缩工作，其中包括剪枝、量化、知识蒸馏和低秩分解等方法。
现有方法及其缺陷
在模型剪枝中，剪枝方法主要分为两大类：全局剪枝（Global Pruning）和局部剪枝（Local Pruning）。
全局剪枝
全局剪枝旨在对整个模型应用统一的稀疏化掩码（sparsity mask），以最小化未压缩模型与压缩模型之间的全局损失。虽然理论上全局剪枝可以为模型提供最优的性能，尤其是在高稀疏度情况下，但它的主要缺点在于其内存开销巨大。为了实现全局剪枝，整个模型必须被加载到同一个GPU中，这对于如今规模巨大的LLMs（如GPT和LLaMA）来说是不现实的。
局部剪枝
为了规避全局剪枝的内存瓶颈，局部剪枝通过将模型压缩分解为每一层的子问题来减少内存消耗。局部剪枝的方法通常会对每一层的输出进行独立的稀疏化，并构造局部损失来衡量未压缩和压缩层之间的差异。虽然局部剪枝方法在资源使用方面更加高效，但由于它只关注每一层的局部最优解，导致了全局上次优的模型性能，尤其是在高稀疏度下表现尤为显著。
方法缺陷
局部剪枝的局限性在于它无法很好地捕捉模型层之间的相互依赖关系，只针对各层的稀疏化进行优化。这种过度约束每层激活值的行为，容易导致全局性能的下降。因此，尽管局部剪枝在某些情况下能带来效率提升，但它往往无法达到全局剪枝所能带来的最优性能。
技术方法
在SparseLLM框架中，我们的目标是实现全局剪枝。SparseLLM通过将全局剪枝目标分解为多个子问题，每个子问题可以使用较少的资源解决，并且可以协同实现全局剪枝目标。SparseLLM的优势在于它能够在内存消耗较低的情况下实现全局剪枝。
动机
SparseLLM的开发基于以下观察：LLMs可以被表述为一个复合函数，后一个模块的输出是下一个模块的输入。这使得可以通过辅助变量将全局剪枝目标重新表述为等价形式，从而将其分解为多个子问题。然后，开发了一种高效的算法，通过交替优化每个子问题来实现全局最优解。
剪枝问题的重新表述
表示模型的总层数，
表示需要剪枝的层集合，
和
分别代表原始预训练模型的中间变量值。 是模型的最终输出结果 。Remark：
时，解决该问题等同于全局剪枝。而当
时，问题简化为局部剪枝，独立地考虑每一层。通过调节
的大小，我们可以在全局和局部剪枝之间无缝切换。SparseLLM优化过程
惩罚项，原始的全局目标函数被转化为无约束问题，具体表达如下：
和
是超参数。这样，我们将每一层的剪枝问题解耦为局部子问题，同时保留了层之间的依赖性。
设置为
，其中
表示每个解码器层FFN模块中线性层的索引集（具体数学细节详见论文）。也就是说，同一个FFN模块中的线性层被全局剪枝，而剩下的线性层仍按照局部剪枝。对于每一层的FFN模块，我们的目标是最小化以下无约束的目标函数：
此时，每一层的剪枝问题被表述为线性投影模块（上投影和下投影）的优化问题，并且可以针对每一个子问题进行交替优化。这一过程通过引入辅助变量和惩罚函数，使得我们能够在不消耗过多内存的前提下，求解全局最优解。
OPT 模型中的子问题求解
剪枝权重的优化：首先，我们固定其他变量，优化权重 
和 
。具体来说，优化目标是最小化以下表达式：
为了求解该问题，我们首先对 
进行分解：
，其中 
（ 表示伪逆）。通过将分解后的 带入原始损失函数中，可以得到权重剪枝问题的闭式解，类似于 SparseGPT 等剪枝方法。激活值的更新：对于激活值 
的优化，这是一个类似于权重剪枝的简单最小二乘问题。激活值 
的更新可以通过如下公式求解：
输出 的更新：对于输出 
，我们最小化以下非凸的损失函数：
由于ReLU函数的逐元素操作，问题可以通过简单的if-then逻辑分为两种情况： 
LLaMA 模型中的子问题求解
剪枝权重的优化：对于LLaMA模型中的FFN模块，权重剪枝的优化过程与OPT模型类似。我们通过最小化以下表达式来优化权重： 
并使用类似于OPT模型的分解方法求解。 激活值 的更新：在LLaMA模型中，我们使用SiLU激活函数，因此激活值的更新公式如下： 
输出 的更新：对于LLaMA模型中的输出 
，我们最小化以下损失函数：
通过类似的数学操作，我们可以求得 
的解析解为：
门控投影输出 
的更新：对于门控投影层的 ，我们通过最小化以下表达式来更新：
由于SiLU函数的非凸特性，问题可以通过查找预先计算好的解决方案的查找表来快速求解。
实验部分
实验设置
模型与剪枝策略
OPT 模型：我们针对不同规模的 OPT 模型（如 OPT-1.3b、OPT-2.7b 等）进行全局剪枝，评估 SparseLLM 在不同大小的模型上的效果。 LLaMA 模型：我们也对 LLaMA 系列模型进行了实验，特别是 LLaMA-2 13b 模型，考察其在使用 SparseLLM 框架下的剪枝效果。
比较基线方法
局部幅度剪枝：一种传统的剪枝方法，逐层根据权重幅度进行独立的局部剪枝，不考虑层间的依赖性。 SparseGPT：一种最新的局部剪枝方法，专门针对预训练语言模型进行高效的权重剪枝，能够在高稀疏度条件下维持较好的性能。 Wanda：另一种基于局部剪枝的方法，通过最小化剪枝后的模型输出误差，达到较好的压缩效果。
结果分析
剪枝效果分析
模型大小与性能权衡
收敛速度分析
不同任务的通用性
为了验证 SparseLLM 在零样本设置下的通用性，我们在多个零样本任务上测试了剪枝后的模型性能。实验结果表明，SparseLLM 在大多数任务上保持了优异的性能，特别是在涉及推理和生成任务时，SparseLLM 剪枝后的模型仍然能够生成高质量的文本输出，而其他基线方法的性能在高稀疏度下显著下降。
实验小结
实验结果表明，SparseLLM 框架能够在不同规模的预训练语言模型上实现高效的全局剪枝，同时保持良好的模型性能。无论是在较小的 OPT 模型上，还是在更大规模的 LLaMA 模型上，SparseLLM 均表现出色，特别是在高稀疏度的条件下表现尤为突出。此外，SparseLLM 的收敛速度和剪枝后的通用性也为其在实际应用中的高效性和适用性提供了强有力的支持。
结论
未来研究方向
动态剪枝策略的研究：当前的 SparseLLM 是基于静态的全局剪枝策略，即在训练前决定剪枝的层和稀疏度。未来可以探索动态剪枝策略，结合任务需求和模型训练过程，实时调整剪枝比例和剪枝区域，以进一步提升模型压缩效率。 稀疏性与硬件加速的结合：尽管 SparseLLM 能够大幅减少模型的计算和存储开销，但在实际部署中如何高效利用硬件加速器（如GPU、TPU）的稀疏性支持仍是一个开放问题。未来的工作可以研究如何更好地结合稀疏矩阵乘法和硬件加速技术，以实现更大的性能提升。 剪枝后的模型微调：尽管我们的实验已经展示了 SparseLLM 在高稀疏度下的性能保持能力，但剪枝后的模型是否可以通过微调进一步提升性能仍然值得探索。未来的研究可以结合剪枝和微调技术，探讨如何在稀疏化后进一步优化模型性能。
