In this article, we'll talk about different types of attention that are released after the basic multi headed attention with some changes. We'll discuss about their basic architectures, differences and what changes do they make.
Multi Head Attention :-
This is the very basic type of attention which is generally used most of the time. The foundational idea is very simple, the total no. of dimensions are broken down into a particular no. of heads with their respective dimensions. For example,
input x dim → (batch_size, seq_len, total_dim)
broken down into → (batch_size, seq_len, no_of_heads, head_dim)
Group Query Attention :-
The basic idea of group query attention is to share the some no. of query heads with some no. of keys and values. In this type of attention, instead of just taking the same no. of heads for all key, query and value, we take different no. of value for query heads and (key and value) heads. For example,
input x dim → (batch_size, seq_len, total_dim)
broken down into →
query: (batch_size, seq_len, no_q_heads, head_dim)
key & values: (batch_size, seq_len, no_kv_heads, head_dim)
After this the key and values are repeated for the no. of query head to process the attention part.
Multi Query Attention :-
This can be considered as a part of Group Query Attention where number of kv heads is 1. It just means that we are sharing query heads for a single head of key and values.
input x dim → (batch_size, seq_len, total_dim)
broken down into →
query: (batch_size, seq_len, no_q_heads, head_dim)
key & values: (batch_size, 1, seq_len, head_dim)
Differences
With these changes, there are also some pros and cons attached. The speed does increase with reduction in no. of heads which shows it effect on the quality of results. If they are used in bigger and bigger models, it might affect significantly. There is always a trade-off between speed, efficiency and quality. The GQA has been used in LLaMA-2 and Mistral-7B.
The code for all the three with their dimensions, their time and memory statistics is discussed below.
import math
import torch
import torch.nn as nn
from time import time
import torch.profiler
class multi_head_atnn(nn.Module):
def __init__(self, total_dim: int, n_heads: int):
super().__init__()
self.nh = n_heads
self.dim = total_dim
assert self.dim % self.nh == 0, "must be divisible"
self.head_dim = self.dim // self.nh
self.wq = nn.Linear(self.dim, self.dim)
self.wk = nn.Linear(self.dim, self.dim)
self.wv = nn.Linear(self.dim, self.dim)
self.wo = nn.Linear(self.dim, self.dim)
def forward(self, x: torch.Tensor):
bs, sl, dims = x.shape
# (bs, sl, dim) --> (bs, sl, dim)
q = self.wq(x)
k = self.wk(x)
v = self.wv(x)
# (bs, sl, dim) --> (bs, sl, nh, hd) --> (bs, nh, sl, hd)
q = q.view(bs, sl, self.nh, self.head_dim).permute(0, 2, 1, 3)
k = k.view(bs, sl, self.nh, self.head_dim).permute(0, 2, 1, 3)
v = v.view(bs, sl, self.nh, self.head_dim).permute(0, 2, 1, 3)
# (bs, nh, sl, hd) @ (bs, nh, hd, sl) --> (bs, nh, sl, sl)
attn = (q @ k.transpose(2, 3)) / math.sqrt(self.head_dim)
att = nn.functional.softmax(attn, dim=-1)
# (bs, nh, sl, sl) @ (bs, nh, sl, hd) --> (bs, nh, sl, hd)
atts = att @ v
# (bs, nh, sl, hd) --> (bs, sl, nh, hd) --> (bs, sl, nh * hd)
x = atts.permute(0, 2, 1, 3).contiguous().view(bs, -1, self.dim)
# (bs, sl, nh * hd) --> (bs, sl, dim)
x = self.wo(x)
return x
def repetition(x: torch.Tensor, factor: int):
bs, sl, nh, hd = x.shape
y = x[:, :, :, None, :].expand(bs, sl, nh, factor, hd).contiguous().view(bs, sl, nh * factor, hd)
return y
class multi_query_atnn(nn.Module):
def __init__(self, num_h, dim):
super().__init__()
self.dim = dim
self.h_dim = dim // num_h
self.nh = num_h
self.wq = nn.Linear(self.dim, self.dim)
self.wk = nn.Linear(self.dim, self.h_dim)
self.wv = nn.Linear(self.dim, self.h_dim)
self.wo = nn.Linear(self.dim, self.dim)
def forward(self, x: torch.Tensor):
bs, sl, dim = x.shape
q = self.wq(x) # (bs, sl, dim)
k = self.wk(x) # (bs, sl, hd)
v = self.wv(x)
q = q.view(bs, -1, self.nh, self.h_dim).permute(0, 2, 1, 3) # (bs, sl, dim) -> (bs, sl, nh, hd) -> (bs, nh, sl, hd)
k = k.unsqueeze(1) # (bs, sl, hd) -> (bs, 1, sl, hd)
v = v.unsqueeze(1)
attn = (q @ k.transpose(2, 3)) / math.sqrt(self.h_dim)
att = nn.functional.softmax(attn, dim=-1)
x = torch.matmul(att, v)
x = x.permute(0, 2, 1, 3).contiguous().view(bs, -1, self.dim)
x = self.wo(x)
return x
class group_query_attn(nn.Module):
def __init__(self, q_heads: int, kv_heads: int, total_dim: int):
super().__init__()
self.q_h = q_heads
self.kv_h = kv_heads
self.dim = total_dim
self.head_dim = self.dim // self.q_h
self.factor = self.q_h // self.kv_h
self.wq = nn.Linear(self.dim, self.q_h * self.head_dim)
self.wk = nn.Linear(self.dim, self.kv_h * self.head_dim)
self.wv = nn.Linear(self.dim, self.kv_h * self.head_dim)
self.wo = nn.Linear(self.dim, self.dim)
def forward(self, x: torch.Tensor):
bs, sl, dims = x.shape
q = self.wq(x)
k = self.wk(x)
v = self.wv(x)
q = q.view(bs, sl, self.q_h, self.head_dim).permute(0, 2, 1, 3)
k = k.view(bs, sl, self.kv_h, self.head_dim)
v = v.view(bs, sl, self.kv_h, self.head_dim)
k = repetition(k, self.factor).permute(0, 2, 1, 3)
v = repetition(v, self.factor).permute(0, 2, 1, 3)
attn = (q @ k.transpose(2, 3)) / math.sqrt(self.head_dim)
att = nn.functional.softmax(attn, dim=-1)
a = att @ v
x = a.permute(0, 2, 1, 3).contiguous().view(bs, -1, self.dim)
x = self.wo(x)
return x
def time_mem_stats(model_cls, model_name, **model_kwargs):
print(f"for {model_name}")
batch_size = 50
seq_len = 128
total_dim = 32 * 128
device = 'cuda' if torch.cuda.is_available() else 'cpu'
torch.cuda.empty_cache()
torch.cuda.reset_peak_memory_stats()
x = torch.randn(batch_size, seq_len, total_dim, device=device)
model = model_cls(**model_kwargs).to(device)
with torch.profiler.profile(
activities=[
torch.profiler.ProfilerActivity.CPU,
torch.profiler.ProfilerActivity.CUDA
],
profile_memory=True, with_stack=True, record_shapes=True
) as pro:
torch.cuda.synchronize()
start_mem = torch.cuda.memory_allocated()
start = time()
u = model(x)
end = time()
torch.cuda.synchronize()
end_mem = torch.cuda.memory_allocated()
peak_mem = torch.cuda.max_memory_allocated()
print(f'Output shape: {u.shape}')
print(f'Time for forward pass: {end - start:.4f} s')
print(f'Memory before start: {start_mem / 1e6:.2f} MB')
print(f'Memory after: {end_mem / 1e6:.2f} MB')
print(f'Peak memory: {peak_mem / 1e6:.2f} MB')
print("Stats:")
print(pro.key_averages().table(sort_by="cuda_time_total", row_limit=5))
print("=" * 80)
if __name__ == "__main__":
time_mem_stats(
multi_head_atnn, "Multi-Head Attention",
total_dim=32 * 128, n_heads=32
)
time_mem_stats(
group_query_attn, "Group Query Attention",
q_heads=32, kv_heads=4, total_dim=32 * 128
)
time_mem_stats(
multi_query_atnn, "Multi-Query Attention",
num_h=32, dim=32 * 128
)
Stats
- Multi Query Attention
- Group Query Attention
- Multi Head Attention
Conclusion:
From the results, we can see that the CUDA and CPU time for
Multi Head Attention > Group Query Attention > Muti Query Attention
Also from the memory point of view, it is in the same order.