Tokens, Embeddings, and Attention

Multi-head attention

Many parallel attention patterns, one layer

~11 minLesson 27 of 57Written by Claude Opus 4.8

One attention layer can only look one way at a time. So a transformer runs several in parallel - each in its own subspace - and reads many relationships at once.

You just built a single attention head: every token forms a query, compares it to every key, and a softmax turns those scores into one probability distribution over the sentence. But one distribution can only point one way. When the model reads the word it, does it look at the previous word for grammar, at the subject noun for topic, or all the way back at the noun it refers to? A single head has to pick. Multi-head attention refuses to choose: it runs several heads side by side, each in its own slice of the representation, then stitches their answers back together.

The core idea in one line

Split the model width into $h$ smaller subspaces, run a full attention operation independently in each one, then concatenate the results and pass them through one output projection. Each head is free to specialise - one learns to track word order, another the sentence's subject, another long-range coreference - and the layer as a whole sees all of those views at the same time.

Why one big head is not enough

Attention ends in a softmax, and a softmax produces exactly one distribution per token - the weights sum to 1 and get sharper the more you concentrate them. That is a feature: it forces the model to commit. But a sentence carries several kinds of relationship simultaneously. Averaging "look at the previous token" and "look at the antecedent 6 words back" into a single distribution blurs both into mush. The fix is not a bigger head - it is more heads, each keeping its own sharp, separate distribution.

More heads, not a wider head

Crucially, splitting into $h$ heads does not cost $h$ x the compute. Each head works in a subspace of size $d_{k} = d_{model} / h$ , so the total width is unchanged and the arithmetic is about the same as one full-width head. You get several distinct views for the price of one - the concatenation lines back up to exactly $d_{model}$ .

Split, attend, concat, project

Every head gets its own learned projection matrices that squeeze the full-width vectors down into its private subspace. It runs ordinary scaled dot-product attention there. Then all the head outputs are laid end to end and multiplied by one shared output matrix that mixes them back together for the next layer.

The multi-head formula

With model width $d_{model}$ and $h$ heads, set $d_{k} = d_{v} = d_{model} / h$ . Each head projects $Q, K, V$ into its own subspace and attends there:

head_{i} = Attention (Q W_{i}^{Q}, K W_{i}^{K}, V W_{i}^{V}), W_{i}^{Q}, W_{i}^{K} \in R^{d_{model} \times d_{k}} .

Run all $h$ in parallel, concatenate, and mix with one output projection $W^{O}$ :

MultiHead (Q, K, V) = Concat (head_{1}, \dots, head_{h}) W^{O}, W^{O} \in R^{h d_{v} \times d_{model}} .

Because $h d_{v} = d_{model}$ , the concatenation is exactly the model width again. The original transformer used $d_{model} = 512$ , $h = 8$ , so each head lived in a $d_{k} = 64$ -dimensional subspace.

Meet the heads

Below is a fixed sentence and four hand-crafted heads, each with a distinct, deterministic attention pattern - a caricature of what real heads learn. Click any token to make it the query (the word doing the looking); each coloured arc is one head's attention out of that word. Toggle heads on and off at the bottom, or press Isolate to solo one head and watch its whole pattern light up across every token.

Head explorer click a token - toggle heads - isolate

Query: it · 4 heads

Query “it”: each coloured arc is one head's attention out of this token. Different heads reach different words - the model reads several relationships at once. Click any token to move the query; toggle heads below.

What the demo is showing

With every head on and the query set to it, four arcs fan out to four different words: the blue head steps one token back, the amber head jumps to the subject cat, the green head follows grammar to the verb, and the red head resolves the pronoun all the way back to cat. That is the whole point of multi-head attention in one picture: a single token, read four different ways at once. Isolate the red head and you see a specialist that stays silent everywhere except the one long coreference link - exactly the kind of division of labour real heads discover during training.

All heads at once: one batched matmul

Here is the part that makes transformers fast. The $h$ heads never run in a Python loop - they are stacked into one extra tensor dimension and computed as a single batched matrix multiply. Reshape $Q, K, V$ to (batch, heads, tokens, d_k) and every head's attention is one call, dispatched across the GPU in parallel:

def multi_head_attention(x, Wq, Wk, Wv, Wo, h):
    B, T, d = x.shape                       # batch, tokens, d_model
    dk = d // h                             # per-head subspace: d_model / h

    Q = (x @ Wq).view(B, T, h, dk).transpose(1, 2)   # (B, h, T, dk)
    K = (x @ Wk).view(B, T, h, dk).transpose(1, 2)
    V = (x @ Wv).view(B, T, h, dk).transpose(1, 2)

    scores = (Q @ K.transpose(-2, -1)) / dk ** 0.5  # (B, h, T, T)  one batched matmul, all heads
    A      = scores.softmax(dim=-1)               # each head gets its own distribution

    out = (A @ V).transpose(1, 2).reshape(B, T, d) # concat heads back to d_model
    return out @ Wo                            # final output projection W_O

Concat then mix - don't skip W^O

The concatenation alone just parks the heads' outputs next to each other in disjoint slices. It is the output projection $W^{O}$ that lets information flow between heads, blending "previous token" with "coreference" into a single vector the next layer can use. Drop $W^{O}$ and the heads stay siloed, unable to combine what they each found.

Check yourself

The original transformer used d_model = 512 with h = 8 heads. What is each head's key dimension d_k?

Why run several small attention heads instead of one large one?

Recall: what are the three moves of multi-head attention?

Nail down the split-attend-concat structure. Try to state it, then check.

Lock it in

Multi-head attention splits the model width into $h$ subspaces, runs a full attention operation in each, then concatenates and projects back - several views for the cost of one.
Each head owns its own learned $W^{Q}, W^{K}, W^{V}$ projections and produces its own sharp softmax distribution, free to specialise on a different relationship.
The output projection $W^{O}$ is what lets information flow between heads - without it the heads stay siloed.
All heads compute as one batched matmul - reshape to (batch, heads, tokens, d_k) and the GPU handles every head in parallel.

Primary source

Jay Alammar's The Illustrated Transformer has the canonical picture of eight heads being projected, run in parallel, then concatenated and squeezed through $W^{O}$ - worth studying alongside this demo. For the mathematics and history of attention variants, see Lilian Weng, "Attention? Attention!".

Ask your teacher

Here is the connection that makes multi-head attention so GPU-friendly: the $h$ heads are just an extra dimension on the tensor, so all of them run as a single batched matrix multiplication - the matmul of Lesson 13, dispatched in parallel over the whole batch and all heads at once. This is exactly what the RNNs transformers replaced could not do: an RNN has to process token 1 before token 2 before token 3, sequentially. Attention has no such chain, so a modern GPU chews through the entire sentence - and every head - in one shot. Curious how the causal mask stops a head from peeking at future tokens, or why real heads sometimes learn nearly identical patterns (and get pruned)? Ask and I'll build the follow-up.