发表：ICLR24
机构：ETH Zurich
连接：https://arxiv.org/pdf/2401.15024
ABSTRACT
大型语言模型（Large Language Models, LLMs）已成为自然语言处理的基石，但其使用伴随着在计算和内存资源方面的高昂代价。稀疏化（sparsification）提供了一种缓解这些资源限制的解决方案，近期的研究表明，训练好的模型可以在训练后进行稀疏化。然而，现有的稀疏化技术面临一些挑战：它们通常需要额外的数据结构，并且在现有硬件上只能带来有限的加速效果。
在本文中，我们提出了一种新的训练后稀疏化方法 SliceGPT，它用一个更小的（稠密）矩阵替换每一个权重矩阵，从而降低了网络的嵌入维度。通过大量实验证明，SliceGPT 可在不显著损害性能的前提下，去除多达 25% 的模型参数（包括嵌入部分），适用于 LLAMA-2 70B、OPT 66B 和 Phi-2 模型，在零样本任务中分别保留了 99%、99% 和 90% 的原始性能。
我们切片后的模型可在更少的 GPU 上运行，且无需任何额外的代码优化即可运行更快：在 24GB 的消费级 GPU 上，我们将 LLAMA-2 70B 模型的推理总计算量降低为稠密模型的 64%；在 40GB 的 A100 GPU 上则降低至 66%。
我们还提出了一个新的洞察：Transformer 网络中的计算不变性（computational invariance）。该性质使 SliceGPT 成为可能，并希望这一理念能激发并推动未来在减少预训练模型内存与计算需求方向的研究。
代码地址：
🔗 https://github.com/microsoft/TransformerCompression
1 INTRODUCTION
大型语言模型（Large Language Models, LLMs）是包含数十亿参数的神经网络，在数万亿 token 上进行训练（Zhao et al., 2023）。LLM 的训练成本极高，因此已出现将同一个预训练模型用于多个任务的趋势，即“基础模型范式（foundation model paradigm）”。然而，LLM 的体积庞大，使得部署一个预训练模型成为一项代价高昂的工作。许多模型需要多个 GPU 才能完成一次预测，而且由于模型是自回归的，生成文本响应时需要多次执行神经网络的前向传播。因此，降低这些模型的计算需求成为广泛关注的课题，通常通过训练后技术来实现，这类技术统称为模型压缩（model compression）。
大多数模型压缩方法可归为以下四类之一：蒸馏（distillation）、张量分解（tensor decomposition，包含低秩分解）、剪枝（pruning）和量化（quantization）（Hoefler et al., 2021；Gholami et al., 2021；Zhu et al., 2023；Gupta & Agrawal, 2021）。在本研究中，我们专注于剪枝，但也希望我们的方法论能对其他方向的研究有所启发。
虽然剪枝方法已提出多年，但许多方法在剪枝之后需要进行“恢复微调（recovery fine-tuning, RFT）”以维持模型性能，这使得整个流程成本高昂、难以扩展。而我们的 SliceGPT 方法能够使用单张 GPU 在短短几个小时内对大型模型进行压缩，且即使不进行 RFT，也能在生成任务和下游任务上保持较强的性能。
剪枝方法的原理是将 LLM 中部分权重矩阵元素设为零，并（可选地）更新矩阵中其余元素以做出补偿。这种操作产生稀疏模式，使得在神经网络前向传播中进行矩阵乘法时可以跳过某些浮点操作。计算加速的程度取决于稀疏程度及其模式：结构化稀疏通常带来更大的加速。
与其他剪枝方法不同，SliceGPT 剪掉的是整个权重矩阵的行或列（即“切片”掉）。在切片之前，我们对网络进行一次变换，这种变换不会改变预测结果，但能使切片的影响最小化。
最终的结果是，权重矩阵变得更小，block 之间传递的信号也变得更小：我们减小了神经网络的嵌入维度。
图 1 展示了我们的方法与现有稀疏化方法的比较。我们的贡献如下：
我们引入了“计算不变性（computational invariance）”的概念： 我们证明了可以对 Transformer 中的每个权重矩阵施加正交矩阵变换，而不会改变模型行为。
我们使用这种不变性来修改 Transformer 架构中的每个 block， 使 block 之间传递的信号矩阵被投影到其主成分上。我们删除变换后权重矩阵的部分列或行以减小模型规模。我们将这种变换和权重删除的过程称为 SliceGPT。
我们在 OPT（Zhang et al., 2022）和 LLAMA-2（Touvron et al., 2023）模型上进行了大量实验证明， SliceGPT 可以将这些模型压缩最多达 30%，在 WikiText-2 上的困惑度（perplexity）优于现有 2:4 稀疏化方案。在下游任务中，我们还在 Phi-2 模型上进行了实验，显示所有模型最多可被切片 30%，同时保留超过 90% 的稠密模型性能。
2 BACKGROUND
在本节中，我们首先描述一些关于 Transformer 架构的必要背景知识，这使我们能够引入将在证明主要结果中使用的符号。然后我们将描述用于压缩此类架构的稀疏化相关工作。
2.1 TRANSFORMER NETWORKS
2.2 RELATED WORK
在最简单的设置中，可以使用基于幅值的稀疏化（magnitude-based sparsification），其方法是将模型中幅值最小的权重置为零（Han et al., 2016；Zhu & Gupta, 2017；Gale et al., 2019）。虽然这种方法具有良好的可扩展性，但应用于 LLM 时，会导致性能严重下降（Frantar & Alistarh, 2023）。
一种更为复杂的方法是 Optimal Brain Surgeon (OBS)（Hassibi et al., 1993；LeCun et al., 1989），它以系统化的方式移除对损失函数影响最小的权重。该方法通过使用 Hessian 矩阵的逆，对未被剪除的权重进行调整，从而补偿权重移除带来的误差。不幸的是，对于拥有几百万参数的模型而言，OBS 并不实用，因为它需要计算并存储 Hessian 矩阵的逆。
为了解决 OBS 所带来的计算瓶颈，近期研究提出了两个方向：
一是对 Hessian 矩阵的逆进行近似，例如 WoodFisher（Singh & Alistarh, 2020）；
二是对每一层分别应用该方法，例如 Optimal Brain Compression (OBC)（Frantar & Alistarh, 2022），也称为逐层剪枝（layer-wise pruning）。
虽然这些技术在中等规模网络上表现良好，但并不适用于 LLM，因为单层的权重矩阵通常包含超过 10810^8108 个参数。
GPTQ（Frantar et al., 2022）通过一种列按列（column-by-column）的策略，对 LLM 的权重矩阵进行量化（即以更低精度表示参数），并在下一列更新尚未量化的权重，从而解决了上述问题。
SparseGPT（Frantar & Alistarh, 2023）将这种思想应用于剪枝，对 LLM 进行非结构化（unstructured）和半结构化（semi-structured）稀疏化。Sun 等人（2023）则简化了该思路，仅使用了 Hessian 的对角线部分。由于通过非结构化剪枝实现端到端速度提升是一项困难任务，他们还尝试使用如 2:4、4:8（Mishra et al., 2021）等半结构化模式来诱导稀疏性。然而，采用这些结构会损害模型的准确率。
另一种压缩策略是低秩近似（low-rank approximation），即将每个权重矩阵替换为两个内部维度更小的矩阵之积，通常还需后续微调（Hu et al., 2021；Mahabadi et al., 2021；Noach & Goldberg, 2020；Tukan et al., 2020）。为了实现压缩，内部维度必须小于原始维度的一半。相比之下，我们的方法将每个权重矩阵直接替换为一个更小的单个矩阵，通过减小嵌入维度达到目的，且无需微调。
我们提出的方法是删除权重矩阵的行和列，这与卷积网络文献中**滤波器和通道的剪枝（filter/channel pruning）**相似。在该方向上，研究通常在 BatchNorm 参数（Liu et al., 2017）或网络结构（Huang & Wang, 2018）上加入诱导稀疏性的正则化，然后训练或微调，从而实现通道或结构的剪除。
与我们方法最相似的是 ThiNet（Luo et al., 2017；He et al., 2017），其通过在层之间引入线性操作（正如我们也将采用的），并结合更多的微调与正则化来实现压缩。在这类研究中，模型通常比 LLM 小几个数量级，例如 VGG16 网络仅有 1.38 亿个参数，差不多相当于我们考虑的最小 OPT 模型。
而 LLM 的巨大参数规模使得那些依赖大量微调的方案变得不现实，尤其是需要外部循环来选择正则化参数的情况。
近期已有一些研究尝试对 LLM 进行结构化剪枝（structured pruning），然后继续训练（或微调）以恢复性能。例如 LLM-pruner（Ma et al., 2023a）在进一步训练之前剪除了 LLM 中的连接结构；与我们工作同时进行的 LLM Surgeon（van der Ouderaa et al., 2023）则将恢复微调与剪枝过程交替进行。
我们在本文中提供了 SliceGPT 作为一种单次剪枝（single-shot）方法，并同时展示了结合**剪枝后微调（post-slicing recovery fine-tuning）**的实验结果。
3 SLICEGPT
我们的方法 SliceGPT 依赖于 transformer 架构中固有的一种计算不变性（computational invariance）。我们指的是：只要在某个组件的输出中应用一个正交变换（orthogonal transformation），并在下一个组件中“撤销”该变换，整个网络的功能保持不变。我们关键的洞察是，在网络的各个 block 之间执行的 RMSNorm 操作并不会影响这种变换：这两个操作是可交换的（commute）。在本节中，我们首先描述这种不变性如何在 RMSNorm 连接的 transformer 网络中成立，然后说明如何将使用 LayerNorm 的网络转换为 RMSNorm。接着，我们介绍如何使用主成分分析（PCA）计算每一层的变换矩阵，使得 block 之间传递的信号被投影到其主成分上。最后，我们描述如何删除较弱的主成分，从而对应于在修改后的网络中切除行或列。
3.1 COMPUTATIONAL INVARIANCE IN TRANSFORMER NETWORKS
3.2 LAYER NORM TRANSFORMERS CAN BE CONVERTED TO RMSNORM
3.3 A TRANSFORMATION PER BLOCK
3.4 SLICING
4 EXPERIMENTAL VALIDATION
Setup 我们使用 Hugging Face Transformers（Wolf et al., 2019）与 PyTorch（Paszke et al., 2019）实现了我们的代码。正交变换矩阵 Q的计算是在一块 80GB 显存的 H100 GPU 上完成的，以 LLaMA-2 70B 模型为例，计算过程大约耗时 3.5 小时。
在进行 PCA 的过程中，我们使用**双精度（double precision）**来计算协方差矩阵的特征向量。我们发现，若在 PyTorch 中使用单精度计算特征向量，最终模型的精度会出现偏差，详细分析见附录 A.2。
我们使用两个不同的校准数据集进行了实验：WikiText-2 训练集（Merity et al., 2016）和 Alpaca 训练集（Taori et al., 2023）。关于校准集大小与序列长度的消融实验，详见附录 A.3。
Models, Tasks, and GPUs 我们在多个模型族上验证了 SliceGPT 的效果，包括 OPT（Zhang et al., 2022）、LLaMA-2（Touvron et al., 2023），并在零样本任务中额外评估了 Phi-2 模型。
我们排除了 OPT 175B 模型，因为它在性能上已经被更小的 LLaMA-2 模型超越。然而我们预计，这个更大的模型在 SliceGPT 上仍然能取得更好的压缩效果，因为更大的模型通常有更多的冗余，压缩潜力更大（详见第 4.1 节）。
我们在以下两类任务中评估我们的方案：
语言生成任务
热门的零样本（zero-shot）任务
为全面展示 SliceGPT 在推理速度上的提升效果，我们在不同类型的 GPU 上进行了实验：
Quadro RTX6000（24GB 显存）：代表消费级 GPU；
A100（40GB）与 H100（80GB）：代表数据中心级别 GPU。
Baseline Setup 我们最初计划将 SliceGPT 与一种基线方案进行对比：该方案基于最小范数剪除列（或行）。然而我们发现，这种基线效果极差——在仅剪除少数列后，模型在 WikiText-2 上的困惑度（perplexity）飙升至上千。
因此，我们改为将 SliceGPT 与 SparseGPT（Frantar & Alistarh, 2023）进行比较，后者采用了 2:4 的稀疏化比率，这是目前唯一已知的能在实际硬件上带来加速的稀疏化方案（Mishra et al., 2021）。
4.1 RESULTS
生成任务（Generation Task）
我们首先使用 WikiText-2 数据集展示我们的发现。在该实验中，我们评估了 OPT 与 LLaMA-2 模型系列在不同模型规模下，在该数据集上进行剪枝的性能。表 1 展示了不同剪枝比例下模型的困惑度（perplexity）。
在应用于 OPT 模型时，SliceGPT 展现出优于 LLaMA-2 模型的性能，这与我们对这些模型谱分析的直觉一致（详见附录 A.4）。SliceGPT 的性能随着模型规模的增大而提升。
将 SliceGPT 与 SparseGPT 相比，我们发现：在所有 LLaMA-2 模型中，SliceGPT 在剪除 25% 参数的情况下，其性能优于 SparseGPT 采用 2:4 稀疏模式的结果。对于 OPT 模型，除 2.7B 外，其余所有模型在 30% 剪枝比例下的表现也优于 SparseGPT。
零样本任务（Zero-shot Tasks）
我们在五个知名的零样本任务上评估 SliceGPT，包括：
PIQA（Bisk et al., 2020）
WinoGrande（Sakaguchi et al., 2021）
HellaSwag（Zellers et al., 2019）
ARC-e 和 ARC-c（Clark et al., 2018）
遵循相关工作（Frantar & Alistarh, 2023；Dettmers et al., 2022；Frantar et al., 2022；Dettmers et al., 2023），我们使用 LM Evaluation Harness（Gao et al., 2021）进行评估，采用其默认参数。
图 5 展示了剪枝模型在这些任务中的平均得分：
上方图为使用 WikiText-2 校准数据的得分；
下方图为使用 Alpaca 校准数据的得分。
我们观察到的模式与生成任务类似：
OPT 模型更容易压缩，且准确率下降更小；
模型越大，压缩后的性能下降越小。
在该实验中我们也纳入了 Phi-2 模型：我们发现，剪枝后的 Phi-2 模型在性能上与剪枝后的 LLaMA-2 7B 模型相当。OPT 与 LLaMA-2 的最大规模模型均可以被有效压缩，即使剪除 30% 参数，也仅带来几个百分点的性能损失。
恢复微调实验（Recovery Fine-Tuning）
我们在此还进行了恢复微调（Recovery Fine-Tuning, RFT）实验：
我们使用 LoRA（Hu et al., 2021） 对剪枝后的 LLaMA-2 与 Phi-2 模型进行少量微调，方法参照 Ma et al. (2023a)。
对于使用 WikiText-2 剪枝的模型，我们使用约 1000 条序列进行微调；
对于使用 Alpaca 数据剪枝的模型，我们使用 5000 条序列进行微调。
所有 RFT 使用如下参数：
lora_r = 32，lora_alpha = 10，序列长度 = 1024，
其余超参数采用 Hugging Face PEFT 包（Mangrulkar et al., 2022）默认设置。
图 6 展示了微调结果。我们观察到，使用 Alpaca 数据集进行 RFT 的效果显著优于使用 WikiText-2。这一差异可归因于 Alpaca 与评估任务的相似性。
对于被剪除 30% 的最大 LLaMA-2 70B 模型，使用 Alpaca 做 RFT 后，其平均准确率为 74.3%，而原始稠密模型为 76.6%；
该剪枝模型的参数量约为 51.6B，推理吞吐量也有显著提升（详见后文）；
Phi-2 模型若仅用 WikiText-2 做 RFT 无法恢复准确率，但使用 Alpaca 可恢复若干百分点；
剪除 25% 后的 Phi-2 平均准确率为 65.2%，而稠密模型为 72.2%；
剪枝模型约含 2.2B 参数，保留了约 90.3% 的原始模型性能；
这表明，即使是小型语言模型，也能从训练后剪枝中获益。各任务的准确率表格见附录 A.5。
吞吐量基准测试（Benchmarking Throughput）
与传统稀疏方法只对 Win 与 Wout引入稀疏性不同，SliceGPT 还对中间特征 X引入了结构化稀疏性：直接剪除 X的整列，减小了嵌入维度，从而提高了压缩模型的 FLOPs 效率和内存移动效率。
我们将剪枝比例为 25% 和 50% 的模型在 80GB H100 GPU 上与稠密模型进行了吞吐量对比。设定序列长度为 128，持续加倍 batch size，直到 GPU 显存耗尽或吞吐量下降为止。
结果表明：
25% 剪枝模型的 token 吞吐量提升最多可达 1.55 倍；
50% 剪枝后，最大模型只需 1 张 GPU 就可运行（原需 2 张），吞吐量提升巨大：
OPT-66B 提升为 3.13×
LLaMA2-70B 提升为 1.87×
固定 GPU 数量下，相当于达到了：
6.26×（3.13× × 2）
3.75×（1.87× × 2） 的吞吐率增益。
虽然 SliceGPT 在 50% 剪枝下的 WikiText-2 perplexity 差于 SparseGPT（2:4），但在吞吐量上远高于任何不剪中间特征 XXX 的稀疏方法。
对于 13B 模型，batch size 带来的吞吐提升不明显，因为模型本身占用显存不多。但对于消费级 GPU（显存小），小模型的吞吐量也会因剪枝而受益。详见附录 A.6。
推理时间（Inference Time）
我们接着评估了 SliceGPT 压缩模型的端到端运行时间。
表 2 比较了 OPT-66B 与 LLaMA2-70B 模型在 Quadro RTX6000 与 A100 GPU 上生成单个 token 所需时间：
在 RTX6000 上，剪除 25% 后的模型推理时间缩短了 16~17%；
在 A100 上，推理时间减少了 11~13%；
同时也减少了所需 GPU 数量，降低了能耗与部署成本。
例如，LLaMA2-70B 在 RTX6000 上，推理一个 token 所需计算由 1764 GPU-ms 降至 1075 GPU-ms（64%）。
我们将这种提升归因于：
替换为更小权重矩阵；
全程使用稠密计算内核（dense kernel），这是其他剪枝方法做不到的。
目前，SparseGPT（2:4）还无法实现端到端推理加速。因此，我们改为比较 transformer layer 内各操作的相对执行时间。我们发现：对于大模型，SliceGPT（25%）在推理速度与困惑度方面可与 SparseGPT（2:4）竞争甚至超越。详见附录 A.7。
计算成本（Compute cost）
所有 LLaMA-2、OPT 和 Phi-2 模型都可以在 单张 GPU 上在 1 到 3 小时内完成剪枝。
如果加上恢复微调，总时间也仅需 1 到 5 小时。详见表 3。
5 CONCLUSION AND FUTURE WORK
我们提出了 SliceGPT，一种支持大型语言模型结构化剪枝的新方法。
该方法使我们无需任何额外代码优化，即可将 LLaMA-2 70B 模型在 40GB A100 GPU 上的推理成本降低到稠密模型的 66%，所需 GPU 数量也从 4 块减少到 3 块，同时其保留集（held-out set）困惑度优于 SparseGPT 的 2:4 稀疏方案。
在 24GB RTX6000 GPU 上，推理成本被降至稠密模型的 64%，所需 GPU 数量从 7 块减少到 5 块。
在零样本下游任务中，分别对 OPT-66B、LLaMA-2-70B 和 Phi-2 进行 25% 剪枝，其性能可分别保留稠密模型的 99%、96% 和 87%；
在进行恢复微调后（recovery fine-tuning），这两个模型的表现进一步提高，LLaMA-2-70B 和 Phi-2 可分别恢复至 99% 和 90% 的稠密性能。
我们的方法仍有进一步发展的空间：
当前，在同等参数规模下，小型但稠密的语言模型往往优于被剪到同样大小的大模型，但我们认为这种情况不会持续太久；
虽然 SliceGPT 剪枝后的模型参数量通常多于 SparseGPT，但我们的方法能使更大的 batch size 载入 GPU 显存，且无需稀疏结构支持的额外开销；
因此，一种可能的改进方向是将两种方法结合，发挥各自优势；
此外，还可以尝试其他方式计算变换矩阵 QQQ，进一步提升剪枝效果；
若想继续减少推理时间和 GPU 数量，还可以结合量化（quantization）（Xiao et al., 2023；Dettmers et al., 2022；Ashkboos et al., 2023；Frantar et al., 2022）以及其他结构化剪枝方法（如 Ma et al., 2023b）。