在合并注意力头之前对注意力头进行对齐：将 MHA 转换为 GQA 有效方法

大型语言模型 (LLM) 在 diverse natural language processing 任务中展现了卓越的性能。然而，随着模型规模和输入序列长度的增加，线性增长的 key-value (KV) 缓存显著降低推理吞吐量。因此，作为 MHA (多头注意力) 的替代方案，grouped-query attention (GQA) 被广泛引入 LLM。本文提出一种将 MHA 转换为 GQA 的经济有效方法，支持任意 KV 头压缩比例。我们方法的关键在于应用 Procrustes 分析对注意力头进行处理，既保持计算不变性又增强注意力头之间的相似性，从而提高模型的后训练性能。随后，我们采用 $\mathit{L_0}$ 正则化来修剪冗余参数。经修剪后的模型可适配标准 GQA 框架。实验结果表明，我们的策略可将 LLaMA2-7B 模型的 KV 头压缩至 87.5%，Sheared-LLaMA-1.3B 模型压缩至 75%，同时保持可接受的性能损失。我们的代码已发布于 https://github.com/fpcsong/mha2gqa。