分类： AI

引言

在自然语言处理领域，如何在不牺牲模型质量的前提下实现更快的推理速度一直是一个重要的研究课题。本文将介绍一篇名为《GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints》的研究论文，该论文提出了一种从多头检查点训练广义多查询变换模型的方法，旨在解决这一挑战。

论文贡献

1. 现有多头语言模型检查点的再训练

论文提出了一种再训练现有多头语言模型检查点的方法，将其转换为具有多查询注意力（MQA）的模型。MQA使用单个键值头，大幅加快解码器推理速度。通过仅使用原始预训练计算量的5%，再训练现有模型，作者旨在实现更快的推理速度，而无需单独训练一个新模型。

2. 引入分组查询注意力（GQA）

论文引入了GQA，作为MQA的一种广义形式。GQA使用中间数量的键值头，数量介于1和查询头总数之间。该方法旨在平衡MQA的速度和多头注意力的质量。作者通过实验表明，再训练的GQA模型在保持与MQA相当速度的同时，能够达到接近多头注意力的质量。

方法与实验结果

再训练方法

论文详细介绍了再训练现有多头语言模型检查点的具体步骤。通过仅使用原始预训练计算量的5%，再训练现有模型，使其具备MQA的特点，从而实现更快的推理速度。

GQA的实现

GQA通过使用中间数量的键值头，增加了模型的灵活性和适应性。实验结果表明，再训练的GQA模型在多个自然语言处理任务中表现出色，质量接近于多头注意力，同时推理速度与MQA相当。

结论

本文提出的方法为实现更快的推理速度提供了一个有效的解决方案，而无需牺牲模型质量。通过再训练现有多头语言模型检查点，并引入分组查询注意力（GQA），可以在保持高质量的同时，实现高效的推理。这一研究对于提升变换模型在各类自然语言处理任务中的效率和性能具有重要意义。

---

进一步了解：

1. arXiv:2305.13245 – GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints
2. Papers With Code – GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints
3. ACL Anthology – GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints

引言

最近，幻方发布的DeepSeek-V2引起了广泛关注。其1块钱100万token的价格令人惊叹，而背后的关键技术之一——MLA（Multi-head Latent Attention）更是备受瞩目。本文将带大家梳理从MHA、MQA、GQA到MLA的演变历程，并深入介绍MLA的设计思路。

MHA：多头注意力

MHA（Multi-Head Attention）是《Attention is all you need》提出的注意力机制的基础。它通过多个独立的单头注意力拼接而成，广泛应用于当前的主流LLM（大语言模型）。

MHA的设计使得每个注意力头（Head）都有独立的键（Key）、值（Value）和查询（Query）向量，这些向量通过线性变换得到。MHA的计算量和存储开销较大，特别是在长上下文（Context）情况下，KV Cache（键值缓存）会占用大量显存。

瓶颈：为何降低KV Cache大小如此重要？

LLM的推理主要在GPU上进行，而GPU显存有限。一部分显存用于存放模型参数和激活值，另一部分用于存放KV Cache。随着上下文长度增加，KV Cache的大小会逐渐占据主导地位，可能超出单张卡甚至单台机器的显存容量。

减少KV Cache的目的是在更少的设备上推理更长的上下文，或在相同上下文长度下提高批处理大小，从而实现更快的推理速度或更大的吞吐量，最终降低推理成本。

MQA：多查询注意力

MQA（Multi-Query Attention）是减少KV Cache的一次尝试，首次提出于《Fast Transformer Decoding: One Write-Head is All You Need》。MQA的思路是让所有注意力头共享同一个Key和Value，从而将KV Cache减少到原来的1/h（h是头的数量）。

尽管MQA在显存占用上有显著优势，但其效果在某些任务上可能有所下降。为了弥补这一损失，研究人员提出了GQA。

GQA：分组查询注意力

GQA（Grouped-Query Attention）是MHA与MQA之间的过渡版本，出自《GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints》。GQA将所有注意力头分为g个组，每组共享同一对Key、Value。这样既可以减少KV Cache，又能在一定程度上保留多样性。

GQA提供了从MHA到MQA的自然过渡，当g=h时是MHA，g=1时是MQA，1 < g < h时则是GQA。GQA在KV Cache压缩率和效果之间提供了平衡。

MLA：多头潜在注意力

Part 1：增强模型能力

MLA（Multi-head Latent Attention）对GQA进行了改进，采用低秩投影的方式替代GQA的分割、重复。MLA通过不同的投影矩阵增强模型能力，并在推理阶段通过恒等变换技巧减少KV Cache。

Part 2：兼容RoPE

MLA的一个难题是如何兼容RoPE（旋转位置编码）。RoPE是一个位置相关的矩阵，MLA通过引入一种混合方法，每个注意力头的Query和Key新增部分维度用于添加RoPE，从而保持KV Cache的减少效果。

Part 3：减少训练参数量

MLA的最终版本将Query输入改为低秩投影形式，减少训练期间的参数量和相应的显存占用。推理阶段，通过恒等变换技巧减少KV Cache，同时保持高效的计算。

小结

本文概述了多头注意力的演变历程，特别是从MHA到MQA、GQA，最终到MLA的变化理念。MLA通过低秩投影和恒等变换技巧实现了KV Cache的进一步压缩，同时兼容RoPE，称得上是一种非常实用的注意力变体。

转载本文请包括本文地址：https://kexue.fm/archives/10091

如果您还有什么疑惑或建议，欢迎在下方评论区继续讨论。如果您觉得本文不错，欢迎分享或打赏本文。打赏并非为了获得收益，而是希望知道科学空间获得了多少读者的真心关注。再次表示欢迎和感谢！

---

苏剑林. (May. 13, 2024). 《缓存与效果的极限拉扯：从MHA、MQA、GQA到MLA》[Blog post]. Retrieved from 科学空间

---

MLA（Multi-head Latent Attention）是一种对GQA（Generalized Query Attention）进行改进的注意力机制。它采用低秩投影的方式替代了GQA中的分割和重复操作，同时通过恒等变换技巧减少了KV Cache的使用。MLA的核心思想是通过不同的投影矩阵增强模型的能力，并在推理阶段通过恒等变换技巧减少KV Cache的存储和计算开销。

MLA的改进主要解决了推理过程中的KV Cache问题，从而实现在更少的设备上推理更长的上下文，或者在相同的上下文长度下增大批处理大小，以实现更快的推理速度或更大的吞吐量，从而降低推理成本。

与经典的MHA（Multi-head Attention）和GQA、MQA（Multi-query Attention）相比，MLA在优化KV Cache和保证模型效果方面具有显著的优势。MLA通过低秩投影的方式替代了GQA中的分割和重复操作，从而大大减小了KV Cache的大小。与MQA相比，MLA的性能和效果显著优于MQA，甚至强于MHA和GQA，真正实现了降低推理成本并保证模型性能的目标。

MLA的核心是权重矩阵的合并。在传统的MHA中，注意力计算涉及到多个投影矩阵的乘法运算，而MLA通过合并这些投影矩阵，减少了存储和计算的开销。具体来说，MLA将(Q)^(T)K的计算结果合并为一个矩阵，并将合并后的权重应用到输入上，从而减少了存储和计算的开销。

然而，尽管MHA也可以进行合并，但由于其特定的计算方式，无法像MLA那样进行合并。MLA通过恒等变换技巧，将合并后的权重矩阵应用到输入上，从而实现了KV Cache的减少。

综上所述，MLA通过低秩投影替代分割和重复操作，采用恒等变换技巧减少KV Cache的使用，从而在推理过程中降低了存储和计算的开销，实现了更高效的推理和更低的成本。

---

Learn more:

1. 还在用MHA？MLA来了DeepSeek-v2的MLA的总结和思考 – 知乎
2. [2405.04434] DeepSeek-V2: A Strong, Economical, and Efficient Mixture-of-Experts Language Model
3. DeepSeek-V2中的MLA详解 – 知差（chai）