Take a sentence a cricket commentary bot has to process: "Kohli, who anchored the innings after Rohit fell early, finally reached his century in the 47th over." The pronoun "his" needs to resolve to "Kohli" — that is a coreference relationship, spanning eleven tokens. The word "century" needs to connect to "47th over" — that is a topical, numerically-adjacent relationship, spanning four tokens. Both relationships live inside the same sentence, need to be tracked from overlapping query positions, and are structurally nothing alike: one is about matching an entity mention to a pronoun, the other is about matching a scoring milestone to a delivery count. A single scaled dot-product attention head, as you have already derived, produces exactly one probability distribution per query token — one softmax, one weighted average over the sequence. If "his" needs to attend strongly to "Kohli" for coreference and, separately, a nearby position needs to attend strongly to "47th over" for topical grounding, a single head has to represent both jobs through one shared similarity function Q·Kᵀ. This chapter works out, in full arithmetic, why that is a genuine bottleneck, what splitting the projection into multiple heads actually buys you, and — a result most treatments skip — what it costs in parameters and compute, which turns out to be surprisingly close to nothing.
The problem one attention head can't solve
Recall the single-head mechanism only in outline: given a query vector q and a set of key vectors {k_j}, attention computes scores q·k_j / √d_k, turns them into a probability distribution via softmax, and returns a weighted sum of the value vectors {v_j} using that one distribution. The critical structural fact for this chapter is that the softmax has exactly one set of weights per query per layer per head. If the "ideal" attention pattern for resolving coreference (mostly weight on "Kohli") is different from the ideal pattern for resolving topical adjacency (mostly weight on "over"), a single head is forced into a compromise: it produces one distribution that partially serves both goals and fully serves neither. This is not a vague intuition — it is a direct consequence of softmax being a function from one score vector to one probability vector. You cannot get two different distributions out of one softmax call over the same scores.
Multi-head attention's answer is not "make the head bigger" — a bigger single head still produces one distribution. The answer is to run several independent attention computations in parallel, each in its own lower-dimensional subspace, so each subspace can settle on its own similarity function and therefore its own distribution, and only merge the results afterward. That "merge afterward" step is where the real information — both the coreference signal and the topical signal — survives, undiluted, into the next layer.
The mechanism, formally
Let d_model be the model's embedding width and h the number of heads, with d_k = d_v = d_model / h (the standard convention, though d_k and d_v need not be equal in principle). For each head i = 1 … h, the model learns its own projection matrices W_Q^i, W_K^i ∈ ℝ^{d_model × d_k} and W_V^i ∈ ℝ^{d_model × d_v}, distinct from every other head's. Given the layer's input X ∈ ℝ^{n × d_model} (n tokens), head i computes:
Q_i = X W_Q^i K_i = X W_K^i V_i = X W_V^i
head_i = softmax( Q_i K_iᵀ / √d_k ) V_i ∈ ℝ^{n × d_v}Every head runs this in full independence — different weights, different resulting Q/K/V, different softmax distribution. The h resulting matrices, each n × d_v, are concatenated along the feature axis into a single n × (h·d_v) matrix, and then passed through one more learned matrix W_O ∈ ℝ^{(h·d_v) × d_model} that mixes information across heads back into the model's working dimension:
MultiHead(X) = Concat(head_1, …, head_h) · W_OTwo things to notice before the worked example. First, because d_k = d_model / h, each head operates on a strictly smaller vector than the full embedding — it physically cannot see the whole representation, only its own slice of the learned subspace. Second, W_O is the only place where the heads' outputs interact; everything before concatenation is fully parallel and fully independent.
Worked example: one softmax versus two, on the cricket sentence
To make the compromise argument numeric rather than hand-wavy, set up a minimal but complete instance: d_model = 4, three tokens — "Kohli", "his", "century" — with query token "his" attending over all three (including itself, as in standard self-attention). Assume the projections W_Q, W_K, W_V have already been applied (their construction was the sibling chapter's job), producing these post-projection vectors. Dimensions 1–2 are meant to carry coreference-relevant features (name-ness, pronoun-ness); dimensions 3–4 carry topic/position features (cricket-term-ness, sequence position):
Q(his) = [1, 2, 0, 1]
K(Kohli) = [2, 1, 0, 0] V(Kohli) = [1, 0, 0, 0]
K(his) = [0, 1, 0, 2] V(his) = [0, 1, 0, 0]
K(century) = [0, 0, 2, 1] V(century) = [0, 0, 1, 0]Single full-width head (d_k = 4). Raw dot products: Q(his)·K(Kohli) = 1·2+2·1+0·0+1·0 = 4; Q(his)·K(his) = 0+2+0+2 = 4; Q(his)·K(century) = 0+0+0+1 = 1. Scale by √d_k = √4 = 2: scores = [2, 2, 0.5]. Softmax: e² = 7.389 (twice), e^0.5 = 1.649, sum = 16.427, giving weights [0.450, 0.450, 0.100]. The output is the weighted sum of the V vectors:
out_full = 0.450·[1,0,0,0] + 0.450·[0,1,0,0] + 0.100·[0,0,1,0]
= [0.450, 0.450, 0.100, 0]Read this distribution: 45% to "Kohli" (the correct coreference antecedent), but an equal 45% to itself and only 10% to "century". The full-width dot product mixed the dims-1–2 coreference signal and the dims-3–4 topic signal into a single number per key, and the result is a tie between "attend to Kohli" and "attend to self" — a genuinely ambiguous, compromised distribution, not a clean resolution of either relationship.
Two heads, d_k = 2 each. Split every vector above at the dimension boundary: head A keeps dims 1–2, head B keeps dims 3–4 (a legitimate special case of a learned linear projection — a 2×4 selection matrix is still a linear map, though in a trained network each head's W_Q^i, W_K^i, W_V^i would be learned freely rather than fixed as a slice).
Head A (dims 1–2): Q_A = [1,2], keys [2,1], [0,1], [0,0]. Dot products: 4, 2, 0. Scale by √2 = 1.414: scores [2.828, 1.414, 0]. Softmax: e^2.828=16.92, e^1.414=4.113, e^0=1, sum =22.03, weights [0.768, 0.187, 0.045]. With V_A = [1,0], [0,1], [0,0]:
out_A = 0.768·[1,0] + 0.187·[0,1] + 0.045·[0,0] = [0.768, 0.187]Head B (dims 3–4): Q_B = [0,1], keys [0,0], [0,2], [2,1]. Dot products: 0, 2, 1. Scale by √2: scores [0, 1.414, 0.707]. Softmax: e^0=1, e^1.414=4.113, e^0.707=2.028, sum =7.141, weights [0.140, 0.576, 0.284]. With V_B = [0,0], [0,0], [1,0]:
out_B = 0.140·[0,0] + 0.576·[0,0] + 0.284·[1,0] = [0.284, 0]Concatenating: [0.768, 0.187, 0.284, 0] — same total width as out_full, but qualitatively different. Head A committed to a sharp, largely resolved coreference judgment: 76.8% to "Kohli" versus 45.0% in the single-head case, with the tie to "self" broken decisively. Head B, working in a subspace the coreference computation cannot see, independently registered a topical pull toward "century" (28.4%) — a signal the single head only weakly expressed and had to trade off against the coreference signal to do so. Splitting did not add information that wasn't there; it removed the forced interference between two unrelated similarity computations that shared one dot product.
Verifying the arithmetic in code
The hand computation above is fully reproducible; nothing in it depends on hidden state or an "assumed helper" — every array is defined before use.
import numpy as np
def softmax(x):
e = np.exp(x - np.max(x))
return e / e.sum()
Q_his = np.array([1, 2, 0, 1], dtype=float)
K = {
"Kohli": np.array([2, 1, 0, 0], dtype=float),
"his": np.array([0, 1, 0, 2], dtype=float),
"century": np.array([0, 0, 2, 1], dtype=float),
}
V = {
"Kohli": np.array([1, 0, 0, 0], dtype=float),
"his": np.array([0, 1, 0, 0], dtype=float),
"century": np.array([0, 0, 1, 0], dtype=float),
}
keys = ["Kohli", "his", "century"]
# single full-width head, d_k = 4
scores_full = np.array([Q_his @ K[k] for k in keys]) / np.sqrt(4)
w_full = softmax(scores_full)
out_full = sum(w * V[k] for w, k in zip(w_full, keys))
print("full weights:", np.round(w_full, 3))
print("full output :", np.round(out_full, 3))
# two heads, d_k = 2 each, split on the dimension boundary
def head_output(dims):
q = Q_his[dims]
scores = np.array([K[k][dims] @ q for k in keys]) / np.sqrt(len(dims))
w = softmax(scores)
out = sum(wt * V[k][dims] for wt, k in zip(w, keys))
return w, out
w_A, out_A = head_output([0, 1])
w_B, out_B = head_output([2, 3])
print("head A weights:", np.round(w_A, 3), "output:", np.round(out_A, 3))
print("head B weights:", np.round(w_B, 3), "output:", np.round(out_B, 3))
print("concatenated :", np.round(np.concatenate([out_A, out_B]), 3))Running this prints values matching the hand derivation exactly: full weights ≈ [0.450, 0.450, 0.100], full output ≈ [0.450, 0.450, 0.100, 0.000]; head A weights ≈ [0.768, 0.187, 0.045], output ≈ [0.768, 0.187]; head B weights ≈ [0.140, 0.576, 0.284], output ≈ [0.284, 0.000]; concatenated ≈ [0.768, 0.187, 0.284, 0.000]. In a full model this 4-vector would next pass through the learned output projection W_O, which mixes the two heads' contributions into the residual stream — but note that the mixing happens strictly after each head has independently committed to its own distribution; W_O cannot undo the fact that two different softmaxes ran.
The parameter-count and FLOPs accounting
Here is the number that surprises most students: splitting into h heads does not increase the parameter budget or the compute cost of the attention layer, provided d_k = d_model / h. Work it out both ways.
Single head, full width (d_k = d_v = d_model = d): W_Q, W_K, W_V are each d × d, and there is no separate output projection needed since the head's output is already d-wide (though a real implementation would still often include one). Parameter count for the three input projections: 3d².
Multi-head, h heads (d_k = d_v = d/h): each head's W_Q^i, W_K^i, W_V^i is d × (d/h), so one head costs 3 · d · (d/h) parameters. Summed over h independent heads: h · 3d(d/h) = 3d² — identical to the single-head total, because the h in the numerator and denominator cancel. The concatenated output is h · (d/h) = d wide, so the output projection W_O is d × d, costing d². Grand total: 3d² + d² = 4d², exactly matching a single-head layer that also includes an output projection.
The FLOPs story is the same shape. Computing the n × n score matrix for one head costs n² · d_k multiply-adds (each of n² entries is a d_k-length dot product). Across h heads: h · n² · (d/h) = n² · d — again exactly the cost of one full-width head computing scores over the whole d-dimensional vectors. Take a concrete instance: n = 100 tokens, d = 512. Single head: 100² · 512 = 5,120,000 multiply-adds for the score matrix. Eight heads of d_k = 64: 8 · 100² · 64 = 8 · 640,000 = 5,120,000 — identical, to the last unit.
So multi-head attention is, in a real sense, "free": for the same parameter budget and the same FLOPs, you trade one full-rank similarity computation for h independent lower-rank ones. The gain demonstrated in the worked example — two clean, specialized distributions instead of one compromised distribution — comes at essentially no additional cost. This is also explicitly the design rationale given in Vaswani et al. (2017), "Attention Is All You Need," which introduced the mechanism with d_model = 512, h = 8, d_k = d_v = 64 — note 8 × 64 = 512, exactly the parameter-preserving split worked out above.
Why the scaling factor doesn't need re-deriving per head
You already know that 1/√d_k keeps the variance of the raw dot product controlled — but does that derivation need to be redone separately for every head, now that each head's d_k is smaller than d_model? Assume each component of q and k is drawn independently with mean 0 and variance σ². The dot product q·k = Σ_{j=1}^{d_k} q_j k_j is a sum of d_k independent, mean-zero terms, each with variance σ⁴ (the variance of a product of two independent mean-zero variables). By additivity of variance for independent terms, Var(q·k) = d_k · σ⁴ — it scales linearly with whatever d_k happens to be, full-width or split. Dividing by √d_k before the softmax gives a scaled score with variance d_k σ⁴ / d_k = σ⁴, which does not depend on d_k at all. This is exactly why multi-head attention can use a smaller d_k per head without retuning the scaling rule: the 1/√d_k factor is self-correcting, automatically re-normalizing the softmax input variance to the same constant σ⁴ regardless of how many pieces d_model is chopped into. Change h, and d_k changes, and the scale factor changes with it in exactly the compensating way — no separate derivation needed per head count.
What heads actually specialize in: evidence, not folklore
The claim that different heads "specialize" is often stated but rarely grounded. Two studies give it real empirical footing. Clark et al. (2019), "What Does BERT Look At? An Analysis of BERT's Attention" (Kevin Clark, Urvashi Khandelwal, Omer Levy, Christopher D. Manning), examined attention patterns in a trained BERT model and found individual heads that consistently attend to specific syntactic relations — for example, heads that reliably attend from a verb to its direct object, or from a noun to a determiner, well above what a positional or random baseline would produce — alongside heads whose behavior is close to purely positional (attending to a fixed offset such as the previous or next token) and heads that behave close to a no-op, dumping most of their attention mass onto a fixed low-information token. Voita et al. (2019), "Analyzing Multi-Head Self-Attention: Specialized Heads Do the Heavy Lifting, the Rest Can Be Pruned" (Elena Voita, David Talbot, Fedor Moiseev, Rico Sennrich, Ivan Titov), went further and showed that in trained Transformer translation models, a small number of heads per layer account for most of the useful behavior — including identifiable positional heads and syntactic heads — while a substantial fraction of the remaining heads can be pruned after training with only a small drop in translation quality, because their function has become redundant with other heads or negligible.
Put together with Michel et al. (2019), "Are Sixteen Heads Really Better than One?" (Paul Michel, Omer Levy, Graham Neubig), which independently confirmed that many trained heads can be removed at inference time with minimal accuracy loss, the picture is consistent: the architecture provides h independent slots capable of specializing (as the worked example shows mechanically), but training does not guarantee every slot ends up doing distinct, useful work. Some heads converge to genuinely different functions — positional tracking, specific syntactic dependencies, rare-token handling — while others converge to overlapping or near-trivial functions. The multi-head mechanism creates the opportunity for specialization; gradient descent decides, empirically and per model, how much of that opportunity gets used.
The misconception to retire
A common misconception: multi-head attention must use more parameters than a single head of the same model width, because "you're running eight attention computations instead of one." The parameter-count derivation above shows this is false whenever d_k = d_model / h: the total is 4d² either way, with or without the split. What changes between the two is not capacity but structure — one full-rank d × d bilinear similarity form becomes h independent rank-(d/h) forms occupying disjoint subspaces. That structural change is precisely what let head A and head B, in the worked example, each converge on a different, sharper distribution instead of forcing one softmax to average two unrelated signals. The benefit of multi-head attention is a representational one bought by reorganizing an unchanged parameter budget, not a capacity increase bought by adding parameters.
Putting it together
Active recall
Attempt each question before reading its answer.
1. In the worked example, if d_model stayed 4 but the model used h = 4 heads instead of 2, what would d_k be, and how would the softmax scaling factor change?
2. BERT-base uses d_model = 768 and h = 12. Compute d_k, and verify the total parameter count of the Q/K/V/O projections equals 4d². What is that number?
3. Starting from the worked example (his/Kohli/century, d_model=4, h=2), suppose V(century) is changed from [0,0,1,0] to [0,0,3,0], tripling its magnitude, while every Q and K stays the same. Trace the full effect: do the attention weights change in either head or the full-head case? Does head A's output change? Head B's? The concatenation? The full-head output?
4. Why does the 1/√d_k scaling factor not need to be re-derived separately for every choice of h?
5. Voita et al. (2019) found that many trained attention heads can be pruned after training with little performance loss. Does this contradict the earlier result that multi-head attention has the same total parameter count as a single full-width head?
6. For n = 100 tokens and d_model = 512, confirm that the FLOPs for computing the score matrix are the same whether you use 1 head of d_k=512 or 8 heads of d_k=64.
Answer 1. d_k = d_model / h = 4/4 = 1. The scale factor becomes √1 = 1, so scores are no longer divided at all — softmax operates directly on the raw single-dimensional dot products. With only one dimension per head, each head's similarity function is a single scalar multiplication per key, which is very low expressiveness: a head this narrow can encode at most a one-dimensional notion of similarity per subspace. This mirrors the real finding in Michel et al. (2019) that adding heads is not free of representational tradeoffs once d_k gets very small — the parameter/FLOPs equivalence proven earlier holds regardless of h, but the quality of what an individual head can represent degrades as its subspace shrinks toward one dimension.
Answer 2. d_k = 768/12 = 64. Total parameters = 4d² = 4 × 768² = 4 × 589{,}824 = 2{,}359{,}296 — about 2.36 million parameters for one multi-head self-attention block's four projection matrices (excluding biases and the feed-forward sublayer), whether that width is split into 12 heads of 64 or left as one head of 768.
Answer 3. Attention weights depend only on Q and K, never on V — so in every case (full head, head A, head B) the softmax weights are completely unchanged: [0.450, 0.450, 0.100] for the full head, [0.768, 0.187, 0.045] for head A, [0.140, 0.576, 0.284] for head B. Head A's value slice for "century" was [0,0] (dims 1–2 of the original [0,0,1,0], both zero) and stays [0,0] after the change — head A's output is completely unaffected: still [0.768, 0.187]. Head B's value slice for "century" was dims 3–4, i.e. [1,0], and becomes [3,0]. Recomputing: out_B = 0.140·[0,0] + 0.576·[0,0] + 0.284·[3,0] = [0.852, 0] — exactly triple the previous [0.284,0], since the weight on "century" (0.284) is unchanged and only its value vector scaled. The concatenation becomes [0.768, 0.187, 0.852, 0] — only the last two coordinates move. For the full head, "century"'s full value vector goes from [0,0,1,0] to [0,0,3,0] with an unchanged weight of 0.100, so its contribution goes from [0,0,0.1,0] to [0,0,0.3,0], making the full-head output [0.450, 0.450, 0.300, 0] — again, only the third coordinate moves, tripled. The general lesson: changing a value vector never touches attention weights, and it only moves the coordinates of the output that fall inside the dimensions that value vector occupies — which head(s) are affected depends entirely on which dimensional slice was changed.
Answer 4. With independent, mean-zero, variance-σ² components, Var(q·k) = d_k σ⁴ for a d_k-dimensional dot product (a sum of d_k independent variance-σ⁴ terms). Dividing by √d_k gives a scaled-score variance of d_k σ⁴ / d_k = σ⁴, independent of d_k. Since d_k shrinks exactly in proportion to 1/h, the 1/√d_k rule automatically compensates for any choice of h — the softmax input variance stays at the same constant σ⁴ whether you use 1 head or 16, so the rule needs deriving once, not once per head count.
Answer 5. No contradiction. The 4d² parameter-count equivalence is a statement about the architecture's capacity budget at design time — it says a multi-head layer costs no more to build or run than a single-head layer of the same width. Whether all h heads end up learning distinct, useful functions is a separate, empirical question about training dynamics: gradient descent may converge multiple heads to redundant or near-identical behavior, which is exactly what pruning studies detect and remove after the fact. The architecture guarantees h independent slots are available and costs nothing extra to provide them; it does not guarantee all h slots get used distinctly.
Answer 6. Single head: n² · d = 100² × 512 = 10{,}000 × 512 = 5{,}120{,}000 multiply-adds. Eight heads of d_k = 64: h · n² · d_k = 8 × 10{,}000 × 64 = 8 × 640{,}000 = 5{,}120{,}000 — identical.
Think About It
Think about this: How would you explain multi-head attention: deep mathematical analysis 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.
Practice Exercises
Now it is time to practice! Complete these challenges to solidify your understanding:
- Exercise 1: Write a short program that demonstrates the core concept from this chapter. Test it with at least 3 different inputs.
- Exercise 2: Find a real-world example where multi-head attention: deep mathematical analysis is used in an Indian company (like TCS, Infosys, Flipkart, or ISRO). Write a paragraph explaining the connection.
- Exercise 3: Create a mind-map connecting multi-head attention: deep mathematical analysis to at least 3 other topics you have studied.
Key Takeaways — Summary and Recap
Let us recap what we covered: the core ideas behind multi-head attention: deep mathematical analysis, how they connect to real-world applications, and why they matter for your journey in computer science. Remember these key points as you move forward. For competitive exam preparation (CBSE, JEE, BITSAT), focus on understanding the WHY behind each concept, not just the WHAT.