大规模语言模型剪枝又一力作，比SparseGPT还快300倍！

news2024/9/21 0:33:57

©Paperweekly 原创 · 作者 | An.

单位 | 中科院自动化所

研究方向 | 计算机视觉、模型压缩

DenseNet、ConvNeXt、Network Slimming 一作刘壮研究员的剪枝新工作，针对 LLMs 特点设计的极低成本无需微调的剪枝算法，耗时接近幅值剪枝，性能表现与 SparseGPT [1] 相当。

论文标题：

A Simple and Effective Pruning Approach for Large Language Models

论文链接：

https://arxiv.org/pdf/2306.11695.pdf

代码链接：

https://github.com/locuslab/wanda

动机&背景

大规模语言模型（LLMs）的实用化部署需要解决高计算成本的问题，模型剪枝是主流且很有潜力的解决方案。现有的剪枝方法大都需要重新训练或漫长的迭代修剪，LLMs 巨大的微调成本阻碍了这些方法的应用。今年年初出现的 SparseGPT 不再需要传统的重新训练，但仍需要一个计算密集型的权重补偿过程。

SparseGPT 工作中的实验表明——幅值剪枝在较高剪枝率下在 LLMs 的性能退化严重，这与过往幅值剪枝在小模型的趋势不符。作者敏锐地观察到了这一差异，提出了一种简单有效的方法来解决幅值剪枝在 LLMs 下的性能退化，称为 Wanda（Pruning by Weights and activations）。

这一方法的提出受到了最近研究 [2] 的启发——在 LLMs 达到一定规模时（大约为 6B 参数），一小组隐藏特征的幅值将显著大于其他特征，这些异常特征对于 LLMs 的预测能力至关重要，当这些特征被置零时，语言模型的性能会出现严重退化。具体来说，本文引入了一种新的剪枝度量，每个权重的重要性由其幅值大小和相应输入激活的范数的乘积来决定。该方法利用这个新的度量，在线性层的每个输出内局部比较权重并删除较低的优先级权重来获得稀疏的 LLMs。

方法：Wanda

该方法由两个简单但必不可少的组件构成——剪枝度量和剪枝粒度。首先，作者提出了一种新的剪枝度量，将权重和输入激活结合到权重重要性的计算中。其次，作者在每个输出的基础上比较权重，而不是整个层，并发现这一因素对于有效地修剪 LLM 至关重要。图 1 是 Wanda 的概述。

▲ 图1. Wanda和幅值剪枝的方法对比

2.1 关键组件

剪枝度量：权重形状为，输入激活形状为。对于每个权重，作者通过其大小和相应的输入特征范数的乘积来评估其重要性。具体地，当前权重的得分可定义为：

剪枝粒度：作者认为，除了剪枝度量之外，剪枝粒度，即用于比较重要性的权重集，在修剪 LLM 中发挥着重要作用。现有的剪枝方法通常采用分层或全局修剪策略，这意味着在网络中的每个层内或所有层之间比较权重。本文采用了一种全新的剪枝粒度——在每个输出的基础上（图 1 中的每行）比较和删除权重，其中在每个输出神经元内局部比较权重重要性得分。权重的比较组被定义为连接到输出的所有权重：

这一剪枝粒度有助于以均匀的比率去除每个输出的权重，以防止不同输出特征之间的不平衡剪枝。作者也在图像分类模型上做了类似的实验，并没有类似的趋势。这表明以“每个输出”为粒度进行剪枝可能是由于语言模型任务的某些特性导致的。

2.2 算法流程

Wanda 的具体流程如算法 1 所示。先在参考样本上进行单次前向传递，用于估计特征范数，接下来可以根据式（1）直接计算每个权重的重要性，并按照式（2）的剪枝粒度进行修剪。Wanda 不需要权重更新和微调，在修剪完毕后可以直接使用。

2.3 扩展&对比

N:M 结构化稀疏：Wanda可以自然地扩展到结构化的 N:M 稀疏性，只需要在每 M 个连续权重中使用相同的度量来比较与输出连接的所有权重的权重。

对比分析：SparseGPT 是最优脑手术 OBS [3] 在 LLM 剪枝的复兴，而 Wanda 可以视为最优脑损伤 OBD [4] 在 LLM 剪枝的复兴。具体对比见表 1。

▲ 表1. Wanda与LLM上现有剪枝算法的对比

实验

1. 语言模型：如表 2 所示，Wanda 在任何剪枝模式下都大大优于幅值剪枝，且在完全不更新权重的情况下和 SparseGPT 持平。从图 2 中也可以看出 Wanda 解决了幅值剪枝的性能退化问题，随着稀疏率的增加，Wanda 和 SparseGPT 呈现相近的退化趋势。

▲ 表2. WikiText数据集上LLaMA模型的剪枝方法对比

▲ 图2. Wanda和其他剪枝方法在LLaMA上的性能退化趋势

2. Zero-shot 任务：表 3 对比了不同剪枝方法在下游 zero-shot 任务的性能表现，Wanda 远超幅值剪枝，与 SparseGPT 各有优劣。

▲ 表3. 7个zero-shot任务下剪枝算法的性能对比

3. 剪枝成本：表 4 对比了在不同模型规模下 Wanda 和 SparseGPT 的剪枝成本，与 SparseGPT 相比 Wanda 的时间开销可以忽略不计，在 65B 的模型上至需要 5.6 秒。

▲ 表4. Wanda和SparseGPT的时间开销（单位：秒），不包括共享的前向传递过程

4. 消融实验：表 5 分析了不同剪枝度量和剪枝粒度的性能对比，证明了适当的剪枝粒度对于修剪 LLMs 的重要性。图 3 分析了不同参考样本下剪枝算法的性能变化，样本越多性能越好，Wanda 比 SparseGPT 对少样本更鲁棒。表 6 尝试将 SparseGPT 的权重更新策略与 Wanda 结合，实验发现权重更新不能帮助 Wanda 进一步提升性能。

▲ 表5. 剪枝度量和剪枝粒度的消融实验

▲ 图3. Wanda对少样本更鲁棒

▲ 表6. 权重更新对Wanda没用

5. 扩展-图像分类：图 4 分析了 Wanda 在 ViT 图像分类任务的性能，在图像分类任务上，Wanda 仍优于幅值剪枝，但没有有效缓解性能退化的问题，我猜测这是因为实验所采用的 ViT 过小，没有出现类似 LLM 的“紧凑大幅值特征”的特性。

参考文献

[1] Elias Frantar and Dan Alistarh. SparseGPT: Massive language models can be accurately pruned in one-shot. arXiv preprint arXiv:2301.00774, 2023.

[2] Tim Dettmers, Mike Lewis, Younes Belkada, and Luke Zettlemoyer. LLM.int8(): 8-bit matrix multiplication for transformers at scale. In Advances in Neural Information Processing Systems, 2022.

[3] Babak Hassibi, David G Stork, and Gregory J Wolff. Optimal brain surgeon and general network pruning. In IEEE International Conference on Neural Networks, 1993.

[4] Yann LeCun, John S Denker, and Sara A Solla. Optimal brain damage. In Advances in Neural Information Processing Systems, 1989.

更多阅读