LLM-Pruner: 剪枝+少量数据+少量训练 = 高效的LLM压缩

概要

大语言模型（LLMs, Large Language Models）在各种任务上展现出了惊人的能力，这些能力很大程度上来自于模型庞大的模型规模以及海量的训练语料。为了应对这些模型部署上存在的挑战，许多研究者开始关注大语言模型的轻量化问题。由于 LLM 模型庞大的参数量，我们希望能以最小的成本完成对模型的压缩，尽可能的减少压缩后训练的开销，实现高效的模型压缩。

因为，我们提出了一种基于自动化的结构化剪枝的方案，采用结构化剪枝可以尽最大可能保留模型已习得的知识，实现高效的大语言模型压缩。在实验中，我们仅使用了 5 万条训练语料以及单张 4090 24GB 显卡，就可以在 3 小时内完成 LLaMA，Vicuna 和 ChatGLM 等大语言模型的剪枝和训练。

论文标题：

LLM-Pruner: On the Structural Pruning of Large Language Models

论文作者：

Xinyin Ma, Gongfan Fang, Xinchao Wang

单位：

新加坡国立大学

论文地址：

https://arxiv.org/abs/2305.11627

Github地址：

https://github.com/horseee/LLM-Pruner

大语言模型压缩面临的问题

首先，大语言模型的压缩与之前的语言模型（例如 BERT，RoBERTa 等）的压缩等有什么差异呢。这需要从模型/数据/任务三个角度来分析，

模型规模：第一个主要差异来自 LLM 参数量规模远超之前的语言模型，这导致许多依赖重新训练的压缩方案，例如知识蒸馏，重新训练的开销较大。
海量训练语料：许多 LLMs 经历了 1 万亿甚至更大规模的 tokens 上的训练，这导致许多依赖于原始数据或收集替代数据的方案变得尤其昂贵。
任务无关的模型压缩：现有的压缩算法通常针对单一、特定的任务进行压缩，而 LLMs 是很优秀的多任务处理器，在压缩过程中我们不希望折损 LLM 的通用性和多功能性。

上述三种问题实际上对应了三个基本要求：

降低训练规模：压缩算法需要尽可能少的依赖大规模的重新训练。
减少数据依赖：压缩训练所需要的数据需要尽可能少
保留模型原始能力：压缩模型过程中需要保留一定的原始模型能力。

由此，我们需要一种能够避免大规模重新训练、且能保持模型原有能力的压缩方法。现有的较为可行的两种方案是模型量化和结构化剪枝，他们的最主要特点都在于不需要完全从零开始训练。其中模型量化侧重于降低推理阶段的存储开销以及提升计算速度，而结构化剪枝则直接移除部分参数实现压缩，两种方案可以相互结合达到最优性能。本文主要介绍基于结构化剪枝的 LLM 压缩方法，通过保留 LLM 中的重要参数，降低压缩训练成本。

方法

本文提出一种简洁高效的结构化剪枝方法，遵循了经典的“重要性估计-剪枝-微调”的策略，能够在有限资源下完成大语言模型的压缩：

重要性估计：在 LLM 中存在着复杂的子结构，剪枝算法需要对网络的参数进行分组，找到可移除的最小单元（通常称为 Group），并对各个组进行重要性评估。不同于之前的研究，我们解决了依赖图自动的挖掘神经网络中存在依存关系的子结构，也因此我们的方法更易拓展到更多的大规模语言模型。
剪枝：在得到各个分组的重要性后，我们将冗余的组整个移除，从而降低模型的参数量。
微调：在剪枝后的模型上应用 LoRA 等高效微调策略，恢复模型性能。

3.1 参数组重要性分析

首先我们分析了 LLM 模型中参数之间的依赖关系，从而找到最小可移除的分组。这些依赖可以归纳为三类，下图展示了 LLaMA 模型中存在的层耦合关系，主要由 MLP 内部的耦合、MHA 内部的耦合（主要为 QKVO 四个映射层）与整个网络中的通道耦合。这些耦合导致虚线显示的参数、神经元需要被同时移除，从而确保剪枝后结构的正确性。

为了自动化上述依赖分析，我们需要对神经元的依赖进行形式化建模。假设和是模型中的两个神经元，和表示被指向和指向的所有神经元。结构之间的依赖性可以定义为：

其中表示神经元的入度。注意这里的依赖关系并非是前向计算流向上的依赖关系，而是神经元结构上的依赖关系。这种依赖性是方向性的，我们因此可以得到另一种依赖性：

其中表示神经元的出度。依赖性在网络剪枝中的重要性在于，如果当前神经元（例如）仅依赖于另一个神经元（例如），且神经元被剪枝，那么神经元也必须被剪枝。此时如果不裁剪，那么网络结构中的维度就会不匹配。基于上述规则，我们可以实现一个自动识别参数依赖的程序，帮助我们剪枝各种 LLM 模型。同时，我们注意到 LLM 的最小可移除子结构通常包含了多个层，这就要求我们设计一种针对整个组的重要性估计策略。

3.2 组重要性估计

目前，我们已经对模型内部的耦合结构进行了建模。为了找到冗余分组，我们需要对将整个组作为一个最小单元进行重要性估计。鉴于对训练数据集的访问权限有限，我们探索使用可用的公共数据集或手动创建的样本作为替代资源。尽管这些数据集的领域可能与训练集不完全一致，但它们仍提供了评估结构组重要性的宝贵信息。

a）权重向量的一阶重要性估计：给定一个数据集，其中 N 是样本数量。在我们的实验中，我们设置 N 等于 10，也就是仅需 10 个样本。一个组（如前所述，被定义为一组耦合结构）包含一组互相耦合的参数，其中 M 是一个组中耦合结构的数量，是组内第i层的权重。

在修剪时，我们的目标是移除对模型预测影响最小的结构，因此，参数的重要性可以通过损失函数的剪枝后扰动来判断。我们对损失函数进行泰勒展开并进行移项，可以得到一种简单的重要性估计指标，这类重要性评估策略需要我们尽可能精准的对损失函数进行建模，因此我们采用了二阶的泰勒展开：

其中，表示 Hessian 矩阵，表示 Next-token Prediction Loss。通常来说，重要性的一阶项由于模型在训练数据集已完全收敛，此项通常为 0，即。然而，由于我们选取的数据并非来自原始数据集，重要性评估中我们通常可以得到非 0 的一阶项，即（实验中发现我们的数据上的 Loss 高于充分收敛样本上的 Loss）。

重要性估计的第二项要求我们对 Hessian 矩阵进行计算。然而，由于 Hessian 矩阵的计算复杂度过高，这对参数了巨大的 LLM 而言是不现实的。因此，我们不直接对参数计算 Hessian 矩阵，而是仅考虑估计 hessian 矩阵的对角线元素，这需要引入单个参数标量的重要性评估。

b）单个参数的二阶重要性估计：上述过程对整个权重向量的进行了估计。实际上，我们可以在更细粒度上得到另一种重要性度量，其中参数矩阵内的每个数值都被独立地执行重要性估计。

这里，用于索引中的第 k 个参数。我们使用 Fisher Information Matrix 来近似 Hessian 矩阵的对角线，重要性可以被定义为：

3.3 组重要性聚合

在获得权重向量重要性以及单个参数的重要性后，我们还需要将其进行聚合，得到整个分组的重要性。本文提出四种重要性聚合的策略：

求和 Sum: 或者
乘积 Prod: 或者
极值 Max: 或者 ;
门控 Last-Only: 或者

上述聚合方式实际上包含了不同的偏置，例如求和策略认为不同参数的贡献是独立且可叠加的，乘积策略则假设不同层的重要性会相互影响，最大值策略的偏置在于层的重要性由某一层主导。最后，门控策略则认为组内的最后一层主导了整个组的重要性，因为通过将该层设为 0 我们可以使得整个组不再参与网络预测。在评估每个组的重要性之后，我们根据预定义的修剪比例对每个组的重要性进行排序，并修剪重要性较低的组。

3.4 剪枝模型的低秩近似快速恢复

为了加速模型的压缩后训练，同时提高有限数据下的训练效率，我们需要降低优化的参数量。为了实现这一目标，我们将 LoRA 与剪枝模型相结合。模型中的每个可学习的权重矩阵，表示为，包含 LLM 中所有剪枝和未剪枝的线性投影。的更新值可以被分解为，其中 and 。前向计算现在表达为：

其中 b 是稠密层中的偏置。仅训练 P 和 Q 减少了整体训练复杂性，从而减少了对大规模训练数据的需求。此外，额外的参数 P 和 Q 可以被重新参数化为，这不会在最终压缩模型中造成额外的参数。

3.5 算法总结

上述过程给出了 LLM 模型的依赖分析、重要性评估以及后训练的完整方案，通过聚合一阶和二阶泰勒展开，我们可以得到更加鲁棒的重要性评估策略。对于大模型而言，重要性评估指标是尤其重要的，因为剪枝造成的性能损失越大，后训练恢复所需要的数据量、训练时间也就越多

实验

本文对三种开源的 LLM 进行剪枝实验，包括 LLaMA-7B，Vicuna-7B 和 ChatGLM-6B，剪枝前后的模型参数量、MACs 和内存占用如表所示。

4.1 LLaMA-7B剪枝后模型的Zero-shot能力验证

对于大模型而言，保留原始的多任务处理以及零样本能力是尤其重要的，我们完整评测了不同多种方案，包括权重剪枝、随机剪枝以及本文方法的剪枝效果。我们发现在剪枝 20% 参数的情况下，大模型依旧能保持一定的 zero-shot 能力，同时经过少量微调，zero-shot 性能可以快速提升，约达到基座模型的 94%。

4.2 Vicuna-7B剪枝后模型的Zero-shot能力验证

4.3 ChatGLM-6B剪枝后模型的Zero-shot能力验证

4.4 剪枝后模型的生成结果

更多分析结果请参考论文。

总结

在本文中，我们提出了 LLM-Pruner，一种用于大型语言模型的结构化剪枝方法。LLM-Pruner 旨在以任务无关的方式压缩庞大的语言模型，同时尽量减少对原始训练语料库的依赖，并保留 LLM 的语言能力。LLM-Pruner 通过迭代地检查模型中的每个神经元作为识别依赖组的触发器，从而构建 LLM 的依赖图。随后，LLM-Pruner 使用参数级和权重级估计来评估这些组的重要性。

最后，我们利用 LoRA 对被剪枝模型进行快速恢复和调整。我们使用多个 zero-shot 数据集评估了 LLM-Pruner 在三个不同模型（LLaMA，Vicuna 和 ChatGLM）上的有效性。我们的实验结果表明，LLM-Pruner 成功地剪枝了模型，在保留 zero-shot 能力的同时减轻了计算负担。

欢迎试用我们的项目：

https://github.com/horseee/LLM-Pruner

更多阅读