Cross-layer Attention Sharing for Pre-trained Large Language Models

📄 arXiv: 2408.01890v2 📥 PDF

作者: Yongyu Mu, Yuzhang Wu, Yuchun Fan, Chenglong Wang, Hengyu Li, Jiali Zeng, Qiaozhi He, Murun Yang, Fandong Meng, Jie Zhou, Tong Xiao, Jingbo Zhu

分类: cs.CL

发布日期: 2024-08-04 (更新: 2025-10-17)

备注: A version accepted by TACL, prior to its publication by MIT Press

💡 一句话要点

LISA：通过跨层注意力共享，高效压缩预训练大语言模型中的冗余计算。

🎯 匹配领域: 支柱九：具身大模型 (Embodied Foundation Models)

关键词: 大语言模型 注意力机制 模型压缩 跨层共享 低秩近似 模型加速 推理优化

📋 核心要点

  1. 现有大语言模型注意力机制效率提升方法主要集中于KV缓存压缩和注意力头分组，忽略了层间注意力模式的冗余。
  2. LISA通过微型前馈网络对齐相邻层注意力头，并使用低秩矩阵近似层间注意力权重差异，从而实现跨层注意力共享。
  3. 实验表明，LISA在保持响应质量的同时，显著减少了冗余注意力计算，并实现了Q和K矩阵的压缩和吞吐量提升。

📝 摘要（中文）

为了提升大语言模型（LLMs）中注意力机制的效率，以往工作主要集中于压缩KV缓存或分组注意力头，而忽略了层间的冗余。本文通过对多种LLMs的全面分析表明，大多数层中存在高度相似的注意力模式。直观上，可以通过跨层共享注意力权重来减少冗余。然而，进一步的分析揭示了两个挑战：（1）在没有仔细重新排列注意力头的情况下直接共享权重矩阵是无效的；（2）浅层对注意力权重的微小偏差非常敏感。基于这些发现，我们提出了LISA，一种用于良好训练的LLMs中自注意力的轻量级替代方案。LISA采用微小的feed-forward网络来对齐相邻层之间的注意力头，并使用低秩矩阵来近似层间注意力权重的差异。包含13个典型基准的评估表明，LISA在准确性和困惑度方面保持了较高的响应质量，同时减少了总层数中53%-84%的冗余注意力计算。LISA的实现实现了注意力机制中Q和K矩阵的6倍压缩，LLaMA3-8B、LLaMA2-7B和LLaMA2-13B的最大吞吐量分别提高了19.5%、32.3%和40.1%。

🔬 方法详解

问题定义：论文旨在解决预训练大语言模型中注意力机制计算冗余的问题。现有方法主要关注KV缓存和注意力头，忽略了层间注意力模式的相似性，导致模型效率提升受限。直接共享注意力权重又面临注意力头不对齐和浅层对权重变化敏感的问题。

核心思路：论文的核心思路是利用层间注意力模式的相似性，通过跨层共享注意力权重来减少冗余计算。为了解决直接共享权重带来的问题，论文提出了LISA，通过对齐注意力头和近似层间权重差异来实现有效的跨层共享。

技术框架：LISA作为自注意力的替代方案，嵌入到预训练LLM中。其主要包含两个模块：1) 注意力头对齐模块，使用微型前馈网络对齐相邻层的注意力头；2) 注意力权重近似模块，使用低秩矩阵近似层间注意力权重的差异。整体流程是先对齐注意力头，然后使用低秩矩阵进行权重近似，最后替代原始的自注意力机制。

关键创新：LISA的关键创新在于其跨层注意力共享的策略，以及解决直接共享权重问题的具体方法。通过注意力头对齐和低秩矩阵近似，LISA能够在保持模型性能的同时，显著减少注意力计算的冗余。与现有方法相比，LISA更关注层间冗余，并提出了针对性的解决方案。

关键设计：注意力头对齐模块使用两层MLP，输入是相邻层的注意力头的输出，输出是对齐后的注意力头。低秩矩阵近似模块使用两个低秩矩阵分解层间注意力权重的差异，秩的大小是一个可调参数。损失函数主要包括语言建模损失和正则化损失，用于约束低秩矩阵的参数。

🖼️ 关键图片

img_0

📊 实验亮点

实验结果表明，LISA在13个典型基准测试中保持了较高的响应质量，同时减少了53%-84%的冗余注意力计算。LISA实现了Q和K矩阵的6倍压缩，LLaMA3-8B、LLaMA2-7B和LLaMA2-13B的最大吞吐量分别提高了19.5%、32.3%和40.1%。这些结果验证了LISA在压缩和加速大语言模型方面的有效性。

🎯 应用场景

LISA可应用于各种预训练大语言模型的压缩和加速，尤其是在资源受限的场景下，如移动设备或边缘计算。通过减少计算冗余，LISA能够降低模型的部署成本，提高推理速度，并促进大语言模型在更广泛的应用场景中的普及。

📄 摘要（原文）

To enhance the efficiency of the attention mechanism within large language models (LLMs), previous works primarily compress the KV cache or group attention heads, while largely overlooking redundancy between layers. Our comprehensive analyses across various LLMs show that highly similar attention patterns persist within most layers. It's intuitive to reduce the redundancy by sharing attention weights across layers. However, further analysis reveals two challenges: (1) Directly sharing the weight matrix without carefully rearranging the attention heads proves to be ineffective; (2) Shallow layers are vulnerable to small deviations in attention weights. Driven by these insights, we introduce LISA, a lightweight substitute for self-attention in well-trained LLMs. LISA employs tiny feed-forward networks to align attention heads between adjacent layers and low-rank matrices to approximate differences in layer-wise attention weights. Evaluations encompassing 13 typical benchmarks demonstrate that LISA maintains high response quality in terms of accuracy and perplexity while reducing redundant attention calculations within 53%-84% of the total layers. Our implementations of LISA achieve a 6x compression of Q and K matrices within the attention mechanism, with maximum throughput improvements 19.5%, 32.3%, and 40.1% for LLaMA3-8B, LLaMA2-7B, and LLaMA2-13B, respectively.