Multi-head attention

Many attentions in parallel

Sources: Attention Is All You Need — Vaswani et al., 2017; GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints — Ainslie et al., 2023

Imagine you had to summarize a paragraph and you were only allowed to track one type of relationship: say, “which noun is the subject of which verb.” You’d lose everything about tense, modifiers, references, sentiment. Real reading comprehension requires tracking many relationships at once. The model has the same problem: one attention head uses one set of projections and one learned subspace, even though its weights can vary across positions and inputs. The fix is to run many heads side-by-side, each free to learn its own patterns. That’s multi-head attention.

Splitting d_model into heads

Each attention computation is called a headheadOne independent attention computation. Multi-head splits d_model into N parallel heads, each learning its own pattern.See in glossary →. If $d_{\text{model}} = 4096$ and we want 32 heads, we split into 32 chunks of 128. Each head gets its own slice of the query, key, and value vectors, and runs the attention formula independently. The number 128 here is called the head dimensionhead_dimd_model / num_heads — the dimension each attention head operates in.See in glossary → ($d_k$ or $d_{\text{head}}$).

Concretely:

num_heads = 32
d_model   = 4096
head_dim  = d_model / num_heads = 128

Each head produces an output of size 128. The 32 outputs are concatenated back to a single vector of size 4096, then passed through one more learned projection ($W_O$) to mix the heads’ outputs.

In code (schematic):

def multi_head_attention(x):
    # x: (seq, d_model)
    q = (x @ W_Q).reshape(seq, num_heads, head_dim)
    k = (x @ W_K).reshape(seq, num_heads, head_dim)
    v = (x @ W_V).reshape(seq, num_heads, head_dim)

    # heads run in parallel — same formula, just over the last axis per head
    scores = einsum("shd, thd -> sht", q, k) / sqrt(head_dim)
    scores = scores.masked_fill(causal_mask, -inf)
    weights = softmax(scores, dim=-1)
    out = einsum("sht, thd -> shd", weights, v)

    # concat heads, project
    out = out.reshape(seq, d_model)
    return out @ W_O

What different heads actually do

After training, researchers have observed heads whose attention patterns often correlate with different relationships. These patterns are useful clues, not a clean decomposition of every model into independently interpretable parts. Examples include:

• Previous-token heads: attend to position $i-1$. Useful for tracking local syntax.
• Same-word-class heads: nouns attend to other nouns; verbs to other verbs.
• Coreference heads: pronouns (“he”, “she”, “it”) attend to the noun they refer to.
• Punctuation heads: track sentence and clause boundaries.
• Induction heads: implement the pattern “if A B has appeared earlier, and we just saw A again, attend to the position right after the earlier A.” They are associated with some forms of in-context pattern completion.
• Attention sinks: heads that concentrate probability on token 0 or another early token. This pattern has been observed in trained models, but attention weights alone do not establish why a head uses it.

You can switch between a few of these patterns in the heatmap from the previous section. None of this behavior is hand-coded; it can emerge from training on next-token prediction.

How big does this get?

Take Llama-3-8B: 32 layers, 32 query heads per layer, head_dim = 128, and 8 key/value heads. Because it uses grouped-query attention, the Q and O matrices are each 4096×4096 = 16.8M parameters, while the K and V matrices are each 4096×(8×128) = 4.2M parameters. That is about 41.9M attention-projection parameters per layer, or 1.34B across 32 layers. The MLP blocks (next section) take most of the rest.

We’ve only described attention here in its “vanilla” form, where every head has its own query, key, and value projection. Modern models use a memory-saving variant called grouped-query attention (GQA) that has many fewer KV heads than Q heads, sharing keys and values across groups of query heads. We’ll come back to that in section 12: it substantially reduces KV-cache size and is especially useful for long contexts.

Attention is the “how do tokens look at each other” part of a transformer. But there’s something we glossed over at the start of section 4: the math you’ve just seen treats its inputs as a bag of vectors. Permute the input tokens and you get a permuted output: the cat and the mat are interchangeable. Real language clearly cares about order, so we need to tell the model where each token sits in the sequence. That’s the topic of the next section: positional encoding.