Multi-head attention: Question bank — AI-ML deep dive

Definition Symbols used on this page (all built in the parent note)

$d_{model}$ — the width of each token's vector coming into MHA (e.g. 512).
$n$ — the sequence length, i.e. how many tokens are in the input row-by-row. Every $Q$ , $K$ , $V$ has one row per token, so they are $[n, d_{model}]$ matrices.
$m$ — a second sequence length, used only in cross-attention, where the keys and values come from a different sequence (the encoder's $m$ tokens) than the queries (the decoder's $n$ tokens).
$h$ — the number of parallel heads.
$d_{k} = d_{v} = d_{model} / h$ — the width each head works in. This requires $h$ to divide $d_{model}$ evenly, otherwise $d_{k}$ is not a whole number and the heads cannot be concatenated back to width $d_{model}$ .
Per head: $Q_{i}, K_{i}, V_{i} \in R^{[n, d_{k}]}$ (the projected, per-head query/key/value matrices), and the scores $Q_{i} K_{i}^{T} \in R^{[n, n]}$ .

What to observe in the figure above: the left grid is $Q$ with one row per token ( $n$ rows) and its columns sliced into coloured blocks of $d_{k}$ each — one block per head; the right grid is a single head's score matrix, square and $[n, n]$ (one cell per pair of tokens, not per feature). Carry that "rows = tokens, score matrix = token × token" picture into every item below.

What to observe in the figure above: the pink measured curve (spread of $q \cdot k$ for random vectors) lands exactly on the yellow dashed $d_{k}$ line as $d_{k}$ grows from 1 to 128 — confirming that the typical logit magnitude is $d_{k}$ , so dividing by that constant (blue arrow) restores a spread of about 1.

True or false — justify

Each line: statement, then the verdict with a real reason.

Multi-head attention costs roughly

h

times more compute than one full-width head.

False — one head does an

n \times n

score step costing

O (n^{2} d_{k})

and an

n \times d_{k}

value-weighting step also

O (n^{2} d_{k})

; multiplying by

h

heads gives

h \cdot O (n^{2} d_{k}) = O (n^{2} \cdot h d_{k}) = O (n^{2} d_{model})

, the same total FLOPs as one full-width head.

If every head used the same projection matrices

W^{Q}, W^{K}, W^{V}

, MHA would still help because you average more attention maps.

False — identical projections produce identical heads, so concatenating

h

copies and blending with

W^{O}

is no richer than one head. The diversity comes entirely from the projections being different.

The

d_{k}

scaling factor changes between single-head and multi-head attention.

True in value — inside a head you scale by

d_{k} = d_{model} / h

, not

d_{model}

, because each head's dot products live in the smaller

d_{k}

-dimensional subspace.

Concatenating the heads already lets them share information with each other.

False — concatenation just stacks them side by side into separate channels; no head has yet "seen" another. The mixing happens only when

W^{O}

multiplies the concatenation.

The output projection

W^{O}

is optional — you could feed the raw concatenation to the next layer.

False in practice — without

W^{O}

the heads stay in fixed, non-interacting slots, so the model cannot learn to combine syntactic + semantic + positional signals. It also fixes which output dimensions each head "owns".

Multi-head attention and single-head attention with the same

d_{model}

have the same number of parameters in their projection matrices.

True — the

h

per-head matrices

W_{i}^{Q}

each of shape

d_{model} \times d_{k}

stack to one

d_{model} \times d_{model}

matrix, identical in parameter count to one full-width projection.

Each head must always attend most strongly to a token's own position.

False — self-attention weight is just softmax over learned scores; a head can legitimately put near-zero weight on the diagonal if the useful information lives elsewhere (e.g. a co-reference head looking backward).

In Cross-Attention, multi-head attention still splits into

h

heads the same way it does in Self-Attention.

True — the head-splitting is independent of where

Q

K

V

come from; cross-attention only differs in that

K, V

come from the encoder while

Q

comes from the decoder.

Any

h

works as long as it is at most

d_{model}

False —

h

must divide

d_{model}

so that

d_{k} = d_{model} / h

is a whole number; e.g.

h = 6

fails for

d_{model} = 512

even though

6 < 512

Spot the error

Each line states a flawed claim; the reveal names the specific mistake. (Recall from the definition box: $Q_{i}, K_{i} \in R^{[n, d_{k}]}$ , with $n$ = number of tokens.)

"With

d_{model} = 512

and

h = 6

heads, each head gets

d_{k} = 512/6

The error is choosing an

h

that does not divide

d_{model}

—

512/6

is not an integer, so concatenation would not return exactly

d_{model}

h

must divide

d_{model}

"The scores matrix

Q_{i} K_{i}^{T}

has shape

[d_{k}, d_{k}]

Wrong shape —

Q_{i}, K_{i}

are

[n, d_{k}]

(one row per token), so

Q_{i} K_{i}^{T}

[n, n]

: one score per pair of token positions, not per pair of feature dimensions.

"Softmax is applied over the whole

[n, n]

scores matrix at once so all entries sum to 1."

The error is the axis — softmax is applied per row, so each query token's weights over the

n

keys sum to 1, giving

n

separate probability distributions.

"We sum the head outputs to keep dimension at

d_{model}

The error is summing instead of concatenating — summing

h

vectors of size

d_{v}

gives one

d_{v}

-vector and destroys the distinct patterns. Concatenation gives

h \cdot d_{v} = d_{model}

while preserving each head.

W^{O}

has shape

d_{v} \times d_{model}

Wrong input width —

W^{O}

multiplies the concatenation of width

h \cdot d_{v} = d_{model}

, so

W^{O} \in R^{d_{model} \times d_{model}}

"Adding more heads always improves the model."

The error ignores the trade-off — pushing

h

up shrinks

d_{k} = d_{model} / h

, so each head has too few dimensions to represent anything useful, and quality falls.

"Positional patterns can't be captured because attention is permutation-agnostic, so MHA is useless for order."

Half-true premise, wrong conclusion — attention is order-blind by itself, which is exactly why Positional Encoding is added to inputs first; a head can then learn positional patterns from those encodings.

Why questions

Why do we use different projection matrices per head instead of one big shared projection?

Because different

W_{i}^{Q}, W_{i}^{K}, W_{i}^{V}

send the same input into different subspaces, so each head can specialise on a different relation type (grammar, meaning, position). Shared matrices would collapse all heads into one.

Why divide by

d_{k}

rather than

d_{k}

or nothing at all?

As the intuition box shows, a dot product of two random

d_{k}

-dimensional vectors (mean 0, variance 1 per entry, independent) has variance

d_{k}

, hence typical magnitude

d_{k}

; dividing by

d_{k}

rescales this back to variance

1

. Dividing by

d_{k}

would instead give variance

d_{k} / d_{k}^{2} = 1/ d_{k}

— the logits shrink toward

0

, softmax becomes nearly uniform, and the head loses its ability to focus. Dividing by nothing leaves variance

d_{k}

, so softmax saturates and gradients vanish.

Why concatenate the heads before the final projection rather than projecting each head and adding?

Concatenating first lets

W^{O}

learn arbitrary cross-head mixtures — every output dimension can draw on all heads. Projecting-then-adding would fix each head's contribution independently and forbid cross-head interaction.

Why does MHA use the same scaled dot-product attention formula in every head instead of different mechanisms?

The richness is meant to come from the learned projections, not the mechanism. Reusing one formula keeps every head interpretable in the same way and keeps the implementation a single batched operation. (See Linear Projections in Neural Networks.)

Why is

W^{O}

needed even though each head already produced a good context vector?

Because the heads sit in disjoint channels of the concatenation; only

W^{O}

can blend their specialised outputs into a single unified representation for the next layer.

Why can two heads still learn different things even though they start from the identical

Q, K, V

input?

Their projection matrices are initialised differently and trained by gradient descent, so they descend into different local specialisations — the shared input is filtered through distinct learned lenses. (See Attention Visualization to observe this.)

Edge cases

What happens when

h = 1

MHA reduces exactly to single-head attention:

d_{k} = d_{model}

, the "concatenation" is just the one head, and

W^{O}

becomes an ordinary output projection.

What happens as

h \to d_{model}

(each head one-dimensional)?

Each head has

d_{k} = 1

, so its

Q_{i} K_{i}^{T}

scores come from a single scalar feature — heads become extremely weak and can barely encode a relationship, degrading quality despite maximal "diversity".

If a query token's scores are all equal across positions, what does that head output?

Softmax gives a uniform distribution, so the head returns the plain average of all value vectors — a "no-preference" state that ignores content and just pools everything.

If one score is vastly larger than the rest, what does softmax do?

It saturates toward a one-hot weight, so the head effectively copies a single token's value — hard attention as a limiting case of soft attention.

In cross-attention where the encoder has

m

tokens (a second sequence length, defined in the symbol box) and the decoder has

n

, what shape are the per-head scores?

Q_{i}

[n, d_{k}]

(decoder) and

K_{i}

[m, d_{k}]

(encoder), so

Q_{i} K_{i}^{T}

[n, m]

— non-square, one row per decoder token attending over

m

encoder tokens. (See Cross-Attention.)

What if

d_{v} \neq = d_{k}

(heads project values to a different width)?

Attention still works —

d_{k}

only controls the score computation and scaling, while

d_{v}

controls the output width. Then the concatenation is

h \cdot d_{v}

, so

W^{O}

must have that as its input dimension.

At the very first layer, before any attention has run, where do

Q, K, V

come from?

From the token embeddings plus Positional Encoding, each linearly projected — there is no "previous layer" yet, so the raw embedded sequence is the shared input to all heads.

Recall Fast self-test

One-sentence answers — cover and recall. Does adding heads increase total attention FLOPs? ::: No — shrinking $d_{k}$ keeps it $O (n^{2} d_{model})$ . What single matrix lets heads talk to each other? ::: $W^{O}$ , applied after concatenation. Why not shared projections across heads? ::: They would produce identical, redundant heads. Softmax normalises along which axis? ::: Each query row, over the keys. Which numbers must $h$ divide? ::: $d_{model}$ , so that $d_{k}$ is a whole number.