Multi-Head Attention Mechanism

A panel of specialists, not one generalist

Watch a cricket broadcast during a review of one delivery. The former bowler talks about seam position and length. The batting coach talks about footwork and bat swing. The fielding analyst talks about where the ball eventually went. Three specialists, watching the exact same six seconds of video, each extracting a different relationship between the same set of events, and none of them is wrong. A single commentator forced to cover all three angles at once would end up giving a blended, watered-down account of each. The broadcast works because the panel runs in parallel, and a producer stitches their separate notes into one segment afterward.

That is the exact problem multi-head attention solves inside a transformer layer, and it is worth being precise about why one attention head cannot just do the job of several. Take the sentence "Virat hit the ball because it was full and outside off." The pronoun "it" sits at the center of two separate relationships at once: it corefers with "ball" (they are the same object), and it is the subject of the modifier "full and outside off" (a property being asserted about the delivery). A single attention head, as you have already studied, produces exactly one probability distribution over the sequence for a given query, computed once via one softmax. If that one distribution has to serve both the coreference link to "ball" and the modifier link to "full", it can only strike some compromise between the two, spreading weight across both without representing either relationship cleanly. Multi-head attention removes the compromise by computing several independent attention distributions in parallel, each in its own learned subspace of the embedding, and only combining them after each has been allowed to specialize. One head can lean toward tracking coreference; another can lean toward tracking syntactic modification; a third might track word order; a fourth might track something with no clean linguistic name at all. None of this specialization is hand-programmed. It emerges from gradient descent because the architecture gives each head room to be different.

Recap: scaled dot-product attention in one head

You have already seen the core operation. Given a matrix of token embeddings X (shape T x d_model, one row per token), a single attention head learns three projection matrices W^Q, W^K, W^V, each of shape d_model x d_k, and computes:

Q = X W^Q      K = X W^K      V = X W^V
Attention(Q, K, V) = softmax( Q K^T / sqrt(d_k) ) V

Each row of the resulting matrix is a new representation of one token, built as a weighted average of every token's value vector, where the weights come from how strongly that token's query matches every other token's key. The division by sqrt(d_k) keeps the dot products from growing so large that softmax saturates into a near one-hot distribution as d_k increases; you can verify this yourself by noting that if Q and K entries are independent with unit variance, a dot product over d_k terms has variance proportional to d_k, so scaling by sqrt(d_k) brings the variance back to roughly 1 regardless of dimension.

Multi-head attention does not replace this operation. It runs it h times, each time with its own small W^Q, W^K, W^V, and then recombines the results.

Formal definition

For h heads, define per-head projections W_i^Q, W_i^K of shape d_model x d_k and W_i^V of shape d_model x d_v, for i = 1 .. h. Each head computes its own attention independently:

head_i = Attention(X W_i^Q, X W_i^K, X W_i^V)   for i = 1 ... h
MultiHead(X) = Concat(head_1, ..., head_h) W^O

W^O has shape (h · d_v) x d_model, and its job is exactly the producer's job in the broadcast booth: concatenation alone just places the h specialist reports side by side in one long vector, with no interaction between them. W^O is the learned linear pass that lets the model blend what head 3 noticed with what head 5 noticed into a single, coherent output vector per token.

The standard convention, used in the original transformer and in every major LLM since, is d_k = d_v = d_model / h, so h must divide d_model evenly. This convention has a consequence worth deriving explicitly, because it corrects a natural assumption that multi-head attention is more expensive than single-head attention. Count the parameters in the three input projections across all h heads: each head contributes d_model x d_k parameters to W^Q, the same to W^K, and d_model x d_v to W^V. Summed over h heads with d_k = d_v = d_model/h:

h * (d_model * d_k) = h * d_model * (d_model / h) = d_model^2

That total is independent of h. Three such projections (Q, K, V) cost 3 · d_model² combined, and W^O costs one more d_model², for 4 · d_model² total, exactly what a single "head" using the full d_model as its own d_k would cost. Splitting one big head into h narrow heads does not add parameters or change the asymptotic compute of the score matrix (still O(T² · d_model) total across all heads, since h heads each do O(T² · d_model/h) work). What changes is the structure of the representation: instead of one T x d_model attention pattern, you get h independent T x T attention patterns, each free to organize itself around a different kind of relationship, at no extra parameter cost.

The multi-head attention pipeline

Worked example: two heads disambiguating a pronoun

Take the mini-sequence ball, full, it from the sentence discussed above, and give each token a toy 4-dimensional embedding (illustrative numbers chosen to make the arithmetic checkable by hand, not values from a trained model). Say the four dimensions loosely correspond to "entity-ness", "adjective-ness", "number/verb-ness", and a constant bias term:

x_ball = [3, 0, 0, 1]
x_full = [0, 3, 0, 1]
x_it   = [1, 1, 0, 1]

We will compute the multi-head attention output for the query token it, attending over all three tokens, using h = 2 heads with d_k = d_v = 2. To keep the projections easy to check by hand, Head 1's W^Q and W^K are built to read off only the "entity-ness" and "bias" coordinates (columns 1 and 4), while Head 2's read off only "adjective-ness" and "bias" (columns 2 and 4). In a real network these matrices are learned by backpropagation and never look this clean, but the toy version lets you trace every multiplication and see exactly why the two heads land on opposite answers.

import numpy as np

x_ball = np.array([3,0,0,1], dtype=float)
x_full = np.array([0,3,0,1], dtype=float)
x_it   = np.array([1,1,0,1], dtype=float)
X = np.stack([x_ball, x_full, x_it])   # rows: ball, full, it
d_k = 2

# Head 1 reads columns (entity, bias)
W_Q1 = np.array([[1,0],[0,0],[0,0],[0,1]], dtype=float)
W_K1 = np.array([[1,0],[0,0],[0,0],[0,1]], dtype=float)
# Head 2 reads columns (adjective, bias)
W_Q2 = np.array([[0,0],[1,0],[0,0],[0,1]], dtype=float)
W_K2 = np.array([[0,0],[1,0],[0,0],[0,1]], dtype=float)
# both heads share one value projection here, purely to shorten the arithmetic
W_V  = np.array([[1/3,0],[0,1/3],[0,0],[0,0]], dtype=float)

def head(Wq, Wk, Wv):
    q = x_it @ Wq                       # query for "it" only
    K = X @ Wk                          # keys for ball, full, it
    V = X @ Wv                          # values for ball, full, it
    scores = (K @ q) / np.sqrt(d_k)
    w = np.exp(scores - scores.max())
    w = w / w.sum()
    return q, K, V, scores, w, w @ V

q1,K1,V1,s1,w1,o1 = head(W_Q1, W_K1, W_V)
q2,K2,V2,s2,w2,o2 = head(W_Q2, W_K2, W_V)
z = np.concatenate([o1, o2])            # W^O = identity for this example

Running this by hand, column by column: q1 = x_it @ W_Q1 = [1, 1], and the three keys under Head 1's projection are K1 = [[3,1], [0,1], [1,1]] for ball, full, it respectively. The raw dot products K1 @ q1 are 3·1+1·1=4 for ball, 0·1+1·1=1 for full, and 1·1+1·1=2 for it. Dividing by √2 ≈ 1.4142 gives scaled scores [2.8284, 0.7071, 1.4142]. Feeding these through softmax (exponentiate, then normalize by the sum 16.9188+2.0281+4.1132 = 23.0602) gives attention weights:

Head 1 weights on (ball, full, it) = [0.7337, 0.0879, 0.1784]

Head 1 puts nearly three-quarters of its attention mass on ball: this is the coreference-leaning head. With value vectors V1 = [[1,0],[0,1],[0.3333,0.3333]], the weighted sum is 0.7337·[1,0] + 0.0879·[0,1] + 0.1784·[0.3333,0.3333] = [0.7931, 0.1474].

Head 2 is built the same way but reading the adjective column instead of the entity column, which swaps which token looks like a match: the raw dot products become 1 for ball, 4 for full, 2 for it, exactly mirroring Head 1 with ball and full exchanged. Its softmax weights come out to [0.0879, 0.7337, 0.1784], putting three-quarters of the attention on full instead: this is the modifier-leaning head. Its output is [0.1474, 0.7931], the mirror image of Head 1's.

Concatenating the two 2-dimensional head outputs gives a 4-dimensional vector [0.7931, 0.1474, 0.1474, 0.7931] for the token it. With W^O taken as the identity for simplicity, that concatenation is the final multi-head output. Notice what happened: the first two coordinates (from Head 1) encode "this token is strongly linked to the entity ball", and the last two (from Head 2) encode "this token is strongly linked to the modifier full", both carried simultaneously in one output vector. No single softmax distribution could have represented both facts about it at once without blending them into a muddier compromise; two heads represented both cleanly, in parallel, at the parameter cost derived above.

Common misconception: "more heads means each head attends to fewer tokens"

A very natural but wrong mental model treats h heads as if they divided the sequence among themselves, so that with h = 8 heads each head only looks at roughly one-eighth of the tokens, the way splitting a task among 8 people means each person does one-eighth of the work. Attention does not work that way, and the worked example above is direct evidence against it. Look again at Head 1's softmax weights: [0.7337, 0.0879, 0.1784]. Every one of the three tokens receives strictly positive weight. Head 1 does not "skip" full or it; it computes a full T-length probability distribution over every token in the sequence, exactly as a single-head attention layer would, and so does Head 2. What differs between the heads is not how many tokens each one looks at (both look at all of them, always), but which relationships the projection matrices make that distribution sensitive to. Splitting d_model into h subspaces of size d_model/h splits the representation dimension, not the sequence dimension. Each head's attention matrix is still T x T, computed over the full sequence; you get h such T x T matrices, one per head, each potentially shaped differently, not one T x T matrix chopped into h pieces. This is precisely why multi-head attention is useful for exactly the kind of sentence used above: both the coreference relationship and the modifier relationship needed access to the entire three-token context, and both heads got it, just weighted differently.

Choosing the number of heads

h is a hyperparameter fixed before training, and it trades off two things that pull in opposite directions. More heads means more independent subspaces available for the model to specialize into, which is good up to a point. But because d_k = d_model/h shrinks as h grows, each individual head's query and key vectors get lower-dimensional, and a query-key dot product computed in a very low-dimensional space has less room to encode a rich matching criterion before it starts colliding with unrelated signals. The original transformer used d_model = 512 with h = 8, giving d_k = 64 per head. Later large language models scale both together: d_model = 12288 with h = 96 (again d_k = 128) is representative of GPT-3-scale configurations. In practice h is chosen empirically alongside d_model and depth, but the constraint that d_model must divide evenly by h, together with the parameter-count invariance shown earlier, means the choice is really about how finely to partition a fixed representation budget, not about buying extra capacity by adding heads.

Active recall

Attempt each question before reading its answer.

1. In the worked example, if W^O were NOT the identity matrix, would the final
   output for "it" still contain the information from both heads separately?

2. Head 1 and Head 2 in the worked example use different W_Q and W_K but the
   SAME W_V. Does sharing W_V break the definition of multi-head attention,
   or is it just a simplification made for this example?

3. Suppose d_model = 512 and h = 8. What is d_k? If h were changed to 16
   with d_model unchanged, what happens to the total parameter count in the
   Q/K/V projections, and what happens to d_k?

4. True or false: because Head 1's softmax weight on "full" (0.0879) is much
   smaller than its weight on "ball" (0.7337), Head 1 does not use any
   information from "full" at all.

5. Why is the division by sqrt(d_k) applied inside each head separately,
   rather than once after concatenating all h heads' scores together?

6. A classmate says "multi-head attention is just single-head attention run
   8 times to get 8 different answers, and then you average them." Identify
   the two things wrong with this description.

Answers

1. Yes, and this is the important case, not the exceptional one. Concat still
   places head_1's [0.7931, 0.1474] and head_2's [0.1474, 0.7931] side by
   side first; W^O then takes that 4-vector and applies a learned linear map
   to it. As long as W^O is a genuine (invertible-ish, or at least
   non-degenerate) linear map, information from both halves of the
   concatenation survives in the output, just recombined. It is only if W^O
   were pathological (for example, all zero in the columns reading Head 2's
   half) that Head 2's contribution would be discarded, and training would
   have no reason to learn that unless Head 2 were truly useless.

2. It is a simplification, not a structural requirement. The formal
   definition gives every head its own W_i^V, so in a real trained model
   Head 1 and Head 2 would learn different value projections too, meaning
   the "content" each head retrieves, not just the weighting pattern, would
   differ. Sharing W_V here only kept the arithmetic in the worked example
   short; it does not change the point being illustrated, which is that the
   two heads produce different attention WEIGHTS over the same tokens.

3. d_k = 512 / 8 = 64. Moving to h = 16 with d_model still 512 makes
   d_k = 512 / 16 = 32, but the total parameter count in the Q, K, or V
   projections stays h * d_model * d_k = 16 * 512 * 32 = 262144, identical
   to 8 * 512 * 64 = 262144. Only the split changes, not the budget.

4. False. Softmax weights are never exactly zero (only asymptotically as a
   score dominates), so 0.0879 still contributes 0.0879 * v_full to Head 1's
   output. It is small relative to Head 1's weight on "ball", which is
   exactly why we call Head 1 "coreference-leaning" rather than
   "coreference-exclusive". Head 1 uses information from every token, just
   unevenly.

5. Because each head has its own d_k and its own scale of raw dot products,
   which depends on that head's own W_Q and W_K (different heads can even
   use different d_k in principle). Scaling must happen before each head's
   own softmax, using that head's own d_k, so that head's particular
   distribution is well-behaved. Scaling once after concatenation would mix
   scores from different heads that were never compared against each other
   in the first place, since each head's softmax normalizes only over its
   own T scores, not over a combined pool.

6. First, the h heads are not run on the same projection and then averaged;
   each head has its OWN learned W_Q, W_K, W_V, so they see different
   projections of X before any attention is computed, not just different
   random outcomes of an identical process. Second, the outputs are not
   averaged at the end; they are concatenated into a longer vector and then
   passed through one learned linear map W^O, which lets the model combine
   the heads with arbitrary learned weights rather than a fixed 1/h average.

Think About It

Think about this: How would you explain multi-head attention mechanism to a friend who has never seen a computer? What real-world analogy would you use? Imagine you had to build a system using these concepts — what would be your first step? Try this: before moving on, write down three things you learned and one question you still have.