Align Attention Heads Before Merging Them: An Effective Way for Converting MHA to GQA

📄 arXiv: 2412.20677v2 📥 PDF

作者: Qingyun Jin, Xiaohui Song, Feng Zhou, Zengchang Qin

分类: cs.CL

发布日期: 2024-12-30 (更新: 2025-07-26)

备注: 13 pages, 3 figures

🔗 代码/项目: GITHUB

💡 一句话要点

提出一种高效方法将多头注意力转换为分组查询注意力

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

关键词: 大型语言模型 多头注意力 分组查询注意力 Procrustes分析 参数修剪 推理效率 自然语言处理

📋 核心要点

  1. 现有的多头注意力机制在处理大规模输入时，KV缓存的线性增长导致推理效率显著下降。
  2. 本文提出通过Procrustes分析对注意力头进行对齐，从而提高其相似性并保持计算不变性，进而实现MHA到GQA的转换。
  3. 实验结果显示，所提方法在LLaMA2-7B和Sheared-LLaMA-1.3B模型上分别实现了87.5%和75%的KV头压缩，且性能下降在可接受范围内。

📝 摘要（中文）

大型语言模型（LLMs）在多种自然语言处理任务中表现出色。然而，随着模型规模和输入序列长度的增加，线性增长的键值（KV）缓存显著降低了推理吞吐量。因此，分组查询注意力（GQA）作为多头注意力（MHA）的替代方案被广泛引入LLMs。本文提出了一种经济有效的方法，将MHA转换为任意压缩比的KV头。我们的方法的关键在于对注意力头应用Procrustes分析，以增强注意力头之间的相似性，同时保持计算不变性，从而提高模型的后训练性能。随后，我们采用$ extit{L_0}$正则化来修剪冗余参数。实验结果表明，我们的策略能够在可接受的性能下降下，压缩LLaMA2-7B模型的KV头达87.5%，Sheared-LLaMA-1.3B模型的KV头达75%。

🔬 方法详解

问题定义：本文旨在解决大型语言模型在推理过程中，由于KV缓存线性增长导致的效率低下问题。现有的多头注意力机制在处理长序列时，推理吞吐量显著下降，影响了模型的实际应用。

核心思路：我们提出了一种通过Procrustes分析对注意力头进行对齐的方法，以增强注意力头之间的相似性，同时保持计算不变性。这种设计旨在有效地将多头注意力转换为分组查询注意力，降低计算复杂度。

技术框架：整体架构包括对注意力头的对齐、应用L0正则化进行参数修剪，以及将修剪后的模型适配到标准的GQA框架。主要模块包括Procrustes分析模块和参数修剪模块。

关键创新：最重要的技术创新在于将Procrustes分析引入到注意力头的对齐过程中，这一方法显著提高了注意力头的相似性，并有效降低了计算成本，与传统的MHA方法相比，具有更高的效率和灵活性。

关键设计：在参数设置上，我们采用了L0正则化来修剪冗余的KV头，确保模型在压缩后仍能保持良好的性能。网络结构方面，注意力头的对齐过程是通过计算相似性矩阵实现的，确保了模型的计算不变性。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

实验结果表明，所提方法在LLaMA2-7B模型上实现了高达87.5%的KV头压缩，而在Sheared-LLaMA-1.3B模型上则达到了75%的压缩率，且性能下降在可接受范围内。这一结果显示了方法在保持模型性能的同时，显著提高了推理效率。

🎯 应用场景

该研究的潜在应用领域包括自然语言处理中的大型语言模型优化、实时推理系统以及资源受限环境下的模型部署。通过提高推理效率，能够更好地满足实际应用需求，推动智能助手、对话系统等技术的发展。

📄 摘要（原文）

Large language models (LLMs) have demonstrated exceptional performance across diverse natural language processing tasks. However, as the model size and the input sequence's length increase, the linearly increasing key-value (KV) cache significantly degrades inference throughput. Therefore, grouped-query attention (GQA), as an alternative to multi-head attention (MHA), has been widely introduced into LLMs. In this work, we propose a cost-effective method for converting MHA into GQA with any compression ratio of KV heads. The key point of our method lies in the application of Procrustes analysis to the attention heads, which enhances the similarity among attention heads while preserving computational invariance, thereby improving the model's post-training performance. Subsequently, we employ $\mathit{L_0}$ regularization to prune redundant parameters. The model after pruning can be adapted to the standard GQA framework. Experimental results show that our strategy can compress up to 87.5\% KV heads of LLaMA2-7B model and 75\% KV heads of Sheared-LLaMA-1.3B with acceptable performance degradation. Our code is released at https://github.com/fpcsong/mha2gqa.