《Transformer 关键要素：非所有注意力皆必要》论文解读

摘要：随着基于Transformer的大型语言模型（LLMs）在各种任务中展现出强大性能，其规模的扩大也带来了冗余结构的问题。本文通过基于相似性的度量方法，研究了Transformer中不同模块（包括块、多层感知机（MLP）和注意力（Attention）层）的冗余情况。发现大量注意力层存在过高的相似性，可在不降低性能的情况下进行修剪，从而降低内存和计算成本。还提出了一种联合丢弃注意力和MLP层的方法，提高了性能和丢弃比例。

一、研究背景和动机
（一）LLMs的发展与挑战Transformer-basedLLMs在AI研究中取得了显著进展，但模型的扩展导致冗余模块增加，这不仅增加了部署成本和资源需求，还限制了其在边缘设备上的应用。例如GPT-3的参数规模巨大，对内存要求过高。
（二）对模型结构冗余性的关注以往的研究虽有优化LLM效率的工作，但常忽略Transformer架构的独特性。Transformer由多个堆叠的块组成，每个块包含MLP层和Attention层，不同模块可能存在不同类型和程度的冗余，这激发了作者对这些不同模块冗余性的研究。

二、研究方法
（一）基于相似性的度量方法通过计算模块输入和输出之间的余弦相似性来确定模块的重要性。相似性越高，模块的重要性得分越低，越有可能被丢弃。公式为，其中为输入，为输出。
（二）不同模块的丢弃方法块丢弃（BlockDrop）计算每个块的输入和输出的相似性得分，从浅层到深层迭代计算每个块的重要性得分，丢弃得分最低的块。层丢弃（LayerDrop）MLP丢弃（MLPDrop）：考虑MLP层及其后续的层归一化（LayerNorm）作为一个整体来计算相似性得分，同时丢弃不重要的MLP层和相关的LayerNorm层。注意力丢弃（AttentionDrop）：类似MLPDrop，考虑注意力层和相关LayerNorm层的整体输出，计算重要性得分。联合层丢弃（JointLayerDrop）先分别获取注意力层和MLP层的重要性得分和，然后将它们连接起来，根据综合得分丢弃层，以优化丢弃过程。

三、实验结果
（一）不同模块丢弃的比较块丢弃在Llama-2-13B上进行实验，丢弃8个块时，平均性能分别下降了7.5%，性能下降明显，说明块对于维持性能较为重要。MLP层丢弃同样在Llama-2-13B和Mistral-7B上实验，丢弃8个MLP层时，性能分别下降了6.3%和6.9%，也会导致性能大幅下降。注意力层丢弃与前两者不同，在Llama-3-13B和Mistral-7B上丢弃8个注意力层后，模型仍能保持超过99%的原始性能。而且在不同的丢弃比例实验中，当丢弃比例低于50%时，性能保持相对稳定。
（二）模块丢弃的效率提升内存使用注意力层丢弃可以有效减少键值（KV）缓存，例如对于Llama-2-13B的一个输入，可将约70GB的KV缓存减少一半，显著降低内存需求。速度提升与MLP层相比，注意力层丢弃对速度提升的贡献更大，如在Mistral-7B上丢弃8层时，注意力层丢弃的速度提升倍数为1.17×，而MLP层丢弃为1.09×。
（三）联合层丢弃的效果联合层丢弃的性能明显高于单独的注意力层丢弃或MLP层丢弃。它首先只丢弃注意力层，直到达到一定数量后再考虑丢弃MLP层，提高了丢弃比例和性能。

四、进一步分析
（一）注意力层冗余的一致性训练过程中的相似性通过评估MAP-Neo-7B在不同训练阶段的检查点，发现注意力层在所有训练阶段的重要性得分都极低，而MLP层和块的重要性得分随着训练逐渐增加。更深层模块的冗余性可视化不同模型的丢弃层或块，发现无论是Llama-2-13B、Mistral-7B还是更大的Llama-2-70B，都倾向于先丢弃更深层的模块，说明更深层模块往往更冗余。
（二）块丢弃和层丢弃的鲁棒性对样本数量的鲁棒性随着样本数量增加（从4到1024），不同注意力层的特征相似性保持相对稳定，说明使用256个样本计算相似性足以满足实验需求。对校准数据集的鲁棒性改变校准数据集（从预训练数据集到指令调整数据集），特征相似性没有显著变化，模块的丢弃情况也基本相同，证明了块丢弃和层丢弃对校准数据集变化的适应性。
（三）不同模型规模的影响更大模型的鲁棒性在Llama-2-70B上进行实验，发现它对块丢弃和MLP层丢弃也敏感，但相比较小模型更具鲁棒性。例如，丢弃8个MLP层时，Llama-2-70B平均性能仅下降0.6%，而Llama-2-13B丢弃4个MLP层时下降3.9%。注意力层丢弃在更大模型上的表现在Llama-2-70B上丢弃32个注意力层时，性能平均仅下降0.1%，而在Llama-2-13B上丢弃相同比例的注意力层会导致明显的性能下降。五、研究贡献开发了基于相似性的度量方法来衡量LLMs内的冗余，并研究了包括块、MLP和注意力层在内的各种模块的结构冗余。揭示了注意力层令人惊讶的冗余性，并证明了注意力层丢弃对准确性、内存和计算效率的好处。提出了联合层丢弃等技术，进一步提高了模块丢弃的有效性和效率。所提出的层丢弃方法能有效识别重要和冗余层，为现有模型评估提供了关键信息，也为网络架构设计的未来研究提供了有价值的见解。六、研究局限所讨论的丢弃技术虽在研究的模型上提高了效率，但还需进一步探索其在更广泛模型（如视觉Transformer和视觉-语言模型）中的应用。主要关注训练后的丢弃，未涉及重新训练，而重新训练可能进一步恢复或提高性能。更多内容欢迎来卡奥斯智能交互引擎搜索探讨