【大模型理论篇】大模型压缩技术之注意力层剪枝以及与MLP层联合剪枝

1. 背景分析

本来打算写一篇关于大模型蒸馏的文章，但刚好看到近期发表的一篇讨论大模型压缩的文章【1】，是关于注意力机制冗余性的讨论，比较有意思，作者分析得出并不是所有的注意力都是必须的，可以通过对模型去除冗余的结构，来达到缩小模型规模的目标。

2. 论文结论

接下来我们具体来看下该论文的分析。在这篇论文中，作者探讨了在大模型中，并非所有的注意力机制都是必需的，他们通过使用基于相似性的度量方法来研究Transformer模型中不同模块的冗余性，包括块（Blocks）、多层感知机（MLP）和注意力（Attention）层。研究发现，尽管注意力层在区分Transformer和其他架构中扮演着关键角色，但许多注意力层的相似性过高，可以被剪枝而不影响性能。例如，通过剪枝一半的注意力层，Llama-2-70B模型在MMLU任务上仅损失了2.4%的性能，但速度提升了48.4%。此外，通过追踪训练过程中的模型检查点，作者观察到注意力层的冗余是固有的，并且在训练阶段是一致的。还提出了一种同时丢弃注意力和MLP层的方法，允许更积极地丢弃额外的层。例如，在丢弃31层（注意力+MLP）时，Llama-2-13B在MMLU任务上仍保持了90%的性能。

3. Transformer结构及压缩技术讨论

从近几年的AIGC发展来看，基于Transformer的大模型做出了很大的贡献，我们在文章《通用大模型架构分类及技术统一化》中对主流大模型的架构已经做了分析。但扩展这些模型也引入了冗余的架构，导致效率降低，而且增加了部署成本和资源需求。Transformer模型（关于transformer模型的解析可以参考《Transformer原理及关键模块深入浅出》）由多个块组成，包括注意力层和MLP层。注意力层计算输入令牌之间的上下文信息，其复杂度与输入序列长度呈二次方关系。KV-Cache 《Transformer KV Cache原理深入浅出》）减轻了计算问题，但导致了过高的内存成本。MLP层独立地转换每个令牌，使用上投影后跟下投影，并贡献了大部分模型参数，我们在《Transformer原理及关键模块深入浅出》中给出了FFN参数量理论值，FFN 参数数量是自注意力参数数量的两倍，即 FFN 占据了大约 $\frac{2}{3}$ 或 66.7% 的参数比例。

之前有一些工作提出了通过去除冗余参数或架构来提高LLM效率的方法。量化和剪枝是压缩LLMs最广泛使用的技术。具体来说，量化将数据类型转换为低比特，但仍然可能存在冗余的架构和参数。剪枝可以分为无结构剪枝和结构化剪枝。虽然无结构剪枝比结构化剪枝保持更好的性能，但它不能有效地应用于硬件，限制了其实际使用。

但这些方法通常采用通用技术，忽视了Transformer架构的独特特性。具体来说，Transformer架构由多个堆叠的块组成，每个块包含一个MLP层和一个注意力层，它们执行不同的功能，并表现出相应的不同冗余水平。

4. Transformer关键组成冗余性研究及丢弃策略实验

作者探索了Transformers中的三个关键组成部分：块、MLP和注意力层的冗余性。使用基于相似性的度量方法，评估了每个组成部分的重要性，并逐步剪枝那些被识别为冗余的部分。

相似性度量方法如下，余弦相似性越高的模块表现出越低的重要性分数，表明冗余。根据预定义的剪枝比例识别并剪枝最低重要性分数的模块。

       计算每个块的输入和输出之间的相似性。对于第l个块，重要性得分计算如下：

$S_l^B = 1 - \text{CosineSim}(X_l^B, Y_l^B)$

其中， $X_l^B$ 和 $Y_l^B$ 分别表示第l个块的输入和输出。

同理，该公式也可以用在注意力层的输入输出相似性度量。注意力层在残差连接中操作。第l个注意力层的输出计算如下：
   $\mathbf{Y}_A^l = \mathbf{X}_A^l + \text{Attention}(\text{LayerNorm}(\mathbf{X}_A^l))$
        其中， $\mathbf{X}_A^l$ 是相应LayerNorm层的输入，而 $\mathbf{Y}_A^l$ 是涉及残差连接的整体输出。像MLP Drop一样，将注意力层及其关联的LayerNorm视为一个单元。注意力层的重要性得分为：
   $S_A^l = 1 - \text{CosineSim}(\mathbf{X}_A^l, \mathbf{Y}_A^l)$
        对于MLP和注意力层，Layer Drop以一次性方式执行，仅计算一次重要性得分并在单一步骤中去除冗余层。

首先应用“块丢弃”方法，但观察到去除整个块会导致显著的性能下降。这表明需要更细粒度的策略。下图展示了块丢弃以及块的重要性分数。

进一步探索MLP和注意力层的单独剪枝。发现虽然丢弃MLP层会对性能产生负面影响，但相当一部分注意力层，即Transformer架构的核心，可以被剪枝而不降低模型的性能。例如，在Llama-2-70B中丢弃50%的注意力层，结果与完整模型的性能相当，表明这些层中存在高度的冗余。

基于上述分析，论文提出一种更灵活的方法，“联合层丢弃”，针对MLP层和注意力层。通过结合这些层的重要性分数，发现联合丢弃低重要性的注意力和MLP层在高稀疏条件下比仅剪枝一种类型的层获得更好的性能。注意力层的冗余性不仅明显，而且在整个训练过程中保持一致，表明这种冗余性是Transformer架构的固有属性。

为了研究丢弃不同目标模块的影响，量化性能下降与加速之间的权衡，引入了一种新的度量标准，即加速退化比率（SDR），定义为：

$\gamma = \frac{\Delta \text{Avg.}}{\Delta \text{Speedup}}$

其中 ΔAvg 表示评估任务的平均性能变化百分比，而 ΔSpeedup 表示每种方法实现的相应加速百分比。因此，γ衡量每增加1%加速所导致的性能下降。较低的 γ 值表明模型在保持最小性能损失的情况下实现加速，使其更高效。相反，较高的 γ值表明相对于获得的加速，性能损失显著，意味着权衡不太理想。

表1显示了丢弃固定数量模块（例如4和8层）对性能的影响。Block和MLP的丢弃都导致两种模型的性能显著下降，尽管实现了适度的加速。例如，丢弃4个块导致Llama-2-13B的平均性能下降2.4%（从68.2降至65.8），加速为1.11×，对应的γ为0.22。然而，丢弃8个块会导致7.5%的性能下降（降至60.7），仅实现1.24×的适度加速，且γ为0.31。

不过实验发现，尽管Attention层在Transformer架构中扮演关键角色，但丢弃Attention层的效果非常好。无论是Llama-2-13B还是Mistral-7B，在丢弃8个Attention层后，仍保持超过99%的原始性能。举例来说，在丢弃8层后，Attention Drop几乎维持基线性能（69.8对比70.3），加速为1.23×，γ为0.02。丢弃12个Attention层仅导致轻微的性能下降（67.3），同时获得1.40×的显著加速，γ为0.08。Attention层的冗余性较高，去除它们对模型准确性影响最小，使得Attention Drop成为一种有效的剪枝策略。

更大模型在Attention Drop下表现出一致的鲁棒性。使用Llama-2-70B，在表3中丢弃不同丢弃比例（从5%到60%）的模块。与小型模型的发现类似，Llama-2-70B对Block Drop和MLP Drop也表现出敏感性，仅丢弃10%到20%的块或MLP层就会导致显著的性能下降。

相比之下，Attention Drop在Llama-2-70B上表现更佳。当丢弃80层中的40层时，Llama-2-70B实现了1.48的加速和γ。Llama-3中的类似趋势在表2中得到了体现。这种鲁棒性表明，更大模型在移除大量Attention层的情况下仍能保持性能。

结果表明，Attention Drop是有效的剪枝方法，能够在性能影响最小的情况下实现显著的加速。接下来通过Attention Drop实现的效率提升的分析，进一步研究Attention层的逐层重要性，深入了解模型架构。

5. Attention Drop的效率评估

进一步地，论文评估Attention Drop在内存使用和推理速度方面的效率。具体来说，通过检查由于键值（KV）缓存导致的内存开销的减少，并测量整个生成阶段的速度提升。结果表明，Attention Drop在这两个效率指标上均提供了显著的改进，同时保持高性能。

5.1 KV缓存内存减少

考虑到LLM的自回归特性，输出是逐个生成的，KV缓存用于存储输入序列的中间表示。该缓存通过防止冗余计算来加速推理，但在较长序列或更大批次的情况下，内存成本显著。所提出的Attention Drop方法有效地移除了不重要的注意力层，从而减少了相应的KV缓存。表4提供了不同模型在应用Attention Drop前后的16位精度KV缓存内存使用情况的比较。在70B模型中使用了8块Nvidia RTX A6000 Ada GPU，其他较小模型则使用4块。结果显示，Attention Drop在所有测试模型中均实现了显著的内存节省。例如，在Llama-2-13B中，KV缓存从52GB减少到26GB，减少了50%。

5.2 速度测量

论文还评估通过Attention Drop实现的运行时速度提升。推理速度在整个生成过程中进行测量，从输入提示开始到最后一个标记的生成。为确保结果准确反映速度提升，在设置中遵循两个关键原则：（1）所有操作在单个Nvidia RTX A6000 Ada GPU上执行，避免多GPU设置引起的通信开销；（2）增大批次以最大化每个模型的GPU利用率。具体而言，对于Llama-2-70B，由于其较大的模型尺寸，采用了4位量化，同时注意到Attention Drop与C中展示的量化是正交的。对于Llama-2-13B和Mistral-7B，使用16位精度。在序列长度方面，使用2048个标记的输入序列，并自回归生成额外的2048个标记。此设置使我们能够捕捉完整的推理过程，确保预填充和生成阶段均包含在速度测量中。结果表明，Attention Drop在保留超过95%原始模型性能的同时，提供了高达40%的速度提升。此外，Attention Drop的γ值显著低于MLP Drop和Block Drop，尤其是在更高的速度提升比率下。这表明，Attention Drop在速度和性能之间实现了更高效的权衡。

为了进一步分析丢弃的模块，作者可视化了不同丢弃比率下的丢弃层或块。图示可视化了随着丢弃模块数量增加而剩余和丢弃的层/块。Llama-2-13B和Mistral-7B在Layer Drop和Block Drop中表现出相似的模式：最初，这两种模型倾向于丢弃深层，随后是浅层。深层往往冗余更高。较大的模型（例如Llama-2-70B）也展现了类似的趋势。这个也比较容易理解，一开始肯定输入和输出经过层学习后差异更大，后面随着模型参数训练趋于收敛稳定，相似度更高。我们早在2019年使用bert模型的时候就采用了去除更深的块来提速，效果降低的幅度在可控的范围内。

5.3 注意力层在训练过程中的一致冗余

鉴于深层表现出高冗余，为了研究这种模式的形成，重新查看历史检查点，以追踪层级重要性评分的动态变化。具体而言，使用由MAP-Neo-7B Zhang et al. (2024)发布的检查点，因其在训练阶段发布的连续检查点。

图10展示了不同训练阶段Block和Layer的重要性评分，MAP-Neo-7B预训练过程中检查点的重要性评分可视化，其中较亮区域表示低重要性评分（即高相似性评分）。其中Attention层在所有训练阶段均表现出明显低于MLP和Block的重要性评分。

尽管MLP层和Block的重要性评分随着训练的推进逐渐上升，Attention层的重要性评分变化则相对缓慢。Attention层在整个训练过程中一致性地表现出较高的冗余性，作者认为这种模式源于Attention层的固有特性。尽管相当比例的注意力层显示出高冗余，但也同样发现一些MLP层的重要性较低。为了进一步优化模型效率，引入了联合层丢弃方法，结合Attention Drop和MLP Drop策略。这种方法利用了注意力和MLP层的冗余性，以提升模型整体性能。

联合层丢弃方法首先为注意力层（ $S_{Al}$ ）和MLP层（ $S_{Ml}$ ）单独计算重要性评分。这些评分基于识别冗余层的相似性指标进行计算。一旦获得了每种层的重要性评分，将其合并为一个数组： $S = [S_{Al}, S_{Ml}]$ 。从这个合并的重要性评分集中，丢弃值最低的层，无论它们是注意力层还是MLP层。这种联合方法使得能够同时移除两个层类型中最冗余的组件，从而提高模型效率，同时保持性能。

如图所示，联合层丢弃始终比单独的Attention Drop或MLP Drop表现出更好的性能。该过程首先独立丢弃注意力层，这些层通常比MLP层更冗余。此过程持续到Mistral-7B丢弃的注意力层数量超过14层，Llama-2-13B超过18层。在剪枝的初始阶段，联合层丢弃的性能与Attention Drop相重叠。

然而，随着丢弃比例的增加，更冗余的注意力层被修剪后，MLP层开始成为下一个冗余组件。此时，联合层丢弃开始移除MLP层，从而在不显著损失性能的情况下进一步减少冗余层，例如，在丢弃31层（Attention + MLP）后，Llama-2-13B在MMLU任务上仍保留90%的性能。

总结：注意力层中存在的高冗余性，尤其是在深层，可以减少注意力层的数量而不损害性能，而不是保持与MLP层的相对一致性。此外，与MLP层不同，注意力层在训练过程中表现出一致的冗余性。这种一致性可能成为训练大型模型的瓶颈。近期有研究提出将Transformer与Mamba结合，比如Jamba来提升吞吐量【2】，后续会带来分享。

扩展阅读：

《全方位解读大模型：多样知识点的深度探讨与技术分享小结》