【NLP】MHA、MQA、GQA机制的区别

Note

• LLama2的注意力机制使用了GQA。三种机制的图如下：
  

MHA机制（Multi-head Attention）

MHA（Multi-head Attention）是标准的多头注意力机制，包含h个Query、Key 和 Value 矩阵。所有注意力头的 Key 和 Value 矩阵权重不共享

MQA机制（Multi-Query Attention）

MQA（Multi-Query Attention，Fast Transformer Decoding: One Write-Head is All You Need）是多查询注意力的一种变体，也是用于自回归解码的一种注意力机制。与MHA不同的，MQA 让所有的头之间共享同一份 Key 和 Value 矩阵，每个头只单独保留了一份 Query 参数，从而大大减少 Key 和 Value 矩阵的参数量。

GQA机制（Grouped-Query Attention）

GQA（Grouped-Query Attention，GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints）是分组查询注意力，GQA将查询头分成G组，每个组共享一个Key 和 Value 矩阵。GQA-G是指具有G组的grouped-query attention。GQA-1具有单个组，因此具有单个Key 和 Value，等效于MQA。若GQA-H具有与头数相等的组，则其等效于MHA。GQA介于MHA和MQA之间。GQA机制，多头共用 KV Cache。

Reference

[1] 一文通透各种注意力：从多头注意力MHA到分组查询注意力GQA、多查询注意力MQA
[2] Transformer系列：注意力机制的优化，MQA和GQA原理简述
[3] Navigating the Attention Landscape: MHA, MQA, and GQA Decoded