先日から、Multi-Head Attention とその関連技術について書いています。
本記事では、Grouped-Query Attention (GQA)1 について個人的な学習として調べてみます。
Grouped-Query Attention
原論文では、以下のように説明されています。
Grouped-query attention divides query heads into G groups, each of which shares a single key head and value head. GQA-G refers to grouped-query with G groups. GQA-1, with a single group and therefore single key and value head, is equivalent to MQA,while GQA-H,withgroupsequaltonumber of heads, is equivalent to MHA.
Grouped-Query Attention（GQA）は、Multi-Head Attention（MHA）と Multi-Query Attention（MQA）の中間的な設計として提案された Attention 機構です。
通常の Multi-Head Attention では、各 Query Head が独立した Key Head と Value Head を持ちます。一方、GQA では Query Head を複数のグループに分割し、各グループ内で1組の Key Head と Value Head を共有します。
グループ数が1つである場合は Multi-Query Attention (MQA) と等価であり、グループ数が Head 数と等しい場合は Multi-Head Attention と等価です。
たとえば、32個の Head を8つのグループに分ける場合は、以下のようになります。
Head0 - 3 K V
Head4 - 7 K V
...
Head28 - 31 K V
↓
8個のKV Cache
この場合、Key Head と Value Head は 32組から8組へ削減されます。その結果、KV Cache サイズは Multi-Head Attention の約 $\frac{1}{4}$ になります。
GQA は、MQA ほど極端に Key と Value を共有しないためモデル精度を維持しやすく、一方で MHA よりも KV Cache サイズや GPU メモリ帯域の使用量を削減できます。このバランスの良さから、現在では多くの LLM で採用されているようです。