Multi-Head Attention — Parallel Views

Why this session exists

One attention head produces one distribution over keys for each query. Multiple heads supply several separately normalized retrievals before their outputs are combined. Splitting a fixed model width controls parameter and arithmetic growth, but head count changes the function class, score storage and kernel behavior. We will implement MHA/GQA with masks and a compact cache, compare against independent oracles, and distinguish learned head behavior from random patterns or guaranteed semantic roles.

Prerequisites

S107 attention, S108 Q/K/V derivatives and S102 multiple filters. Know reshape versus transpose and that a tensor's shape does not identify its semantic axes.

(a) Intuition · 5 min

Several readers can examine the same poem with different lenses. Their notes are concatenated and linearly mixed, not necessarily averaged like an ensemble vote. Heads train jointly and can become redundant; the architecture permits different patterns but does not assign one reader 'syntax' and another 'coreference'. One softmax can attend to multiple regions at once, so the point is independently addressable mixtures, not a theorem that one head can focus on only one item.

What multiple heads add beyond one big head

Let head i output $H_i=A_iV_i$. Then

Each row block $W_O^{(i)}$ maps one head into the common output coordinates. Separate distributions can preserve information from distinct retrievals that a single mixture might confound, at the cost of narrower projections. This is a useful architectural bias, not a universal performance ordering. Downstream layers can mix concatenated slices even without this particular W_O; removing it changes expressiveness/interface but does not make collaboration impossible throughout the entire network.

Why splitting d_model across heads is genius

Under the conventional $d_h=d/h$ and bias-free square input/output projections, Q/K/V use $3d^2$ parameters and W_O uses $d^2$: total $4d^2$, independent of h. Eight full-width heads plus concatenation-to-d output instead cost $4hd^2$. At d=512 and h=8, standard attention has 1,048,576 projection weights; the full-width alternative has 8,388,608. Comparing only Q/K/V would give $3d^2$ versus $3hd^2$, not the complete module count.

Dense pairwise multiply work scales as $hT^2d_h=T^2d$, but softmax work and explicitly stored maps scale with $hT^2$. Thus equal projection count and leading matmul work do not mean equal latency, memory or representation capacity. Each linear QK score form has rank at most $d_h$; after nonlinear softmax the attention matrix need not have that rank. No universal 64-dimensional floor exists—small test models work below it. Unknown proprietary head-count rumors are not design evidence.

Multi-head attention pipeline

A head axis is an additional batch-like axis for matmul, not an invitation to reinterpret contiguous (B,T,d) bytes directly as (B,h,T,dh). That wrong reshape can retain a plausible shape while mixing token and head coordinates.

The shape-wrangling recipe

Project all input features into concatenated head coordinates; each head sees a learned projection of the full input, not just an original input slice. Reshape (B,T,h,dh), transpose time/head, compute scores (B,h,L,S) and normalize the key axis. Concatenate by transposing back then reshaping; contiguous().view(...) explicitly copies if needed, whereas reshape may copy. Transpose does not always make every possible view illegal; the particular layout and requested reshape decide.

Anatomy of the multi-head layer

For self-attention with equal widths, packed projections or separate Q/K/V Linear calls are mathematically equivalent when weights match. Several separate small matmuls are also valid, though often less convenient; 'never use h matmuls' is not a correctness rule. nn.Linear stores (out,in), so packing into PyTorch in_proj_weight concatenates along its output-row axis. Scale by $\sqrt{d_h}$, not $\sqrt d$. Output projection mixes head features independently at each query position; it does not itself route between positions.

MHA vs MQA vs GQA — the modern variants

For h query heads, g KV heads and width $d_h=d/h$:

This implementation requires positive h/g, d divisible by h, and h divisible by g. Other architectures can choose different projection widths. Compact K/V state shrinks by h/g relative to MHA at the same widths; QK score count still follows h query heads. Explicitly repeating K/V tensors is a useful reference implementation but expands temporary memory; cache the unrepeated tensors and use a suitable fused/grouped kernel for actual memory efficiency. MQA saves $(1-1/h)$ of the K/V projection parameters, not always 50%; at h=4 it saves 75% of those matrices and 37.5% of all four projections.

The GQA paper studies sharing and checkpoint conversion/uptraining. Quality/speed tradeoffs depend on architecture, training and workload; it does not prove a universal sub-one-percent loss or a free conversion. Multi-head Latent Attention in DeepSeek-V2 compresses K/V through a learned latent representation with additional positional handling; it is a different cache factorization, not just setting g smaller. This lesson explains that mechanism conceptually, without reproducing a DeepSeek kernel or benchmark.

For MLA's content path, a token state h is down-projected to a latent c, then content keys/values are linear up-projections of c. Without position-dependent rotations between those maps, key up-projection can be absorbed into a query map and value up-projection into the output map at inference. DeepSeek-V2 separately carries a rotary-position key component, so the cache is not just the latent: its report gives per-token per-layer elements as latent width plus the decoupled rotary-key width. This explains why simply applying ordinary RoPE after expanding every cached key would defeat the intended compression/absorption path.

What different heads actually learn (Clark et al. 2019)

Clark et al. analyze BERT attention patterns including delimiters, positional behavior, syntax and coreference. Those observations are checkpoint/task/layer-specific, not universal head IDs such as the earlier internally inconsistent 'head 8-10, layer 6'. A random-initialized matrix has no measured linguistic specialization. To study a trained model, retain checkpoint/tokenizer/input and per-head maps, test hypotheses on many examples, then intervene and measure task effects. Visualization alone cannot establish causality.

A committee with a fixed budget: choosing heads and sharing

Few wide heads allow higher-rank score projections; many narrow heads allow more independently normalized mixtures. Equal parameter count does not settle which wins. Head pruning can reveal redundancy in a trained task, but zeroing a head does not necessarily speed a dense kernel; structural removal and a compatible implementation are separate work. GQA is attractive when cache bandwidth/storage binds, yet changes learned architecture and may require uptraining. Benchmark serving with fixed batch/context/dtype and quality goals, not a universal 64–128-width rule.

(c) Hands-on · 25 min

The exact CPU program implements MHA as the g=h case of a grouped module, preserving the original separate projection and reshape teaching. It checks a framework MHA oracle (including gradients), a per-head GQA oracle, rectangular cached masks, invalid configurations and a wrong-axis reshape negative control. The repeated K/V tensors are temporary reference work; returned caches remain compact.

import math,copy
import torch
from torch import nn
from torch.nn import functional as F
 
torch.set_num_threads(1);torch.manual_seed(23)
class GroupedQueryAttention(nn.Module):
    def __init__(self,d_model,n_heads,n_kv_heads=None):
        super().__init__();g=n_heads if n_kv_heads is None else n_kv_heads
        if min(d_model,n_heads,g)<=0 or d_model%n_heads or n_heads%g:
            raise ValueError('positive divisible dimensions required')
        self.d=d_model;self.h=n_heads;self.g=g;self.dh=d_model//n_heads
        self.W_q=nn.Linear(d_model,d_model,bias=False)
        self.W_k=nn.Linear(d_model,g*self.dh,bias=False)
        self.W_v=nn.Linear(d_model,g*self.dh,bias=False)
        self.W_o=nn.Linear(d_model,d_model,bias=False)
    def forward(self,x,causal=False,cache=None,valid=None):
        B,T,D=x.shape
        if D!=self.d or T==0:raise ValueError('input width/length')
        q=self.W_q(x).reshape(B,T,self.h,self.dh).transpose(1,2)
        k=self.W_k(x).reshape(B,T,self.g,self.dh).transpose(1,2)
        v=self.W_v(x).reshape(B,T,self.g,self.dh).transpose(1,2)
        prefix=0
        if cache is not None:
            if not causal:raise ValueError('cache supports causal inference only')
            ck,cv=cache
            if ck.shape!=cv.shape or ck.shape[:2]!=(B,self.g) or ck.shape[-1]!=self.dh:
                raise ValueError('cache shape')
            prefix=ck.shape[-2];k=torch.cat([ck,k],-2);v=torch.cat([cv,v],-2)
        S=k.shape[-2];allowed=torch.ones(B,1,T,S,dtype=torch.bool,device=x.device)
        if causal:allowed &= torch.arange(S,device=x.device)[None,:]<=torch.arange(prefix,prefix+T,device=x.device)[:,None]
        if valid is not None:allowed &= valid[:,None,None,:]
        if not allowed.any(-1).all():raise ValueError('all-masked query')
        kr=k.repeat_interleave(self.h//self.g,1);vr=v.repeat_interleave(self.h//self.g,1)
        a=(q@kr.transpose(-2,-1)/math.sqrt(self.dh)).masked_fill(~allowed,-torch.inf).softmax(-1)
        out=(a@vr).transpose(1,2).reshape(B,T,D)
        return self.W_o(out),a,(k,v)
 
MultiHeadAttention=GroupedQueryAttention
B,T,D,H=2,5,12,4;x=torch.randn(B,T,D,dtype=torch.double,requires_grad=True)
m=MultiHeadAttention(D,H).double();out,a,_=m(x,causal=True)
ref=nn.MultiheadAttention(D,H,bias=False,batch_first=True,dropout=0.).double()
with torch.no_grad():
    ref.in_proj_weight.copy_(torch.cat([m.W_q.weight,m.W_k.weight,m.W_v.weight]))
    ref.out_proj.weight.copy_(m.W_o.weight)
blocked=torch.ones(T,T,dtype=torch.bool).triu(1)
expected,ar=ref(x,x,x,attn_mask=blocked,need_weights=True,average_attn_weights=False)
assert torch.allclose(out,expected,atol=1e-12) and torch.allclose(a,ar,atol=1e-12)
g=torch.randn_like(out)
g1=torch.autograd.grad((out*g).sum(),[x]+list(m.parameters()),retain_graph=True)
g2=torch.autograd.grad((expected*g).sum(),[x,ref.in_proj_weight,ref.out_proj.weight])
assert torch.allclose(g1[0],g2[0],atol=1e-12)
assert torch.allclose(torch.cat(g1[1:4]),g2[1],atol=1e-12)
assert torch.allclose(g1[4],g2[2],atol=1e-12)
print('per-head random routing',a[0].detach().round(decimals=3))
for kv in [4,2,1]:
    model=GroupedQueryAttention(D,H,kv).double();full,aa,cache=model(x,causal=True)
    with torch.no_grad():
        first,_,c=model(x[:,:2],causal=True)
        last,_,c=model(x[:,2:],causal=True,cache=c)
    assert torch.allclose(full,torch.cat([first,last],1),atol=1e-12)
    assert c[0].shape==(B,kv,T,D//H)
    assert sum(p.numel() for p in model.parameters())==2*D*D+2*D*kv*(D//H)
    # Per-head loop oracle reads compact KV; no repeat_interleave.
    q=model.W_q(x).reshape(B,T,H,D//H).transpose(1,2)
    k,v=cache;heads=[]
    for h in range(H):
        j=h//(H//kv);scores=q[:,h]@k[:,j].transpose(-2,-1)/math.sqrt(D//H)
        weight=scores.masked_fill(blocked,-torch.inf).softmax(-1)
        heads.append(weight@v[:,j])
    oracle=model.W_o(torch.cat(heads,-1));assert torch.allclose(full,oracle,atol=1e-12)
    changed=x.detach().clone();changed[:,3:]+=100
    assert torch.allclose(full[:,:3],model(changed,causal=True)[0][:,:3],atol=1e-12)
    print('KV heads',kv,'compact cache elements',2*c[0].numel())
# A transpose cannot be replaced by a reshape with the same displayed shape.
raw=torch.arange(B*T*D).reshape(B,T,D)
correct=raw.reshape(B,T,H,D//H).transpose(1,2)
wrong=raw.reshape(B,H,T,D//H);assert not torch.equal(correct,wrong)
for bad in [(12,0,1),(13,3,1),(12,3,2)]:
    try:GroupedQueryAttention(*bad);raise AssertionError('invalid dimensions accepted')
    except ValueError:pass
for heads in [1,2,3,6]:
    model=MultiHeadAttention(D,heads);assert sum(p.numel() for p in model.parameters())==4*D*D
# Same projection bytes, different head partition: not the same function.
one=MultiHeadAttention(D,1).double();one.load_state_dict(m.state_dict())
assert not torch.allclose(one(x,causal=True)[0],out)
print('MHA gradients, grouped loop, cache, parameter and shape controls passed')

Anatomy of the script

All parameters are explicit; the g=h and g=1 endpoints implement MHA and MQA. The per-head loop independently selects each group's compact K/V rather than repeating the same grouping expression. The cache test consumes two tokens then three, so a wrong upper-left rectangular causal mask cannot pass by coincidence. Cache reuse is for inference with compatible model/version/dtype/batch ordering; detach or use no-grad to avoid retaining training graphs. Padding keys require a mask spanning the entire cached plus new sequence. Padded query positions still need exclusion from the task loss.

The nn.MultiheadAttention comparison deliberately inverts the allowed/blocked mask convention. It copies weight orientation exactly and checks input and all projection gradients. Different head counts with the same weights produce different functions without implying different quality; parameter invariance is not a language-model benchmark.

Head-count practice: a bounded next-token experiment

This second block runs after the first. It trains one- and four-head causal models on a deterministic inline sequence with shared initial weight values and data order. It is a synthetic memorization/optimization fixture, not a trained BERT study, a linguistic-specialization result or proof that four heads win. A real comparison needs independent documents, comparable tuning budgets, multiple seeds and held-out quality/latency.

class TinyLM(nn.Module):
    def __init__(self,heads):
        super().__init__();self.emb=nn.Embedding(5,16);self.pos=nn.Parameter(torch.randn(1,8,16)*.01)
        self.attn=MultiHeadAttention(16,heads);self.norm=nn.LayerNorm(16);self.head=nn.Linear(16,5)
    def forward(self,ids):
        x=self.emb(ids)+self.pos[:,:ids.shape[1]];z,_,_=self.attn(x,causal=True)
        return self.head(self.norm(x+z))
 
torch.manual_seed(29);base=TinyLM(1);state=copy.deepcopy(base.state_dict())
seq=torch.tensor([[0,1,2,3,4,0,1,2,3],[1,2,3,4,0,1,2,3,4]])
for h in [1,4]:
    model=TinyLM(h);model.load_state_dict(state);opt=torch.optim.Adam(model.parameters(),lr=.02)
    losses=[]
    for step in range(25):
        opt.zero_grad(set_to_none=True);z=model(seq[:,:-1])
        loss=F.cross_entropy(z.flatten(0,1),seq[:,1:].flatten());loss.backward();opt.step();losses.append(loss.item())
    assert all(math.isfinite(v) for v in losses)
    print(h,'heads: synthetic train loss first/last',losses[0],losses[-1])

(d) Production reality · 15 min

Serving: compact cache bytes are $2BNH_{kv}Td_hb$. For B=1, N=80, Hkv=64, T=2,048, dh=128 and two-byte elements, full MHA takes 5,368,709,120 bytes = 5 GiB; Hkv=8 takes 671,088,640 bytes = 0.625 GiB. State exactly which heads are queries versus KV. These are arithmetic configurations, not an assertion that a particular shipped LLaMA checkpoint uses full MHA. Increasing batch size multiplies this component. Model weights, temporary expanded KV, allocator and other activations remain additional memory.

MQA/GQA papers motivate lower incremental decoding bandwidth. The original PaLM anecdote's particular head count, 12× speedup and universal quality loss are not established here. A changed head grouping usually needs compatible training or conversion/uptraining; it is not just a serving flag.

Induction circuits: the 2022 study analyzes a previous-token/induction-head composition implementing patterns like A B ... A → B. It gives stronger causal evidence in small attention-only models and qualified evidence for larger models. It does not establish one universal emergence token count or that every head in every model implements a named algorithm. Effects on next-token prediction must be measured, not inferred solely from attention mass.

Pruning: Michel et al. (2019) find many heads removable in their studied settings, with important differences among attention types. Do not attribute that result generically to 'Meta' or assign every future model the same removable percentage. Evaluate one-at-a-time and joint pruning on validation, then test structurally pruned kernels for real latency. Ablating weights without shrinking operations is not a speedup.

Where this shows up in the rest of the plan

Position information may be added once to inputs, or applied inside layers/per-head (for example RoPE or relative biases); the earlier 'once shared by all heads' rule was too narrow. FFN work can dominate at some widths/lengths, so MHA is not always the compute-dominant block. Encoder versus causal-decoder objectives alter visibility and state. Sparse kernels, GQA and paged caches each address different costs; speculative decoding interacts with cache state but is not simply another head-count variant.

(e) Recall + stretch · 10 min

1. Why split? Keep projection count/leading pairwise matmul width bounded while providing several normalized mixtures; not free in softmax/storage/latency.
2. 512 and eight heads? dh=64; (B,T,512) → (B,T,8,64) → (B,8,T,64).
3. W_O? Linear mixing of concatenated head features; later layers could mix too, but this module provides an explicit common output basis.
4. GQA? Groups of query heads share K/V; cache shrinks h/g, quality and runtime are empirical.
5. Cross-attention? Queries from one sequence and keys/values from another; same head operations with possibly rectangular lengths. Not the original CLIP dual-encoder interface.
6. Every row selects key zero? Inspect masks, key variation, scores, scale, dtype and data; it may be a mask bug or learned pattern, not necessarily overflow.
7. Worked cache? Five GiB for the stated full-MHA configuration, 0.625 GiB for eight KV heads, before other memory.
8. Equal importance? No general guarantee; measure task effects and joint redundancy rather than transplant a pruning percentage.
9. Induction head? A circuit-supported pattern completion mechanism, with scope/evidence tied to studied models.
10. Different context lengths per head? Mix global and width-W local visibility. Pairwise work is proportional to $h_{local}T W d_h+h_{global}T^2d_h$, not $h_{local}W^2$. A dense implementation still computes masked-out pairs, so actual savings require sparse/window kernels. This is a proposed design, not a claim that Mistral uses this exact head mixture.

Explain-out-loud test

Derive both full and grouped projection counts, demonstrate a shape-correct wrong transpose, and explain why a compact cache plus repeated temporary K/V does not establish peak-memory savings.

What comes next

Positional encoding supplies order information under particular architectural conventions. Learning-plan hub.

Primary references

• Transformer §3.2.2, GQA, DeepSeek-V2 MLA.
• BERT attention analysis, head pruning, induction circuits.