Multi-head attention

Direct answer

Multi-head attention runs several attention operations in parallel after applying different learned projections. Each head can form a different score matrix and value mixture. Their outputs are concatenated and projected back into the model width.

This expands the set of relationships a layer can represent, but it does not guarantee that every head learns a human-readable specialization.

Equation and shapes

For head (i):

The heads are recombined as:

In the original Transformer configuration, model width $d_{model}$ is divided across $h$ heads. With 512 model dimensions and 8 heads, each query and key head uses 64 dimensions. Concatenation restores the combined width before the output projection.

One-head shape walkthrough

Assume a batch of 2 sequences, 10 tokens per sequence, model width 512, and 8 heads.

The head dimension affects the scaling term. Each head divides its score by $\sqrt{64}$, not $\sqrt{512}$ in this example.

A transparent shape check

batch, tokens, model_width, heads = 2, 10, 512, 8

if model_width % heads != 0:
    raise ValueError("model width must divide evenly across heads")

head_width = model_width // heads
score_elements = batch * heads * tokens * tokens

print(f"head width: {head_width}")
print(f"score tensor: ({batch}, {heads}, {tokens}, {tokens})")
print(f"score elements: {score_elements}")

This example checks dimensions; it is not a neural-network implementation. Production kernels typically fuse projection and attention steps and may avoid materializing the full score matrix.

Why multiple heads help

A single attention distribution supplies one weighted mixture per query. Multiple learned projections allow a layer to preserve several mixtures before recombination. One head may emphasize nearby syntax while another uses a longer-range association, but such interpretations must be verified rather than assumed.

Increasing head count is not free. At fixed model width, more heads make each head narrower. The score tensor still scales with sequence length, and extra heads introduce projection and memory-layout costs. Head count, width, hardware kernels, and task behavior should be evaluated together.

Common mistakes

• Applying softmax across the head dimension instead of the key positions.
• Forgetting to transpose tokens and heads before matrix multiplication.
• Scaling by the full model width rather than the per-head key width.
• Assuming every head is necessary or semantically distinct.
• Comparing INT8 and floating-point implementations without checking the accumulator and scale strategy.

Interview takeaway

Multi-head attention creates several learned query-key-value subspaces, performs attention independently in each, concatenates the results, and uses an output projection to mix them.