GQA（Grouped Query Attention）：分组注意力机制的原理与实践《一》

GQA（Grouped Query Attention）是近年来在大语言模型中广泛应用的一种注意力机制优化方法，最初由 Google 在 2023 年提出。它是对 Multi-Query Attention (MQA) 的扩展，旨在平衡模型性能与计算效率。

---

🌟 GQA 是什么？

简单定义：

GQA 是一种将查询头（Query Heads）分组，并共享键（Key）和值（Value）头的注意力机制变体。

它试图在 标准的多头注意力（MHA） 和 多查询注意力（MQA） 之间找到一个折中点：

注意力类型	Query Heads	Key/Value Heads	共享情况
MHA	多个	多个	不共享
GQA	多个	少于 Query 的多个	分组共享
MQA	多个	1	完全共享

---

🧠 原理详解

1. 回顾标准 Multi-Head Attention (MHA)

在标准的 Transformer 中：

• 每个 token 的 Q、K、V 都是由输入线性变换得到。
• 如果有 H 个 attention heads，则每个 head 都有自己的 Q、K、V 向量。

公式如下：

$$Q=X{W}_Q,K=X{W}_K,V=X{W}_V$$

其中 $ W_Q, W_K, W_V $ 是可学习参数。

每个 head 的 Q/K/V 是从这些矩阵中切出来的。

---

2. 引入 GQA：Query 分组 + Key/Value 共享

在 GQA 中：

• Query heads 被分成若干组（比如 4 组）
• 每组共享一组 Key 和 Value 向量（即每组对应一个 K 和 V）

例如：

• 总共 32 个 query heads
• 分成 4 组，每组 8 个 heads
• 每组使用相同的 Key 和 Value 向量
• 所以只需要 4 组 K/V，而不是 32 组

这样做的好处是：

• 减少了 Key/Value 的数量，降低了内存占用（尤其是 KV Cache）
• 保留了比 MQA 更多的表达能力

---

⚙️ GQA 的优势

优势	描述
✅ 推理速度更快	更少的 Key/Value 向量意味着更小的 KV Cache，减少解码时的内存访问延迟
✅ 内存占用更低	特别是在批量生成或长文本生成时，KV Cache 占用显著降低
✅ 比 MQA 表现更好	相比完全共享 Key/Value 的 MQA，GQA 保留了部分多样性，模型表现通常更优
✅ 适合部署	对硬件资源友好，特别适合在有限算力设备上运行的大模型

---

🔍 示例说明（来自 Llama 3）

Llama 3 使用了 GQA 技术来提升推理效率。

• 总共 32 个 query heads
• 只使用了 8 个 key/value heads（即每组 4 个 queries 共享一个 key/value）

这意味着：

• 每个 group 包含 4 个 query heads
• 这些 query 共享同一个 key 和 value

这样可以保持大部分 MHA 的表达能力，同时节省内存和计算开销。

---

📈 MHA vs GQA vs MQA 性能对比（大致）

指标	MHA	GQA	MQA
表达能力	最强	中等	最弱
推理速度	较慢	快	最快
内存占用（KV Cache）	最高	中等	最低
部署友好度	一般	高	最高

---

🧩 应用场景

GQA 特别适用于以下场景：

• 大模型推理优化（如 Llama 3、PaLM 2、Gemini Nano）
• 移动端/边缘端部署
• 需要长上下文处理的任务
• 大批量生成任务

---

💡 总结

项目	GQA
类型	注意力机制变体
核心思想	Query 分组 + Key/Value 共享
优点	提升推理速度、降低内存消耗、兼顾模型表现
缺点	表达能力略低于 MHA
应用	大语言模型部署、高效推理系统

---