分組查詢注意力 (Grouped Query Attention) 是一種在大型語言模型中的多查詢注意力 (MQA) 和多頭注意力 (MHA) 之間進(jìn)行插值的方法，它的目標(biāo)是在保持 MQA 速度的同時(shí)實(shí)現(xiàn) MHA 的質(zhì)量。

這篇文章中，我們將解釋GQA的思想以及如何將其轉(zhuǎn)化為代碼。

GQA是在論文 GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints paper.中提出，這是一個(gè)相當(dāng)簡單和干凈的想法，并且建立在多頭注意力之上。

GQA

標(biāo)準(zhǔn)多頭注意層(MHA)由H個(gè)查詢頭、鍵頭和值頭組成。每個(gè)頭都有D個(gè)維度。Pytorch的代碼如下：

from torch.nn.functional import scaled_dot_product_attention
 
 # shapes: (batch_size, seq_len, num_heads, head_dim)
 query = torch.randn(1, 256, 8, 64)
 key = torch.randn(1, 256, 8, 64)
 value = torch.randn(1, 256, 8, 64)
 
 output = scaled_dot_product_attention(query, key, value)
 print(output.shape) # torch.Size([1, 256, 8, 64])

對(duì)于每個(gè)查詢頭，都有一個(gè)對(duì)應(yīng)的鍵。這個(gè)過程如下圖所示:

而GQA將查詢頭分成G組，每組共享一個(gè)鍵和值。可以表示為:

使用可視化的表示就能非常清楚的了解GQA的工作原理，就像我們上面說的那樣，GQA是一個(gè)相當(dāng)簡單和干凈的想法

Pytorch代碼實(shí)現(xiàn)

讓我們編寫代碼將這種將查詢頭劃分為G組，每個(gè)組共享一個(gè)鍵和值。我們可以使用einops庫有效地執(zhí)行對(duì)張量的復(fù)雜操作。

首先，定義查詢、鍵和值。然后設(shè)置注意力頭的數(shù)量，數(shù)量是隨意的，但是要保證num_heads_for_query % num_heads_for_key = 0，也就是說要能夠整除。我們的定義如下：

import torch
 
 # shapes: (batch_size, seq_len, num_heads, head_dim)
 query = torch.randn(1, 256, 8, 64)
 key = torch.randn(1, 256, 2, 64)
 value = torch.randn(1, 256, 2, 64)
 
 num_head_groups = query.shape[2] // key.shape[2]
 print(num_head_groups) # each group is of size 4 since there are 2 kv_heads

為了提高效率，交換seq_len和num_heads維度，einops可以像下面這樣簡單地完成:

from einops import rearrange
 
 query = rearrange(query, "b n h d -> b h n d")
 key = rearrange(key, "b s h d -> b h s d")
 value = rearrange(value, "b s h d -> b h s d")

然后就是需要在查詢矩陣中引入”分組“的概念。

from einops import rearrange
 query = rearrange(query, "b (h g) n d -> b g h n d", g=num_head_groups)
 print(query.shape) # torch.Size([1, 4, 2, 256, 64])

上面的代碼我們將二維重塑為二維：對(duì)于我們定義的張量，原始維度8(查詢的頭數(shù))現(xiàn)在被分成兩組(以匹配鍵和值中的頭數(shù))，每組大小為4。

最后最難的部分是計(jì)算注意力的分?jǐn)?shù)。但其實(shí)它可以在一行中通過insum操作完成的

from einops import einsum, rearrange
 # g stands for the number of groups
 # h stands for the hidden dim
 # n and s are equal and stands for sequence length
  
 scores = einsum(query, key, "b g h n d, b h s d -> b h n s")
 print(scores.shape) # torch.Size([1, 2, 256, 256])

scores張量和上面的value張量的形狀是一樣的。我們看看到底是怎么操作的

einsum幫我們做了兩件事:

1、一個(gè)查詢和鍵的矩陣乘法。在我們的例子中，這些張量的形狀是(1,4,2,256,64)和(1,2,256,64)，所以沿著最后兩個(gè)維度的矩陣乘法得到(1,4,2,256,256)。

2、對(duì)第二個(gè)維度(維度g)上的元素求和——如果在指定的輸出形狀中省略了維度，einsum將自動(dòng)完成這項(xiàng)工作，這樣的求和是用來匹配鍵和值中的頭的數(shù)量。

最后是注意分?jǐn)?shù)與值的標(biāo)準(zhǔn)乘法:

import torch.nn.functional as F
 
 scale = query.size(-1) ** 0.5
 attention = F.softmax(similarity / scale, dim=-1)
 
 # here we do just a standard matrix multiplication
 out = einsum(attention, value, "b h n s, b h s d -> b h n d")
 
 # finally, just reshape back to the (batch_size, seq_len, num_kv_heads, hidden_dim)
 out = rearrange(out, "b h n d -> b n h d")
 print(out.shape) # torch.Size([1, 256, 2, 64])

這樣最簡單的GQA實(shí)現(xiàn)就完成了，只需要不到16行python代碼:

最后再簡單提一句MQA：多查詢注意(MQA)是另一種簡化MHA的流行方法。所有查詢將共享相同的鍵和值。原理圖如下:

可以看到，MQA和MHA都可以從GQA推導(dǎo)出來。具有單個(gè)鍵和值的GQA相當(dāng)于MQA，而具有與頭數(shù)量相等的組的GQA相當(dāng)于MHA。

GQA的好處是什么?

GQA是最佳性能(MQA)和最佳模型質(zhì)量(MHA)之間的一個(gè)很好的權(quán)衡。

下圖顯示，使用GQA，可以獲得與MHA幾乎相同的模型質(zhì)量，同時(shí)將處理時(shí)間提高3倍，達(dá)到MQA的性能。這對(duì)于高負(fù)載系統(tǒng)來說可能是必不可少的。

在pytorch中沒有GQA的官方實(shí)現(xiàn)。所以我找到了一個(gè)比較好的非官方實(shí)現(xiàn)，有興趣的可以試試：

https://github.com/fkodom/grouped-query-attention-pytorch

GQA論文：

https://arxiv.org/pdf/2305.13245.pdf