AI理论知识系统复习(3)：GQA（Grouped Query Attention）、MQA（Multi-Query Attention）以及与MHA的区别

目录1 传统MHA存在什么弊端2 多查询注意力 MQA2.1 什么是多查询注意力MQA2.2 MQA的问题3 分组查询注意力 GQA3.1 什么是分组查询注意力3.2 分组查询注意力的工作原理4 MHA GQA MQA他们三者之间的区别和联系4.1 GQA是MHA和MQA的推广4.2 GQA、MHA与MQA的比较4.3 将MHA转换为GQA5 总结参考文献abstract:MHA每个头各自一套 KV分组查询注意力GQA就是分组共享键值KV多查询注意力MQA就是所有头共享键值KV联系组数 头数 → 退化为 MHA组数 1 → 退化为 MQA1 传统MHA存在什么弊端传统的多头注意力由于每个头都需要缓存各自的 K V因此多头注意力中的KV缓存增长非常快占用大量内存。这使得推理变得缓慢且成本高昂尤其是对于长序列和大型模型。因此我们需要一种更智能的方法在不损失模型输出质量的前提下减少KV缓存的大小。首先让我们理解多查询注意力MQA这是解决该问题的首次尝试。然后我们将学习分组查询注意力GQA这是一个更好的解决方案。2多查询注意力 MQA2.1 什么是多查询注意力MQA多查询注意力MQA是一种策略所有head共享相同的键和值但每个head仍拥有自己的查询。让我们来分解一下这个术语多查询注意力 多查询 单一共享键和值在多头注意力中每个头都有自己的Q、K和V。在多查询注意力中每个头都有自己的Q但所有头共享一个K和一个V。假设我们有8个头多头注意力MHAhead1Q₁、K₁、V₁head2Q₂、K₂、V₂head3Q₃、K₃、V₃...8个独立的K和V套装多查询注意力MQAhead1Q₁、K_shared、V_sharedhead2Q₂K_sharedV_sharedhead3Q₃、K_shared、V_shared...所有head仅设置1K和1V现在KV缓存只需要存储1组密钥和1组值而不是8组。KV缓存体积会变小8倍。2.2MQA的问题MQA大幅减少内存但存在权衡。由于所有头共享相同的键和值模型失去了从不同角度观察输入的能力。输出质量可能会下降训练也可能变得不稳定。需要折中既省内存又尽量保持质量这就是 GQA3分组查询注意力 GQA3.1 什么是分组查询注意力分组查询注意力GQA是一种策略将头项划分为组组内所有头共享相同的键和值但每个头仍拥有自己的查询。分组查询注意力 分组查询 每组共享键和值简单来说GQA不是像MHA那样给每个头有自己的K和V也不是像MQA那样给所有头共用一个K和V而是将头分组每个组共享一个K和V。学习这点的最好方法是举个例子。假设我们有8个头然后把它们分成两组每组4个头第一组头1头2头3头4——共享K_group1V_group1第二组第五个头第六个头第七个头八个头——共享K_group2V_group2所以head1Q₁、K_group1、V_group1head2Q₂K_group1V_group1head3Q₃、K_group1、V_group1head4Q₄、K_group1、V_group1head5Q₅、K_group2、V_group2head6Q₆、K_group2、V_group2head7Q₇、K_group2、V_group2head8Q₈、K_group2、V_group2现在我们不再像MHA那样存储8套K和V而是只存储2套。KV缓存体积会变小4倍。而且我们不是像MQA那样只用一个共享的K和V而是用两个集合。所以模型仍然有一定能力从不同角度看待输入。GQA是MHA和MQA之间的最佳平衡点。它节省了接近MQA的内存同时保持了接近MHA的质量。注在GQA中只有键和值在组内共享。查询仍然是每个头的独立。这很重要因为查询让每个头都能从不同的角度看待输入。通过保持查询分开GQA保持了注意力模式的多样性。KV缓存缩小是因为我们只在推理时存储键和值而不是查询时。3.2分组查询注意力的工作原理让我们一步步走过整个流程。第一步把头分成几组。组数是我们在训练前选择的设定。假设我们有8个头和2组。第二步每个头用自己的权重矩阵计算自己的查询Q。所以这8个头都有各自独立的查询。这和MHA是一样的。第三步每个组使用该组的权重矩阵计算一个共享键K和一个共享值V。第一组计算K_group1和V_group1。第2组计算K_group2和V_group2。第四步每个头部使用自己的查询运行注意力机制但共享其组的键和值。第1到第4头使用K_group1和V_group1。第5到第8号用K_group2和V_group2。步骤5所有头部的输出都被串接并通过最终投影就像MHA一样。结果输出与MHA相同。但KV缓存要小得多因为我们只为每个组存储K和V而不是每个头。4 MHA GQA MQA他们三者之间的区别和联系4.1GQA是MHA和MQA的推广当组数正面数时每个群体恰好有一个头。每个头都有自己的K和V。这正是多头注意力MHA的体现。当组数 1 时所有头颅都在同一组。所有头部共用相同的K和V。这正是多查询注意力MQA。当组数介于1到头数之间时这就是分组查询注意力GQA。4.2GQA、MHA与MQA的比较MHA (8 query heads, 8 KV sets - one per head):[Q1] [Q2] [Q3] [Q4] [Q5] [Q6] [Q7] [Q8]| | | | | | | |v v v v v v v v[K1] [K2] [K3] [K4] [K5] [K6] [K7] [K8][V1] [V2] [V3] [V4] [V5] [V6] [V7] [V8]MQA (8 query heads, 1 KV set - shared by all heads):[Q1] [Q2] [Q3] [Q4] [Q5] [Q6] [Q7] [Q8]\ \ \ | | / / /-----------------------------|v[K_shared][V_shared]GQA (8 query heads, 2 groups - 1 KV set per group):[Q1] [Q2] [Q3] [Q4] [Q5] [Q6] [Q7] [Q8]\ | | / \ | | /-------- ------| |v v[K_group1] [K_group2][V_group1] [V_group2]在MHA中每个查询都有自己的私钥和私值。8个查询8个KV集。最高质量最大内存。在MQA中所有查询都指向一个共享的键和值。8个查询1KV组。内存最少但多样性降低。在GQA中查询被划分为多个组每个组共享一个键和值。8个查询2个KV组。两者兼得。4.3将MHA转换为GQA接下来一个大问题是我们是否总是需要从零开始训练GQA模型答案是否定的。最初的GQA论文表明我们可以将现有的多头注意力模型以非常低的成本转换为GQA模型。这叫做uptraining上训练/继续微调。流程很简单第一步以一个已经训练过的现有MHA模型为例。第二步对于每个组取该组中所有head的键权矩阵并取平均。对Value权重矩阵也同样操作。这样我们每个组就有一个共享的密钥和一个共享的值。步骤3对模型进行短时间微调。原始论文显示仅仅用约5%的预训练计算量进行上训练就足以恢复接近全MHA的质量。这也是GQA被迅速采用的主要原因之一。实验室无需丢弃现有的MHA模型也不必花费大量计算量来训练新模型。他们可以直接把他们升格为GQA。5 总结多头注意力MHA每个head都有自己的查询、键和值。质量最好但KV缓存在推理过程中会变得非常大。KV缓存问题在文本生成过程中模型会为每个head存储每个前单词的键和值。由于多head和长序列这需要大量GPU内存。多查询注意力MQA所有head共享一个键和一个值但每个头仍然有自己的查询。KV缓存会变得非常小但输出质量可能会下降。分组查询注意力GQAhead被划分为多个组。每个组共享一个键和一个值而每个头仍然拥有自己的查询。这是MHA和MQA之间的最佳平衡点。GQA是一个概括当组数等于头数时GQA变为MHA。当组数为1时GQA变为MQA。介于两者之间的都是GQA。重要性GQA节省的内存接近MQA同时保持接近MHA的质量。它现已被应用于许多流行型号如LLaMA 2、LLaMA 3和Mistral 7B。上级培训我们不需要从零开始训练GQA模型。我们可以拿现有的MHA模型平均每个组内的K和V权重然后短时间微调转化为GQA模型。这也是GQA被迅速采用的主要原因之一。参考文献Grouped Query Attention