边缘计算设备部署LLM？Nota提出大模型专用剪枝方法Shortened-LLaMA

目前，大型语言模型（LLM）的一个重要痛点就是高额的计算需求，同一大模型家族中，即便是最小的模型版本也拥有7B的参数量，这给一些实际的落地场景带来了严重的负担。目前学术界和工业界已有一些对LLM效率进行优化的工作，例如对LLM进行结构化剪枝（structured pruning）。目前的大多数研究都集中在纯宽度剪枝或宽度与深度剪枝的混合上，很少分析这两种剪枝模式（宽度与深度）影响 LLM 推理效率的核心原因。

本文介绍一篇来自Nota Inc的最新工作，本文作者对现阶段大模型剪枝技术进行了深入探索，并提出了一种全新的LLM剪枝方法，称为Shortened LLaMA。Shortened LLaMA在zero-shot任务性能方面与目前最新的宽度剪枝方法相媲美，同时提高了模型的整体推理速度。本文作者提到，Shortened LLaMA非常适合与在本地或边缘计算设备上部署LLM。

论文题目： Shortened LLaMA: A Simple Depth Pruning for Large Language Models 论文链接： https://arxiv.org/abs/2402.02834

现有的大模型剪枝技术可以大致分为宽度剪枝和深度剪枝两种，其中宽度剪枝减少了投影权重矩阵的大小（可以通过删除注意力头的数量来实现），同时保持层数。相反，深度剪枝会删除整个层或块，同时保持剩余权重的大小不变。两种剪枝模式的差异如下图所示。

当然，目前也存在一些同时兼顾宽度和深度的剪枝方法，例如Sheared-LLaMA[1]可以在缩小网络宽度的同时降低网络深度。此外，由于LLM在推理阶段采用自回归解码机制，因而在生成预测过程中会占用大量显存。如果想要提高GPU的利用率和吞吐量，最直接的方法是加大模型输入的batchsize，但这种方法在显存容量较小的GPU上是不可行的，尤其是在一些边缘计算设备上。

本文提出了一种简洁的深度剪枝方法Shortened LLaMA，其可以在zero-shot任务方面与一些宽度剪枝方法（LLM-Pruner、FLAP等方法）相媲美。此外，Shortened LLaMA可以有效提升模型的推理速度，尤其在batchsize=1的情况下（如上图所示）。

二、本文方法 2.1 加大批次推理的问题

由于目前大多数LLM本质上都是自回归模型，其中每个token的生成通常涉及大矩阵（权重）与小矩阵或向量（激活）相乘的操作，这种情况推理效率的主要瓶颈并不是数学运算的速度，而是对大型参数的内存访问操作。下图分别展示了在不同 NVIDIA GPU 上运行 LLaMA-7B 模型的GPU计算利用率情况。增加批次大小（batchsize）确实可以提高GPU的利用率，但是过度增加会带来显存OOM的问题。

2.2 模块重要性评估

目前流行的LLM大多仍采用Transformer架构，由一系列基础Transformer块堆叠而成，每个模块包含一个多头注意力（MHA）和前馈网络（FFN）模块。本文选择基础Transformer块作为剪枝目标单元，剪枝策略是先对网络中所有基础模块进行重要性评估，随后根据评估的结果对不重要的区块进行剪枝和轻度的再训练。 作者研究了多种重要性评估指标，包括幅度（Mag）、泰勒值（Taylor）、Mag+、Taylor+和困惑度（PPL）。

（1）幅度(Mag)

Mag是剪枝技术中最常用的基本基准，Mag假设范数较小的权重包含的信息量也较少，在本文的模块重要性分析中，Mag可以直接表示为：

（2）泰勒值(Taylor)

衡量模块重要性也可以通过最直接的loss变化来判断，例如删除模型中某部分权重参数，然后观察模型误差的变化。给定一个经过校准的数据集 ，删除部分权重后的loss值可以计算为：

可以直接对上式计算二阶导数进行优化，因而判断模块重要性分数的Taylor值可以表示如下：

（3）Mag+和Taylor+

Mag+和Taylor+分别是上述两个指标的启发式版本，如果单独使用Mag和Taylor进行判断后，直接整体去除重要性较低的模块会导致模型性能严重下降，因而作者加入了一种启发式策略，对这些低重要性模块，保留其中的少量结构。

（4）困惑度(PPL)

PPL是语言建模任务中常用的性能指标，模型中的冗余模块对模型输出的贡献较小，删除这些模块时，模型的PPL指标会发生微小的退化。这里作者同样使用校准数据集 进行实验，PPL指标定义如下：

其中， 表示不含第 个模块的模型，使用 PPL 可以反映模型的整体行为，因为PPL的计算需要从下一个token的预测损失中进行推导。其只需要前向传递，避免了计算反向传播梯度。如下图所示，有几个模块被删除后，PPL指标仅有轻微变化，因而将其识别为需要进行剪枝。而删除初始块和最终块会明显降低性能，因此需要对它们进行保留。

作者进行了下表的实验后，最终选择了Taylor+ 和 PPL 指标作为Shortened LLaMA的模块重要性指标。

具体来说，一个LLM模型的线性权重矩阵可以表示为 ，其中 代表第 个Transformer模块中的操作类型。权重重要性分数是在输出神经元层面计算的，然后将这些分数相加起来作为整体模块的重要性。

2.3 轻度Retraining

在根据权重重要性指标对模块进行剪枝后，作者对整体模型进行轻度的LoRA再训练，剪枝网络的权重矩阵可以表示为 ，其中 表示初始预训练权重，其维数为 。更新矩阵 被分解成两个可训练的部分 和 ，维数分别为 和 ，其中 代表低秩。

三、 实验结果 3.1 对比实验

本文的实验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]是一种结构化剪枝方法，其基于权重大小和输入激活范数的乘积，这可以作为局部剪枝的重要性指标。 下表展示了本文方法与这些方法的性能对比情况。

3.2 zero-shot性能对比

下表展示了Shortened LLaMA在Vicuna-7B模型上zero-shot下游任务性能的表现，作者同时列出了模型的推理效率。从表中可以看出，本文方法在使用Taylor+和PPL时的性能均超越其他宽度剪枝方法。 这是由于宽度方法通常仅减小权重矩阵的大小，这并不能提高模型的速度，因为在batchsize有限的情况下，token生成的速度很大程度上取决于 LLM 内存访问操作的频率。 由于Shortened LLaMA直接删除了重要性程度较低的Transformer块来实现加速，减少了激活和权重之间的内存访问次数，同时还能保持较好的下游任务适应能力。

3.3 对剪枝结构单元粒度的分析

下表展示了本文方法在不同粒度剪枝情况下的消融实验结果，对于参数规模超过 5B 的模型，删除单个 MHA 和 FFN 模块相比删除整个 Transformer 模块会带来更好的下游任务精度，但 PPL 会更差。而对于参数规模小于 5B 的较小模型，局部剪枝在所有检查指标方面都取得了优异的结果，作者提到，这与通常认为去除更精细的单元会产生更好性能的观点不同。

3.4 一次性剪枝与迭代剪枝哪个更好

作者还探索了Shortened LLaMA在一次性剪枝（One-shot）和迭代剪枝（Iterative）模式中的差别，结果如下图所示，其中分别展示了在剪枝前后模型在训练时的表现，可以看出，迭代剪枝模型比一次性剪枝模型产生更好的后剪枝结果。然而，如果对一次性剪枝后的模型进行轻度再训练，就可以达到与迭代剪枝相似的性能。考虑到对剪枝后的模型进行轻度再训练的成本也不高，因此作者选择一次性剪枝作为最佳剪枝模式。

四、 总结

本文引入了一种全新的局部LLM剪枝方法Shortened LLaMA，本文首先对类Transformer LLM的网络宽度和深度剪枝技术进行了深入的对比分析。具体来说，作者发现，在输入批次有限的情况下，现有的宽度剪枝算法难以实现对自回归LLM的性能提速。

Shortened LLaMA是一种简单而有效的 LLM 深度剪枝策略，其通过多种模块重要性评估手段来确定模型中的冗余模块，对Transformer模块进行多种粒度的剪枝。此外，作者还为Shortened LLaMA设置了低成本的再训练模式，这使得模型在提高推理效率的同时，仍然能够保持一定的下游任务适应能力，本文方法在zero-shot任务方面展现出了优越的性能。

参考

[1] Mengzhou Xia, Tianyu Gao, Zhiyuan Zeng, and Danqi Chen. Sheared llama: Accelerating language model pre-training via structured pruning. In ICLR, 2024.

[2] Xinyin Ma, Gongfan Fang, and Xinchao Wang. Llm-pruner: On the structural pruning of large language models. In NeurIPS, 2023.

[3] Yongqi An, Xu Zhao, Tao Yu, Ming Tang, and Jinqiao Wang. Fluctuation-based adaptive structured pruning for large language models. In AAAI, 2024.

[4] Mingjie Sun, Zhuang Liu, Anna Bair, and J. Zico Kolter. A simple and effective pruning approach for large language models. In ICLR, 2024.

Illustration From IconScout By 22

-The End-

扫码观看！

本周上新！

“AI技术流”原创投稿计划

TechBeat是由将门创投建立的AI学习社区（www.techbeat.net）。社区上线500+期talk视频，3000+篇技术干货文章，方向覆盖CV/NLP/ML/Robotis等；每月定期举办顶会及其他线上交流活动，不定期举办技术人线下聚会交流活动。我们正在努力成为AI人才喜爱的高质量、知识型交流平台，希望为AI人才打造更专业的服务和体验，加速并陪伴其成长。

投稿内容

// 最新技术解读/系统性知识分享 //

// 前沿资讯解说/心得经历讲述 //

投稿须知

稿件需要为原创文章，并标明作者信息。

我们会选择部分在深度技术解析及科研心得方向，对用户启发更大的文章，做原创性内容奖励

投稿方式

发送邮件到

chenhongyuan@thejiangmen.com

或添加工作人员微信（chemn493）投稿，沟通投稿详情；还可以关注“将门创投”公众号，后台回复“投稿”二字，获得投稿说明。

关于我“门”

将门是一家以专注于数智核心科技领域的新型创投机构，也是北京市标杆型孵化器。 公司致力于通过连接技术与商业，发掘和培育具有全球影响力的科技创新企业，推动企业创新发展与产业升级。

将门成立于2015年底，创始团队由微软创投在中国的创始团队原班人马构建而成，曾为微软优选和深度孵化了126家创新的技术型创业公司。

如果您是技术领域的初创企业，不仅想获得投资，还希望获得一系列持续性、有价值的投后服务，欢迎发送或者推荐项目给我“门”:

bp@thejiangmen.com

点击右上角，把文章分享到朋友圈