Multi-Head Attention — Shapes, Specialization and KV Budgets

Several attention distributions, one fixed width

A single attention head already forms a distribution over many keys; it is not forced to pick exactly one token. Multi-head attention gives several independently parameterized distributions and value projections, then mixes their results. At fixed model width d and head count h, the common choice is key/value width d/h. More heads change the function class and inductive bias, not simply multiply parameter capacity.

The original Transformer paper reports that single-head attention was 0.9 BLEU below its best configuration and that too many heads also reduced quality in that ablation. This is not “every benchmark,” eight is not a universal optimum, and going from eight to one is not halving the head count. This lesson restores the actual mechanism instead of promising a fixed validation-loss gain.

Prerequisites: scaled dot-product attention and reshape/broadcast semantics from S036 and S003. Outcomes: implement self/cross MHA with explicit masks; prove shape and parameter contracts; distinguish positional/syntax hypotheses from observed pretrained evidence; and calculate MHA/MQA/GQA/MLA cache memory correctly.

Equations, parameter count and output mixing

Queries have shape (B,h,T_q,d_k), keys (B,h,T_k,d_k), values (B,h,T_k,d_v), and scores (B,h,T_q,T_k). Key and value sequence lengths must match; query length may differ. Softmax is over keys, the final score axis. All projections read the full input feature vector; splitting projected features is not cutting the original embedding into independent raw slices.

For d_k=d_v=d/h and square input/output widths, bias-free projections contain 3d²+d²=4d² parameters. With four ordinary biases, add 4d. This formula assumes h divides d and that widths change with h. Holding per-head width fixed while increasing h does increase parameters. Leading QK/AV arithmetic at fixed total width is O(B T² d), while materialized score storage is O(B h T²); head count still changes softmax work and memory.

Concatenation just places features beside each other. W_O learns a dense write into residual-stream coordinates. It is useful, not logically required for a model to run; without it a later dense layer could still mix features, but the immediate write is constrained. In the row-vector equation above, partition W_O into input-row blocks: each head multiplies its own block, and those projected vectors sum into the stream. PyTorch stores nn.Linear.weight transposed relative to that equation, so the same head occupies input columns in mha.W_O.weight[:, lo:hi]. Zeroing that stored column block ablates the head's write, but does not automatically remove computation or shrink the checkpoint.

Split and merge without scrambling positions

The correct journey is (B,T,d) → (B,T,h,d_k) → (B,h,T,d_k). The reverse is transpose back to (B,T,h,d_v), then reshape to (B,T,h d_v), then project. A direct view into (B,h,T,d_k) usually groups the wrong elements; matching final shape is not a numerical test.

A transpose changes strides and can be noncontiguous. view requires compatible strides, not universal contiguity. contiguous().view(...) is a safe merge here; reshape(...) copies only when necessary. Neither implies a fixed runtime penalty.

Complete MHA and independent loop parity

Save the Python fences as mha_lab.py for S039. allowed is boolean True=can attend and must be explicitly broadcastable to (B,h,T_q,T_k). A padding mask (B,T_k) must first become (B,1,1,T_k); accepting it blindly can mix batch/query axes. Invalid query rows are not silently normalized; either reject them as here or define an explicit zero-output policy separately. Dropout is applied after softmax and before AV. Its sampled rows need not sum to one, although inverted dropout preserves their expectation. Dropping raw scores is a different, not necessarily ineffective, regularizer.

import math
import torch
from torch import nn
from torch.nn import functional as F
 
class MultiHeadAttention(nn.Module):
    def __init__(self,d_model,h,dropout=0.,bias=False):
        super().__init__()
        if h<1 or d_model<1 or d_model%h:raise ValueError('positive divisible width/head count required')
        self.d_model=d_model;self.h=h;self.d_k=d_model//h
        self.W_Q=nn.Linear(d_model,d_model,bias=bias)
        self.W_K=nn.Linear(d_model,d_model,bias=bias)
        self.W_V=nn.Linear(d_model,d_model,bias=bias)
        self.W_O=nn.Linear(d_model,d_model,bias=bias)
        self.dropout=nn.Dropout(dropout)
    def forward(self,x,allowed=None,memory=None,return_weights=False):
        source=x if memory is None else memory
        B,T,C=x.shape;Bk,S,Ck=source.shape
        if B!=Bk or C!=self.d_model or Ck!=C:raise ValueError('input/memory shape contract')
        def split(z,length):return z.reshape(B,length,self.h,self.d_k).transpose(1,2)
        q=split(self.W_Q(x),T);k=split(self.W_K(source),S);v=split(self.W_V(source),S)
        scores=q@k.transpose(-1,-2)/math.sqrt(self.d_k)
        if allowed is not None:
            if allowed.dtype!=torch.bool:raise ValueError('mask must be boolean True=allowed')
            mask=torch.broadcast_to(allowed,scores.shape)
            if not mask.any(-1).all():raise ValueError('no valid key in a query row')
            scores=scores.masked_fill(~mask,-torch.inf)
        weights=scores.softmax(-1)
        out=self.dropout(weights)@v
        result=self.W_O(out.transpose(1,2).reshape(B,T,C))
        return (result,weights) if return_weights else result
 
def loop_oracle(model,x,allowed):
    B,T,C=x.shape;outputs=[]
    for b in range(B):
        heads=[]
        for head in range(model.h):
            lo=head*model.d_k;hi=lo+model.d_k
            q=F.linear(x[b],model.W_Q.weight[lo:hi]);k=F.linear(x[b],model.W_K.weight[lo:hi]);v=F.linear(x[b],model.W_V.weight[lo:hi])
            rows=[]
            for t in range(T):
                keys=[j for j in range(T) if bool(allowed[t,j])]
                score=torch.stack([(q[t]*k[j]).sum()/math.sqrt(model.d_k) for j in keys])
                rows.append(sum(w*v[j] for w,j in zip(score.softmax(0),keys)))
            heads.append(torch.stack(rows))
        outputs.append(F.linear(torch.cat(heads,-1),model.W_O.weight))
    return torch.stack(outputs)
 
torch.manual_seed(37)
mha=MultiHeadAttention(8,2).double();x=torch.randn(2,4,8,dtype=torch.float64,requires_grad=True)
mask=torch.ones(4,4,dtype=torch.bool).tril()
y,w=mha(x,mask,return_weights=True);oracle=loop_oracle(mha,x,mask)
assert y.shape==x.shape and torch.allclose(y,oracle,atol=1e-12)
g1=torch.autograd.grad(y.square().sum(),x,retain_graph=True)[0]
g2=torch.autograd.grad(oracle.square().sum(),x)[0]
assert torch.allclose(g1,g2,atol=1e-12)
assert w[...,~mask].eq(0).all() and torch.allclose(w.sum(-1),torch.ones_like(w.sum(-1)))
changed=x.detach().clone();changed[:,3]+=100
assert torch.allclose(mha(changed,mask)[:,:3],y.detach()[:,:3],atol=1e-12)
labelled=torch.arange(2*4*8).reshape(2,4,8)
assert not torch.equal(labelled.reshape(2,2,4,4),labelled.reshape(2,4,2,4).transpose(1,2))
for heads in (1,2,4,8):assert sum(p.numel() for p in MultiHeadAttention(8,heads).parameters())==4*8*8
try:mha(x,torch.zeros(4,4,dtype=torch.bool))
except ValueError:pass
else:raise AssertionError('all-masked negative control')
print('MHA loop/output/gradient/causal/count contracts: pass')

The loop oracle above checks bias-free self-attention with dropout zero and a shared two-dimensional mask; it is not a reference for the optional biased, cross-attention or stochastic-dropout branches. For SDPA parity, copy these q/k/v tensors to the functional call with the same boolean mask and dropout_p=0. PyTorch nn.MultiheadAttention boolean masks instead use True to disallow attention, so mask conventions cannot be transferred blindly. The functional routine always applies the supplied dropout probability, independently of your containing module's eval mode.

Head patterns, ablations and a fair training comparison

Clark et al. study BERT attention and report delimiter, positional-offset and broad-attention patterns plus correlations with syntax/coreference. Voita et al. study translation encoder heads and prune many while preserving much of the measured BLEU. These are observations in those trained models, not an assignment of permanent roles to your randomly initialized heads. A causal decoder cannot have a next-token-reading head, and attending to the previous token does not make a head an RNN cell. Special-token attention can be a useful null route without proving a scratchpad mechanism.

No BERT weights were downloaded here. For a locally available trusted trained model, put it in eval mode, record pre-dropout attention for a declared sentence, label both axes, then compare several examples/seeds and intervention results. Head importance is not the same as visual sharpness. Ablate one W_O input slice at a time and record held-out NLL/accuracy; validate joint pruning because interactions need not add. Physically changing dimensions is required to turn zeroed weights into resource savings.

The original “60% prunable” and “never train a small-head model” claims were too broad. Voita's abstract reports pruning 38 of 48 encoder heads with a 0.15 BLEU drop on one English–Russian WMT setting. That supports a particular pruning result, not a universal necessity for extra heads during training.

To compare one versus eight heads, use the complete local trainer in S041 with the same vocabulary, split, d=256, four blocks, optimizer/token budget and multiple seeds. Set head counts to 1,8,32; report per-head widths 256,32,8, identical projection counts, actual peak score storage and validation NLL. Do not predict 0.05–0.10 loss gains as guaranteed. A single head is a valid baseline for a small task or explicit constrained architecture. Narrow heads reduce each dot-product subspace but do not prove 32 must always lose to eight.

MQA, GQA and the corrected cache arithmetic

With L layers, batch B, cache length T, g key/value heads, head width d_h and s bytes per element, a conventional unquantized cache has 2 L B T g d_h s bytes. The factor two is K and V. This excludes allocator overhead, padding, activations, weights and replicated shards. MHA uses g=h; MQA uses g=1; GQA uses an intermediate g. Cache grows with KV heads, not query heads alone.

For L=80, B=1, T=4096, d_h=128 and FP16 s=2:

The old per-token table missed a factor two while its total column used different arithmetic. These totals use binary GiB, not decimal GB. Actual Llama 3 70B uses GQA, so the 10 GiB row is its hypothetical MHA comparison, not its deployed cache. Named families and sizes vary; not every Gemma/Llama configuration has exactly eight KV heads.

MQA shares one K/V representation across query heads; GQA shares within groups. Query heads still differ. Parameter matrices shrink for K/V, so these are architecture changes, not free inference toggles on an arbitrary checkpoint. Ainslie et al. describe mean-pooling KV heads followed by uptraining; skipping retraining is not guaranteed equivalent. Shard count may cause replication or uneven partitions; divisibility is an implementation constraint, not a theorem that g must always be a multiple of GPU count.

def kv_bytes(layers,batch,length,kv_heads,head_dim,bytes_per_element=2):
    if min(layers,batch,length,kv_heads,head_dim,bytes_per_element)<1:raise ValueError('positive cache dimensions')
    return 2*layers*batch*length*kv_heads*head_dim*bytes_per_element
assert kv_bytes(80,1,4096,64,128)==10*2**30
assert kv_bytes(80,1,4096,8,128)==int(1.25*2**30)
assert kv_bytes(80,1,4096,1,128)==167772160
# Explicit grouped expansion is a correctness reference, not a memory-saving kernel.
q=torch.randn(1,4,3,2,dtype=torch.float64)
k=torch.randn(1,2,3,2,dtype=torch.float64);v=torch.randn_like(k)
kfull=k.repeat_interleave(2,dim=1);vfull=v.repeat_interleave(2,dim=1)
manual=((q@kfull.transpose(-1,-2))/math.sqrt(2)).softmax(-1)@vfull
sdpa=F.scaled_dot_product_attention(q,kfull,vfull,dropout_p=0.)
assert torch.allclose(manual,sdpa,atol=1e-12)
print('cache and grouped attention arithmetic: pass')

MLA and position are coupled design choices

DeepSeek-V2's MLA compresses joint K/V content into a latent and separates a RoPE key component. A decoder can cache the latent plus that positional component rather than every expanded head. Naively expanding all keys and caching them throws away the benefit. Content projections can be absorbed into adjacent query/output operations; a position-dependent rotation inserted between factors obstructs the same absorption, motivating decoupled RoPE. This is not “rotate half of every latent,” nor automatically identical to GQA.

For an illustrative latent width 512 and positional key width 64 in FP16, cache bytes/token/layer are (512+64)×2=1152. With the hypothetical 80-layer,4096-token budget above this is 377,487,360 bytes, about 0.35156 GiB. DeepSeek's actual layer/configuration must be used for its own total; a latent-only 1024-byte row omits position. No universal 93% saving, quality dominance or claim that no modern model uses MHA follows.

Attention without position and without a fixed non-equivariant mask is permutation-equivariant, not invariant. A causal mask changes the symmetry; cached rectangular queries must use their absolute offsets. A fresh upper-left square triangle is not enough when keys include an older prefix. RoPE, learned positions, score bias and cache layout need joint tests; those mechanisms are developed next.

Practice and worked answers

1. Parameter count. Four bias-free square projections contain 4d² parameters; with biases add 4d. Head count cancels only when d_k=d_v=d/h. More heads still increase the count of stored attention distributions.
2. Six shapes. X is (B,T,d), projected Q is (B,T,d), split is (B,T,h,d_k), transpose is (B,h,T,d_k), per-head result is (B,h,T,d_v), merged/projected output is (B,T,d). Cross-attention uses distinct query/key lengths.
3. Failure fingerprints. A stride-incompatible view raises; a direct (B,h,T,d_k) reshape silently scrambles features/positions. The labeled-integer test catches the latter while output-shape-only tests do not. Softmax on queries and ambiguous mask broadcasting are separate silent bugs.
4. Cache worksheet. Per token/layer at FP16, use 2×g×128×2 bytes: 32768,4096,512 for g=64,8,1. Multiply by batch, layers and tokens, then choose GB or GiB explicitly.
5. When one head is fine. A bounded synthetic task, parameter-controlled baseline or design whose single distribution is sufficient may use one. Measure accuracy/latency; the architecture supplies options, not a universal optimum.
6. Why decoupled RoPE? MLA's content projection absorption is disrupted by position-dependent rotations. A separately cached positional key component preserves the needed relative-position score while allowing compressed content caching. Count that positional component too.

Review shape choreography and the cache worksheet tomorrow; after a week, repeat the loop oracle and one held-out head ablation without relying on a heatmap as explanation.

Sources, optional resources and scope

Primary: Transformer v7, Voita, Clark, MQA, GQA, MLA, PyTorch SDPA, MHA mask convention. Optional original 3Blue1Brown, Umar Jamil, Illustrated Transformer remain optional resource suggestions.

Executed: exact CPU MHA, independent per-head/per-query loop, gradients, causality, mask/layout negative controls and cache arithmetic. Not executed: downloaded BERT, head-pruning quality, Shakespeare rankings, GPU backend speed or serving concurrency. Previous · Next.