文章
- 产品
- 文章
，同一大模型家族中，即便是最小的模型版本也拥有7B的参数量，这给一些实际的落地场景带来了严重的负担。目前学术界和工业界已有一些对LLM效率进行优化的工作，。目前的大多数研究都集中在纯宽度剪枝或宽度与深度剪枝的混合上，很少分析这两种剪枝模式（宽度与深度）影响 LLM 推理效率的核心原因。
本文介绍一篇来自Nota Inc的最新工作，本文作者对现阶段大模型剪枝技术进行了深入探索，并提出了一种全新的LLM剪枝方法，称为Shortened LLaMA。Shortened LLaMA在zero-shot任务性能方面与目前最新的宽度剪枝方法相媲美，同时提高了模型的整体推理速度。本文作者提到，Shortened LLaMA非常适合与在本地或边缘计算设备上部署LLM。
现有的大模型剪枝技术可以大致分为宽度剪枝和深度剪枝两种，其中宽度剪枝减少了投影权重矩阵的大小（可以通过删除注意力头的数量来实现），同时保持层数。相反，深度剪枝会删除整个层或块，同时保持剩余权重的大小不变。两种剪枝模式的差异如下图所示。
当然，目前也存在一些同时兼顾宽度和深度的剪枝方法，例如Sheared-LLaMA[1]可以在缩小网络宽度的同时降低网络深度。此外，由于LLM在推理阶段采用自回归解码机制，因而在生成预测过程中会占用大量显存。如果想要提高GPU的利用率和吞吐量，最直接的方法是加大模型输入的batchsize，但这种方法在显存容量较小的GPU上是不可行的，尤其是在一些边缘计算设备上。
本文提出了一种简洁的深度剪枝方法Shortened LLaMA，其可以在zero-shot任务方面与一些宽度剪枝方法（LLM-Pruner、FLAP等方法）相媲美。此外，Shortened LLaMA可以有效提升模型的推理速度，尤其在batchsize=1的情况下（如上图所示）。
由于目前大多数LLM本质上都是自回归模型，其中每个token的生成通常涉及大矩阵（权重）与小矩阵或向量（激活）相乘的操作，这种情况推理效率的主要瓶颈并不是数算的速度，而是对大型参数的内存访问操作。下图分别展示了在不同 NVIDIA GPU 上运行 LLaMA-7B 模型的GPU计算利用率情况。增加批次大小（batchsize）确实可以提高GPU的利用率，但是过度增加会带来显存OOM的问题。
目前流行的LLM大多仍采用Transformer架构，由一系列基础Transformer块堆叠而成，每个模块包含一个多头注意力（MHA）和前馈网络（FFN）模块。本文选择基础Transformer块作为剪枝目标单元，剪枝策略是先对网络中所有基础模块进行重要性评估，随后根据评估的结果对不重要的区块进行剪枝和轻度的再训练。 作者研究了多种重要性评估指标，包括幅度（Mag）、泰勒值（Taylor）、Mag+、Taylor+和困惑度（PPL）。
Mag是剪枝技术中最常用的基本基准，Mag假设范数较小的权重包含的信息量也较少，在本文的模块重要性分析中，Mag可以直接表示为：
衡量模块重要性也可以通过最直接的loss变化来判断，例如删除模型中某部分权重参数，然后观察模型误差的变化。给定一个经过校准的数据集 ，删除部分权重后的loss值可以计算为：
Mag+和Taylor+分别是上述两个指标的启发式版本，如果单独使用Mag和Taylor进行判断后，直接整体去除重要性较低的模块会导致模型性能严重下降，因而作者加入了一种启发式策略，对这些低重要性模块，保留其中的少量结构。
PPL是语言建模任务中常用的性能指标，模型中的冗余模块对模型输出的贡献较小，删除这些模块时，模型的PPL指标会发生微小的退化。这里作者同样使用校准数据集 进行实验，PPL指标定义如下：
其中， 表示不含第 个模块的模型，使用 PPL 可以反映模型的整体行为，因为PPL的计算需要从下一个token的预测损失中进行推导。其只需要前向传递，避免了计算反向传播梯度。如下图所示，有几个模块被删除后，PPL指标仅有轻微变化，因而将其识别为需要进行剪枝。而删除初始块和最终块会明显降低性能，因此需要对它们进行保留。
作者进行了下表的实验后，最终选择了Taylor+ 和 PPL 指标作为Shortened LLaMA的模块重要性指标。
具体来说，一个LLM模型的线性权重矩阵可以表示为 ，其中 代表第 个Transformer模块中的操作类型。权重重要性分数是在输出神经元层面计算的，然后将这些分数相加起来作为整体模块的重要性。
在根据权重重要性指标对模块进行剪枝后，作者对整体模型进行轻度的LoRA再训练，剪枝网络的权重矩阵可以表示为 ，其中 表示初始预训练权重，其维数为 。更新矩阵 被分解成两个可训练的部分 和 ，维数分别为 和 ，其中 代表低秩。
本文的实验LLM选择了开源的LLaMA-7B和Vicuna-{7B, 13B}-v1.3，在剪枝和轻度再训练阶段，作者使用NVIDIA A100 GPU，此外作者也提到，对于7B模型版本，使用NVIDIA RTX 3090也可以进行。对比的baseline方法主要集中在目前较为流行的宽度剪枝方法，例如LLM-Pruner[2]主要基于Taylor值删除MHA中的注意力头和FFN的中间神经元。FLAP[3]主要根据重要性分数对模型进行全局剪枝。Wanda-sp[4]是一种结构化剪枝方法，其基于权重大小和输入激活范数的乘积，这可以作为局部剪枝的重要性指标。 下表展示了本文方法与这些方法的性能对比情况。
下表展示了Shortened LLaMA在Vicuna-7B模型上zero-shot下游任务性能的表现，作者同时列出了模型的推理效率。从表中可以看出，本文方法在使用Taylor+和PPL时的性能均超越其他宽度剪枝方法。 这是由于宽度方法通常仅减小权重矩阵的大小，这并不能提高模型的速度，因为在batchsize有限的情况下，token生成的速度很大程度上取决于 LLM 内存访问操作的频率。 由于Shortened LLaMA直接删除了重要性程度较低的Transformer块来实现加速，减少了激活和权重之间的内存访问次数，同时还能保持较好的下游任务适应能力。
下表展示了本文方法在不同粒度剪枝情况下的消融实验结果，对于参数规模超过 5B 的模。