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 占据了大约 或 66.7% 的参数比例。
之前有一些工作提出了通过去除冗余参数或架构来提高LLM效率的方法。量化和剪枝是压缩LLMs最广泛使用的技术。具体来说,量化将数据类型转换为低比特,但仍然可能存在冗余的架构和参数。剪枝可以分为无结构剪枝和结构化剪枝。虽然无结构剪枝比结构化剪枝保持更好的性能,但它不能有效地应用于硬件,限制了其实际使用。
但这些方法通常采用通用技术,忽视了Transformer架构的独特特性。具体来说,Transformer架构由多个堆叠的块组成,每个块包含一个MLP层和一个注意力层,它们执行不同的功能,并表现出相应的不同冗余水平。
4. Transformer关键组成冗余性研究及丢弃策略实验
作者探索了Transformers中的三个关键组成部分:块、MLP和注意力层的冗余性。使用基于相似性的度量方法,评估了每个组成部分的重要性,并逐步剪枝那些被识别为冗余的部分。
相似性度量方法如下,余弦相似性越高的模块表现出越低的重要性分数,表明冗余。根据预定义的剪枝比例识别并剪枝最低重要性分数的模块。
计算每个块的输入和输出之间的相似性。对于第l个块,重要性得分计算如下:
其中,和分别表示第l个块的输入和输出。
同理,该公式也可以用在注意力层的输入输出相似性度量。注意力层在残差连接中操作。第l个注意力层的输出计算如下:
其中,是相应LayerNorm层的输入,而是涉及残差连接的整体输出。像MLP Drop一样,将注意力层及其关联的LayerNorm视为一个单元。注意力层的重要性得分为:
对于MLP和注意力层,Layer Drop以一次性方式执行,仅计算一次重要性得分并在单一步骤中去除冗余层。
首先应用“块丢弃”方法,但观察到去除整个块会导致显著的性能下降。这表明需要更细粒度的策略。下图展示了块丢弃以及块的重要性分数。
进一步探索MLP和注意力层的单独剪枝。发现虽然丢弃MLP层会对性能产生负面影响,但相当一部分注意力层,即Transformer架构的核心,可以被剪枝而不降低模型的性能。例如,在Llama-2-70B中丢弃50%的注意力层,结果与完整模型的性能相当,表明这些层中存在高度的冗余。
基于上述分析,论文提出一种更灵活的方法,“联合层丢弃”,针对MLP层和注意力层。通过结合这些层的重要性分数,发现联合丢弃低重要性的注意力和MLP层在高稀疏条件下比仅剪枝一种类型的层获得更好的性能。注意力层的冗余性不仅明显,而且在整个训练过程中保持一致,表明这种冗余性是Transformer架构的固有属性。
为了研究丢弃不同目标模块的影响,量化性能下降与加速之间的权衡,引入了一种新的度量标准,即加速退化比率(SDR),定义为:
其中 Δ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层的冗余性,以提升模型整体性能。
联合层丢弃方法首先为注意力层()和MLP层()单独计算重要性评分。这些评分基于识别冗余层的相似性指标进行计算。一旦获得了每种层的重要性评分,将其合并为一个数组:。从这个合并的重要性评分集中,丢弃值最低的层,无论它们是注意力层还是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】,后续会带来分享。
扩展阅读:
《全方位解读大模型:多样知识点的深度探讨与技术分享小结》
6. 参考文献
【1】What Matters in Transformers? Not All Attention is Needed
【2】Jamba: A Hybrid Transformer-Mamba Language Model