• 1 MQA与GQA 提出动机

  • 1.1 MHA 提出动机

  • 1.2 MQA 提出动机

  • 1.3 GQA 提出动机

  • 1.3.1 MQA 既要快还要好

  • 1.3.2 比 MQA 更好的 GQA

  • 1.3.3 GQA 思想

• 2 MQA与GQA 原理

  • 2.1 MHA 原理

  • 2.2 MQA 原理

  • 2.3 MQA 增量训练原理

  • 2.4 GQA 原理

  • 2.5 GQA 增量训练原理

• 3 MQA与GQA 从零实现

  • 3.1 MQA 从零实现

  • 3.2 MQA 增量训练实现

  • 3.2.1 定义网络结构

  • 3.2.2 转换模型参数

  • 3.2.3 增量训练

  • 3.3 GQA 从零实现

  • 3.3.1 GQA 实现策略

  • 3.3.2 从零实现 GQA

  • 3.4 GQA 增量训练实现

  • 3.4.1 参数转换实现

  • 3.4.2 定义网络结构

  • 3.4.3 增量训练

• 4 复杂度分析

  • 4.1 复杂度分析对比

  • 4.2 参数缓存对比

• 5 总结

• 引用

在前面的一系列文章中，掌柜给大家介绍了多种以降低模型计算资源为动机而提出来的算法模型，可以看出这也成为了大模型研究领域的一个重要方向。

例如到目前为止，我们已经遇到过了多个由类似动机所提出的模型技术：① 希望减少模型参数，但又不太损失精度——模型蒸馏；② 希望加快推理速度，又能保持模型精度—— RMSNorm；③希望减少模型重复计算过程，降低缓存使用同时提高推理效率——KV Cache；④ 希望不仅仅只通过堆资源，而是充分利用训练数据来提高模型性能——LLaMA 1。

MHA、MQA 和 GQA 结构对比图

在接下来的这篇文章当中，掌柜将继续一口气给大家介绍由同一家公司（谷歌）中的不同部门（Google Research 和 Google Brain）分别提出来的两个基于多头注意力机制改造而来的算法（MQA 和 GQA），其中一个作者还是 Transformer 论文的二作。

读完本文你将对大模型中的 MHA 、MQA 和 GQA  的原理和实现过程有着充分地认识和了解，你会发现 MQA 和 GQA 的思路其实就是归一化中层归一化 LN 和组归一化 GN 的想法。当然更重要的是还将教会你如何把现有的 MHA 参数无缝迁移到 MQA 和 GQA 中，以及基于此进行增量训练。

关键词：多头注意力、self-attention、MHA、MQA、GQA、LLaMA、Transformer、大模型

1 MQA与GQA 提出动机

在前面介绍 Transformer 的时候掌柜就详细介绍了多头注意力机制（Multi-Head Attention, MHA）的原理及实现过程，简单来说多头注意力机制包含两个关键元素：自注意力机制和多头。关于自注意力机制的动机原理掌柜这里就不再赘述可以参见掌柜这篇文章

下面，我们稍微来回顾一下提出“多头”的动机是什么。

1.1 MHA 提出动机

30 秒时间，你也思考一下为什么要使用多头？

3

2

1

自注意力机制在对当前位置的信息进行编码时会将注意力过度集中于自身所在的位置，因为在计算注意力权重时自己与自己所在位置的相似性是最高的，而这可能导致模型忽略了其它位置上的信息。

如果大家还有印象的话就知道，上面这段话便是掌柜在介绍“多头”动机时的描述。

图 1. 多头注意力权重可视化结果[1]

如图1所示是一张在真实场景下多头注意力中不同权重矩阵的可视化结果。可以发现，在不同的权重矩阵中注意力值的分配并不相同，并且可以明显地看出在第3张注意力权重可视化结果中，模型将过多的注意力集中到了每个字符自身所在的位置而忽略了其它位置上的信息。

因此，在一定上相同维度中多头  的值越大，整个模型的表达能力就越强，越能使得模型对于注意力权重进行合理分配。

此时，我们终于想起原来使用“多头”的目的是为了解决模型将注意力过度集中于自身位置的问题，即让模型能拥有更多的注意力权重分配方式。

敲黑板，上面加粗是重点，等下要用到！

不过，所谓成也萧何败萧何，多头既让模型增强了表达能力，可同时也显著增加了模型的参数量，当然更重要的一点是增加了推理时 KV 缓存的开销。

所以，在这样的背景下，如何才能做到既可以满足多头的动机，同时又能最大程度的降低所使用的计算资源呢？

1.2 MQA 提出动机

自多头注意力提出以来，它的优越性能就使得其成为了替代 RNN 的不二选择。尽管得益于自注意力机制在序列长度上的并行化计算，模型在训练时的速度通常要远远快于同量级的 RNN 网络；但是，由于在推理过程中生成内容需要逐时刻进行，这就意味着无法实现并行，此时模型需要反复加载大量的的 “keys” 和 “values” 状态，使得内存带宽开销很大 [2]。

While training these layers is generally fast and simple, due to parallelizability across the length of the sequence, incremental inference (where such paralleization is impossible) is often slow, due to the memory-bandwidth cost of repeatedly loading the large "keys" and "values" tensors.

图 2. MHQ 注意力 Key-Value 对示意图

从图2可以看出，在多头注意力当中每个头都有对应的 Key、Value 和 Query，然而在推理过程中，生成每一个新词模型都必须同时使用之前时刻所有的 Key 和 Value，而这些缓存（KV缓存）通常很大且不能并行，所以这就是自注意力机制的瓶颈所在。

更多关于 KV 缓存的内容，可以参见文章 「大模型为什么需要KV缓存是何动机与原理？为什么Q不用缓存？」

在这样的动机下，2019年11月 Noam Shazeer  提出了变种的多头注意力机制——多查询注意力机制（Multi Query Attention, MQA） ，即同一个 Key 和 Value 在不同的头之间共享，这样便显著地降低了在推理过程中对于计算资源的开销 [2]。

We propose a variant called multi-query attention, where the keys and values are shared across all of the different attention "heads", greatly reducing the size of these tensors and hence the memory bandwidth requirements of incremental decoding.

图 3. MQA 注意力 Key-Value 对示意图

从图3我们可以看出，此时虽然只有一个 Key 和 Value ，但是多个 Query 与同一个 Key 一样可以计算得到多个不同的注意力权重分配矩阵，然后再与同一个 Value 作用得到最后多头的计算结果。

你看，这样是不是既满足了多头的动机，又达到了降低模型参数的目的？

那这样行不行呢？感觉上好像行，但是也说不准。

怎么办？实验试一试，如果效果差不多那就证明这样是可以的，如过差很远那就证明这样的做法是不行的。

其实深度学习大多数情况下几乎都是这样。

最后，经过作者实验后发现，通过 MQA 训练得到的模型在解码时的速度将会明显变快，并且与基线模型相比结果只有些许的下降。

We verify exper imentally that the resulting models can indeed be much faster to decode, and incur only minor quality degradation from the baseline.

不过这里需要注意的是，上面的结论仅仅只是在作者自己论文中实验的结论，但是实际情况是 MQA 在大模型推理中的效果并不太理想，因为单一的 Key 和 Value 大大降低了模型的表达能力，因此用得更多的是分组注意力查询（Grouped Query Attention, GQA ）[3]，它是 MHA 和 MQA 之间的一种折衷做法。

不过，这依旧不影响我们去学习和了解 MQA ，因为它也是 GQA 的一种特殊情况。

这里顺便提一下 Noam Shazeer 这个人，他是 Transformer 的作者之一、然后自己提出了 MQA、接着又提出了之前掌柜已经介绍过的 SwiGLU，后来又参与了 Pathways Language Model (PaLM) 等等。

1.3 GQA 提出动机

尽管 MQA 在推理时的速度相比于 MHA 更快，但是从实际情况来看它的效果并不好；其次，从各个技术的出现时间上来看，MHA 在前而GQA 在后，因此还需要考虑的一个点便是如何将 MHA 方便地转换为 GQA 能够适用的情况。

下面，掌柜就分别从这两个方面以 GQA 论文中的角度来为大家介绍 GQA 提出的动机。

1.3.1 MQA 既要快还要好

尽管 MQA 在推理过程中能够极大地加快模型的推理速度，但是由于所有的注意力头都共享这一个 Key 和 Value ，使得模型的效果有了很大程度的降低，同时还影响了模型训练时的稳定性。因此，如果仅仅只是为了使用 MQA 来加快推理速度而再从头单独训练这样一个独立的模型显然是不可取的。

Multi-query attention (MQA), which only uses a single key-value head, drastically speeds up decoder inference. However, MQA can lead to quality degradation  training instability, and moreover it may not be desirable to train a separate model just for faster inference.

什么意思呢，为什么说从头单独训练一个模型又是不可取的呢？

为了说明这个问题，我们先来看看这篇论文提出的背景及时间点。

GQA 是 2023年12月提出的，此时距离 GPT-4 提出已经过去了9个月，距离 1750 亿参数量的 GPT-3 （2020年5月）已经过去了快2年半。作为对比，MQA 提出的时间是 2019年11月，彼此模型的参数规模大约在数十亿这样的量级，例如 2019年2月发布的 GPT-2 参数规模约 15 亿，也就相当于15个BERT-base 左右的规模。

图 4. 模型参数规模发展图 [5]

所以，从千亿规模参数的时代再回过头去看数十亿规模的模型，简直是不要太迷你，说从头开始训练就可以从头开始训练，一点不带犹豫。

然而，现在的大模型动不动就是上百亿千亿的规模，模型已经（采用 MHA ）训练好了，那我们总不能为了使用 MQA 再来重新训练一个模型，这显然是不可取的。

反观 MQA 时代，那时候的模型都还很小，所以可以随时从头训练。

所以，上面所谓的 “从头单独训练这样一个独立的模型” 作者想要表达的正真含义是：先前使用 MQA 的方式都不对，都是基于 MQA 从头开始训练模型的，这样虽然加快了模型的推理速度但是又耗费了资源（重新训练）同时模型效果又不好（共享一个 Key 和 Value）。

因此，在 GQA 这篇论文中，作者的第一个贡献就是提出了通过以增量训练（uptrain）的方式，将 MHA 中的参数复用到 MQA 中，这样既保证了模型推理时的速度，又使得模型的效果比重新从头训练一个 MQA 模型的效果更好。

This work contains two contributions for faster inference with large language models. First, we show that language model checkpoints with multi-head attention (MHA) can be uptrained to use MQA with a small fraction of original training compute.  This presents a cost-effective method to obtain fast multi-query as well as high-quality MHA checkpoints.

1.3.2 比 MQA 更好的 GQA

虽然通过增量训练的方式将预训练得到的 MHA 转换成 MQA 比单纯直接使用 MQA 从头训练得到的效果更好，但是这依旧没能改变一个事实—— MQA 中多头注意力共享一个 Key 和 Value，而这也极大地降低了模型的表达能力，因此，作者又提出了分组查询注意力机制（ Grouped Query Attention, GQA）这一方法 [3]。

最后，经过增量训练以后，GQA 既能在速度上与 MQA 类似，在效果上也与 MHA 几乎接近。

Second, we propose grouped-query attention (GQA), an interpolation between multi-head and multi-query attention with single key and value heads per subgroup of query heads. We show that uptrained GQA achieves quality close to multi-head attention while being almost as fast as multi-query attention.

总结起来这篇论文一共两点贡献：

① 给出了如何通过增量训练将 MHA 中的模型参数迁移到 MQA 中；

② 提出了比 MQA 更优的 GQA 算法，并同样也给出了将 MHA 复用到 GQA 的方法，并对 GQA 的性能进行了分析。

1.3.3 GQA 思想

那什么又是 GQA 呢？

简单来说它是介于 MQA 和 MHA 之间的一种解决方案，算是 MQA 和 MHA 之间的一种折衷处理，即将 Query 进行分组，然后每一组 Query 共享一个 Key 和 Value。可以看出 MQA 就是 GQA 的一种特殊情况，如图5所示。

图 5. MHA、MQA 和 GQA 结构对比图

在图5右侧 GQA 当中，此时的 Queries 被分成了 3 组，因此对应就有 3 组 Keys-Values ，然后每一组内部所有的 Query 共享一个 Key 和 Value。 此时可以看出，MQA 其实就是 GQA 的一种特殊情况，也就是把所有的 Query 看成一组。

看到这里大家脑海有想到了什么吗？

是不是觉得 MQA 和 GQA 的做法有点类似于掌柜之前在深度学习基础中所介绍的两种算法？

想一想……

对，就是 层归一化和组归一化这两种算法，如图6所示。

图 6. LN 和 GN 归一化对比图

在 LN 和 GN 中，LN 是将所有的 feature map 看成是一个整体进行归一化，而 GN 则是分组进行归一化，每一组内的 feature map 共享一组参数，可以看出 LN 也是 GN 的一种特殊情况，即此时每个通道为一个组。

2 MQA与GQA 原理

在介绍完 MQA 和 GQA 各自提出的动机以后，我们再来看一下各自具体的原理部分。为了能够循序渐进地来介绍这两部分内容，掌柜先来带着大家简单回顾一下 MHA 的原理。

2.1 MHA 原理

假设多头输入 Query、Key 和 Value 分别为  、  和   ，其中  和   分别表示源序列和目标序列的长度， 表示模型维度；第  个头对应的线性层映射权重为  、  和  ，其中 ，且通常  。

此时，第  个头的注意力计算过程可以表示为

进一步，多头注意力的计算过程为

其中  。

接着，根据式(1)我们可以得到计算 3 个线性变换的时间复杂度分别为  、  和  ， 对应的时间复杂度为   ，计算注意力权重的时间复杂度为  。因为模型超参数是固定的，MHA 的整体时间复杂度为 。同时，由于 MHA 中每个头拥有独立的 Query 、Key 和 Value ，所以缓存  和  的空间复杂度均为 。

此时，如果仅考虑推理时的场景那么 ，则 MHA 整体的时间复杂度为  ，空间复杂度为  。

以上整个计算过程便是 MHA 的原理，下面我们再来看 MQA 的原理。

2.2 MQA 原理

在介绍 MQA 思想的时候我们谈到，MQA 中同一个 Key 和 Value 会在不同的头之间共享，所以仅有第  个头对应的 Query 有线性层映射权重为  ，其中  ；而 Key 和 Value 则分别只有一个  和  ；通常  。

此时，第  个头的注意力计算过程可以表示为

进一步，多头注意力的计算过程为

其中  。

……

 AI大模型从0到精通全套学习大礼包

我在一线互联网企业工作十余年里，指导过不少同行后辈。帮助很多人得到了学习和成长。

只要你是真心想学AI大模型，我这份资料就可以无偿共享给你学习。大模型行业确实也需要更多的有志之士加入进来，我也真心希望帮助大家学好这门技术，如果日后有什么学习上的问题，欢迎找我交流，有技术上面的问题，我是很愿意去帮助大家的！

如果你也想通过学大模型技术去帮助就业和转行，可以点扫描下方👇👇
大模型重磅福利：入门进阶全套104G学习资源包免费分享！

01.从入门到精通的全套视频教程

包含提示词工程、RAG、Agent等技术点

​

02.AI大模型学习路线图（还有视频解说）

全过程AI大模型学习路线

​​

03.学习电子书籍和技术文档

市面上的大模型书籍确实太多了，这些是我精选出来的

04.大模型面试题目详解

05.这些资料真的有用吗?

这份资料由我和鲁为民博士共同整理，鲁为民博士先后获得了北京清华大学学士和美国加州理工学院博士学位，在包括IEEE Transactions等学术期刊和诸多国际会议上发表了超过50篇学术论文、取得了多项美国和中国发明专利，同时还斩获了吴文俊人工智能科学技术奖。目前我正在和鲁博士共同进行人工智能的研究。

所有的视频由智泊AI老师录制，且资料与智泊AI共享，相互补充。这份学习大礼包应该算是现在最全面的大模型学习资料了。

资料内容涵盖了从入门到进阶的各类视频教程和实战项目，无论你是小白还是有些技术基础的，这份资料都绝对能帮助你提升薪资待遇，转行大模型岗位。

智泊AI始终秉持着“让每个人平等享受到优质教育资源”的育人理念‌，通过动态追踪大模型开发、数据标注伦理等前沿技术趋势‌，构建起"前沿课程+智能实训+精准就业"的高效培养体系。

课堂上不光教理论，还带着学员做了十多个真实项目。学员要亲自上手搞数据清洗、模型调优这些硬核操作，把课本知识变成真本事‌！

如果说你是以下人群中的其中一类，都可以来智泊AI学习人工智能，找到高薪工作，一次小小的“投资”换来的是终身受益！

应届毕业生‌：无工作经验但想要系统学习AI大模型技术，期待通过实战项目掌握核心技术。

零基础转型‌：非技术背景但关注AI应用场景，计划通过低代码工具实现“AI+行业”跨界‌。

业务赋能 ‌突破瓶颈：传统开发者（Java/前端等）学习Transformer架构与LangChain框架，向AI全栈工程师转型‌。

👉获取方式：
😝有需要的小伙伴，可以保存图片到wx扫描二v码免费领取【保证100%免费】🆓