通过对齐注意力头实现高效模型压缩，显著减少显存占用。
Align Attention Heads Before Merging Them: An Effective Way for Converting MHA to GQA
- 用普鲁斯特分析对齐注意力头，提升头间相似性。
- 可压缩LLaMA2-7B 87.5% KV头，1.3B模型75%头。
- 适合需要轻量化部署的LLM应用开发者。
大型语言模型在多种自然语言处理任务中表现优异，但随着模型规模和输入序列增长，键值缓存线性增加，严重降低推理吞吐。为此，分组查询注意力（GQA）作为多头注意力（MHA）的替代方案被广泛引入。本文提出一种低成本方法，可将任意压缩比下的MHA转换为GQA。核心在于使用普鲁斯特分析对齐注意力头，增强头间相似性并保持计算不变性，从而提升微调后性能。随后采用$\mathit{L_0}$正则化剪枝冗余参数，使模型适配标准GQA框架。实验表明，该方法可对LLaMA2-7B模型压缩87.5%的KV头，对Sheared-LLaMA-1.3B模型压缩75%的KV头，且性能下降可接受。代码已开源：https://github.com/fpcsong/mha2gqa。
原文摘要 · Abstract (English)
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.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。