分组查询注意力 GQA 是什么？让多个注意力头共享键值、节省显存

分组查询注意力（Grouped-Query Attention，GQA）是 Transformer 注意力机制的一种改进方式，由谷歌研究者 Joshua Ainslie 等人在 2023 年的论文中提出。它把注意力头分成若干组，同一组内的多个查询头共享同一套键和值，从而大幅减小推理时 KV 缓存的大小，提升大模型的推理速度，同时基本保持模型的效果。如今它被 Llama 2 的 700 亿参数版本、Llama 3、Mistral、Qwen 等大量主流模型采用。

从多头注意力说起

标准的 Transformer 使用多头注意力（MHA）：有多个注意力头，比如 32 个，每个头都有自己独立的一套查询、键和值，可以从不同的角度捕捉词与词之间的关系。在大模型推理时，为了避免重复计算，需要把每个 token 在每一层、每个头上的键和值都保存下来，这就是 KV 缓存。头越多、上下文越长、同时处理的请求越多，KV 缓存就越大，常常占据了大量的显存，并且每生成一个 token 都要从显存中读取这些数据，成为推理速度的主要瓶颈。

多查询注意力：省得多，但有损失

2019 年，Noam Shazeer 提出了多查询注意力（MQA）：所有的查询头共享同一套键和值，只保留一份。这样 KV 缓存一下子缩小为原来的几十分之一，推理速度大幅提升。但代价是模型的效果会有一定的下降，训练也可能变得不太稳定，因为所有的头都只能看到同样的键和值，表达能力受到了限制。

GQA：取一个折中

分组查询注意力在两者之间取了一个折中：把查询头分成若干组，每组共享一套键和值。比如 32 个查询头分成 8 组，每 4 个查询头共享一套键和值，这样只需要 8 套键值，KV 缓存缩小为原来的四分之一。当组数等于头数时，它就变成了标准的多头注意力；当只有一组时，它就变成了多查询注意力。研究表明，选择合适的组数，GQA 的效果可以接近多头注意力，而推理速度接近多查询注意力。论文还提出，已经用多头注意力训练好的模型，可以通过把各组内的键值头取平均，再用少量的计算继续训练，转换成 GQA 模型，不需要从头训练。

后续的发展

减小 KV 缓存一直是大模型推理优化的重点。在 GQA 之后，DeepSeek 在 2024 年的 DeepSeek-V2 中提出了多头潜在注意力（MLA），它不是简单地共享键值，而是把键和值压缩成一个低维的潜在向量进行缓存，在进一步大幅减小缓存的同时，保持甚至提升了模型的效果。