MHA · MQA · GQA✓ Mathematical
◆ The PatternRunning multiple attention computations in parallel
Instead of one big attention operation, we split into h heads, each with dimension d_k = d_model / h. Each head learns different patterns — one might track syntax, another coreference, another positional relationships.
MultiHead(Q,K,V) = Concat(head₁, ..., headₕ) · W_O
Run h parallel attention ops with different projections, concat results, project back to d_model.
MQA: all heads share K,V — only Q varies per head
Multi-Query Attention. Massive KV-cache savings (÷ h). Used in PaLM, Falcon.
GQA: groups of heads share K,V — middle ground
Grouped-Query Attention. G groups, each with h/G query heads. LLaMA 2 70B uses 8 KV heads for 64 query heads.
Standard MHA with 32 heads and d_model=4096 → d_k=128 per head. GQA is now the default for large models: it trades a tiny quality hit for huge inference speedups via smaller KV caches.
Interactive — compare MHA, MQA, GQA head layouts
Python — grouped-query attention#
import torch, torch.nn as nn, torch.nn.functional as F
class GQA(nn.Module):
def __init__(self, d, n_heads, n_kv_heads):
super().__init__()
self.n_heads, self.n_kv = n_heads, n_kv_heads
self.d_k = d // n_heads
self.W_q = nn.Linear(d, n_heads * self.d_k, bias=False)
self.W_k = nn.Linear(d, n_kv_heads * self.d_k, bias=False)
self.W_v = nn.Linear(d, n_kv_heads * self.d_k, bias=False)
self.W_o = nn.Linear(n_heads * self.d_k, d, bias=False)
def forward(self, x):
B, T, _ = x.shape
q = self.W_q(x).view(B, T, self.n_heads, self.d_k).transpose(1, 2)
k = self.W_k(x).view(B, T, self.n_kv, self.d_k).transpose(1, 2)
v = self.W_v(x).view(B, T, self.n_kv, self.d_k).transpose(1, 2)
# Repeat KV heads to match query heads
r = self.n_heads // self.n_kv
k = k.repeat_interleave(r, dim=1)
v = v.repeat_interleave(r, dim=1)
attn = F.scaled_dot_product_attention(q, k, v, is_causal=True)
return self.W_o(attn.transpose(1,2).reshape(B, T, -1))Pattern bridge: Multiple attention heads capture different relationship types in parallel — like running several correlation analyses simultaneously. In markets, combining RSI, MACD, and volume is multi-headed analysis.