Multi-Head Attention: Shapes, Q/K/V, and the Scale

Attention is one line of maths — softmax(QK^T / sqrt(d_k)) V — wrapped in a handful of linear projections and a reshape. Yet that compact expression is where a transformer spends much of its ingenuity, and getting the shapes exactly right is what separates a vague intuition from a working mental model. This article dissects the multi-head self-attention sublayer end to end: how Q, K, and V are produced, why the dot product is scaled, how the computation splits across heads, and how the causal mask enforces autoregression. We assume you've seen the forward-pass map from the previous article.

From residual stream to Q, K, V

The sublayer receives the normalized residual stream: a matrix X of shape n by d, where n is the sequence length and d the model dimension. Attention's first act is to produce three views of X through three learned linear projections. Multiplying X by weight matrices W_Q, W_K, and W_V (each d by d) yields the queries Q, keys K, and values V, all of shape n by d.

These three are not mystical — they are just X seen through three different learned lenses. The query is what each position is looking for, the key is what each position offers as a match target, and the value is the content each position will contribute if attended to. Because the projections are learned, the model discovers for itself what should count as a good query-key match.

The normalized input X is projected by three learned matrices into queries, keys, and values.

Scaled dot-product attention

With Q, K, V in hand, attention is a four-step recipe. Compute the score matrix QK^T of shape n by n — entry (i, j) is the dot product of query i with key j, measuring their compatibility. Divide every score by sqrt(d_k), the square root of the per-head key dimension. Apply softmax along each row so each query's scores become a probability distribution over positions. Finally multiply those weights by V to get a weighted sum of values, an n by d output.

That single chain — score, scale, softmax, weighted sum — is the whole attention operation. The n by n score matrix is the crux: it is where every position meets every other, and it is the object whose size drives attention's quadratic cost in sequence length.

Scores from QK^T are scaled, softmaxed row-wise, then used to weight V into the output.

Why divide by the square root of d_k

The scaling factor is not cosmetic. If query and key components are roughly independent with unit variance, their dot product over d_k dimensions has variance proportional to d_k, so raw scores grow with the head dimension. Feed large-magnitude scores into softmax and it saturates: almost all the probability collapses onto a single position and the gradient through the other entries vanishes.

Dividing by sqrt(d_k) normalizes the score variance back to order one regardless of head size, keeping softmax in a responsive regime where gradients stay healthy. It is a small correction with an outsized effect on trainability — omit it and deep attention stacks become far harder to optimize.

📌

The intuition. Dot products grow with d_k, so dividing by sqrt(d_k) keeps score variance near one and stops softmax from saturating.

Splitting into heads

A single attention operation over the full d dimensions can only express one pattern of mixing at a time. Multi-head attention instead splits the d-wide Q, K, V into h heads, each of width d_k = d / h, and runs the scaled dot-product independently in every head. Head one might track subject-verb agreement, head two coreference, head three positional locality — each a different learned relationship over its own subspace.

Mechanically this is a reshape, not extra parameters: the same n by d projections are viewed as h slices of width d_k. Each head produces its own n by n score matrix and its own n by d_k output. The heads run in parallel and never interact until they are recombined, which is the next step.

The d-wide projections are reshaped into h heads of width d_k, each attending in its own subspace.

Concatenate and project out

After every head has computed its n by d_k output, the heads are concatenated back into an n by d matrix — reassembling the full width. That concatenation is then passed through a final learned output projection W_O (d by d), which lets the model mix information across heads and map the combined result back into the residual stream's coordinate system.

The complete multi-head sublayer is therefore: project to Q, K, V; reshape into heads; scaled dot-product per head; concatenate; project with W_O. The result is added back to the residual stream, exactly the x + Sublayer(Norm(x)) pattern from the forward-pass article.

The causal mask

In a decoder-only model, position i must not attend to positions after it — otherwise the model could peek at future tokens it is supposed to predict. This is enforced with a causal mask: before softmax, every score (i, j) with j greater than i is set to negative infinity, so softmax drives its weight to zero. The result is a lower-triangular attention pattern where each position sees only itself and the past.

This mask is what makes the same architecture usable for autoregressive generation. It also has a practical consequence we'll exploit later: since position i never depends on future positions, the keys and values of past tokens are fixed once computed — the exact property that makes the KV cache possible.

Each query attends only to itself and earlier positions; the upper triangle is masked out.

Why this matters

Every downstream fact about transformer inference traces back to the shapes in this sublayer. The n by n score matrix explains the quadratic scaling of prefill. The per-head d_k sets the scaling constant. The fixed-past property created by the causal mask is the seed of the KV cache. Knowing exactly which matrix has which shape lets you predict memory, cost, and latency instead of guessing.

Next we turn those shapes into numbers. The following article counts the arithmetic of a single token — the FLOPs and byte traffic of one forward pass — and shows why the popular rule of thumb of roughly two floating-point operations per parameter per token holds.

Key takeaways

Attention projects the input X (n x d) into Q, K, V via three learned d x d matrices.
Scaled dot-product attention is softmax(QK^T / sqrt(d_k)) V, producing an n x d output.
The n x n score matrix is where every position meets every other and drives attention's quadratic cost.
Dividing by sqrt(d_k) keeps score variance near one so softmax doesn't saturate.
Multi-head attention reshapes Q, K, V into h heads of width d_k = d/h, run in parallel, then concatenated and projected by W_O.
A causal mask sets future scores to -inf, giving lower-triangular attention and enabling the KV cache.

Glossary

Query (Q): The projection X W_Q representing what each position is looking for.
Key (K): The projection X W_K representing what each position offers as a match target.
Value (V): The projection X W_V representing the content a position contributes when attended to.
Score matrix: The n x n matrix QK^T of query-key compatibilities, one row per query.
Scaling factor: The 1/sqrt(d_k) term that normalizes score variance before softmax.
Head: One of h parallel attention computations, each over a d_k = d/h dimensional subspace.
Output projection (W_O): The final d x d matrix that mixes concatenated head outputs back into the residual stream.
Causal mask: A mask that sets scores to future positions to negative infinity, preventing leakage.