09 · Grouped Query Attention (GQA)

注意力 Ainslie et al. 2023
Ainslie et al., Google · arXiv:2305.13245

迭代了什么

上一代的问题:MQA 省 KV 缓存但质量下降,MHA 质量好但推理贵。找不到中间方案。

这代改了什么:GQA:h 个 Query 头分成 G 组,每组共享一组 KV。GQA-1=MQA,GQA-H=MHA。

效果:Llama 2/3、Mistral、Gemma 2 全部使用 GQA。平均质量 47.1 接近 MHA 47.2,远好于 MQA 46.6。

一、解决的问题

MQA 推理快但质量下降,MHA 质量好但速度慢。

二、核心思想

将 h 个 Query 头分成 G 组,每组共享一组 KV 头。

三、Uptraining

通过 mean-pooling 将 MHA 检查点的 h 组 KV 头合并为 G 组,用原始 5% 计算量继续训练。

四、实际使用

Llama 2 70B, Llama 3, Mistral, Gemma 2 都使用 GQA。