Multi-Head Attention

One attention head is one way of matching queries to keys. Language needs several kinds of match at once: syntax, coreference (“it” pointing at “the invoice”), recency (the last tool result), copying (repeat the id from the JSON).

Multi-head attention splits the model width into h smaller heads, runs attention in parallel, concatenates the outputs, and mixes them with a linear layer.

If the model width is 8 and h = 2, each head works in 4-d. You do not get eight times the compute of one fat head. You get several cheap views. A single softmax is one distribution: it tends to form one or two sharp peaks. Multiple heads can specialize. They are not labeled “syntax” and “copy” in the file. We discover them with probes, when we bother.

You do not twiddle head count per request. Papers pick numbers that train well. Your job is still the sequence those heads look at.

A wrong picture

A wrong picture is: “more heads means the model is more careful.” Too many tiny heads: each subspace is too small to match well. Too few fat heads: you reintroduce the softmax bottleneck. Head count is an architecture pick, not a quality slider.

Another wrong picture is: “one head is one human-readable skill.” A head might copy numbers on Monday and attend to punctuation on Tuesday. Do not write product docs that say “head 7 is the SQL head.”

A third wrong picture is: “cross-attention is a different math.” Cross-attention is the same recipe with queries from one sequence and keys/values from another. Encoder-decoder translation used that. Most agent stacks you will meet later are decoder-only: there is no second encoder. You concatenate docs into the prompt. The heads attend inside one sequence.

Parallel views, then concat

Width d. Heads h. Each head uses d/h dimensions (in the simple split). For each head:

1. Take the slice (or a learned map into that slice).
2. Run the attention recipe: dots, scale, softmax, mix values.
3. Get one output list per token, of length d/h.

Concatenate the h lists back to length d. Multiply by an output matrix so the heads can mix. Residual add comes later, in the block lesson.

Grouped-query and multi-query attention share keys and values across heads to speed decoding and shrink the cache. Same idea, cheaper memory. You will feel that in the KV-cache lesson: cache size is layers × heads × sequence × key width. Sharing keys cuts that bill.

A tiny example in words

Two tokens, width 4, two heads of width 2.

Token 0: [1.0, 0.0, 0.0, 1.0] — id-like in slots 0–1, time-like in slots 2–3. Token 1: [0.9, 0.1, 1.0, 0.0] — similar id, different time.

Head 0 uses slots 0–1 (id). Head 1 uses slots 2–3 (time). Head 0 should agree the two tokens match. Head 1 should disagree. One concatenated fat head forced to average those stories would blur both.

Two heads, two subspaces

Lists of numbers. Print softmax rows per head. No extra libraries beyond math.

Run to execute this in your browser. Nothing is sent to a server.

Head 0 agrees the two tokens match (id-like dims): off-diagonal weights stay healthy. Head 1 disagrees (time-like dims): each token prefers itself more. That is the cartoon of specialization.

If you averaged the two subspaces into one 4-d softmax, the id match and the time clash would fight inside one distribution. Multi-head lets both stories survive until the concat mix.

Change token 1’s last two numbers to match token 0. Head 1 should start to agree too. That is how you debug a toy: change one subspace, watch one head.

What you cannot see from outside

You cannot set head 3 to “always copy the job id.” You can:

• Put the job id in a short, unique span so some head can match it.
• Avoid packing five jobs into one token (“remember the id, the policy, the tone, the schema, and the joke”).
• Spread the job across sentences: a short spec, a short schema, a short observation.

Heads can attend to different sentences. They cannot invent a missing spec. They cannot split a single overloaded sentence into five clean views if you never wrote the five views.

Sharing keys across heads (grouped-query) is a speed/memory move at decode time. It is not a prompt trick. Do not try to “turn on grouped-query” from a chat box.

How agents use this

Do not design prompts that require one token to mean five things. Spread the job across the sequence. A short spec. A short schema. A short observation. Heads can attend to different sentences. They cannot invent a missing spec.

When a tool argument must copy an id and also obey a policy, put those as two spans, not as one tangled clause. You are feeding different keys to different heads.

Decoder-only stacks do not have a hidden encoder “document memory.” Retrieved chunks are more tokens in the same self-attention. They compete for softmax mass with the spec and the history. Multi-head helps that competition a bit. It does not remove it.

• Spread: one idea per short span when you can.
• Do not label heads: you will be wrong next week.
• Cross vs self: agent prompts are usually one concatenated sequence.
• Cache: more heads can mean a bigger KV cache unless keys are shared.
• Missing spec: no head can attend to a sentence you deleted.

Note:Cross-attention is the same recipe with queries from one sequence and keys from another. Most agent stacks hide a decoder-only model: you concatenate docs into the prompt.