The Shapes of Multi-Head Attention

The previous chapter gave the idea; this one follows the shapes. Tracking dimensions is the difference between understanding multi-head attention as a picture and understanding it well enough to build. There is no new concept here — just careful bookkeeping, and one neat implementation trick.

The dimensions in play

Four numbers describe the data inside an attention layer:

• $B$ — the batch size, how many sequences we process at once.
• $n$ — the sequence length, the number of tokens.
• $d$ — the model dimension, the width of each token's vector.
• $h$ — the number of heads, with per-head dimension $d_k = d/h$.

The input to the layer is a block of shape $B \times n \times d$: for each of $B$ sequences, $n$ tokens, each a $d$-vector.

One projection, then a reshape

Conceptually each head has its own $W_Q^{(i)}, W_K^{(i)}, W_V^{(i)}$. In practice we do not run $h$ separate small multiplications. We apply one big projection $W_Q$ of shape $d \times d$ to get all the queries at once — output shape $B \times n \times d$ — and then reshape that last dimension $d$ into $h \times d_k$:

Nothing is computed differently; we have simply relabelled the $d$ numbers as "$h$ heads of $d_k$ numbers each." A final transpose puts the heads next to the batch, giving $B \times h \times n \times d_k$ — which reads as "$B \cdot h$ independent little attention problems, each of $n$ tokens in dimension $d_k$." The same is done for keys and values.

The per-head computation, by shape

With queries, keys and values all shaped $B \times h \times n \times d_k$, the attention from Module 3 runs on the last two axes, identically for every $(B, h)$ pair:

1. Scores $Q K^{\top}$: multiplying $n \times d_k$ by its transpose $d_k \times n$ gives an $n \times n$ score grid per head → shape $B \times h \times n \times n$.
2. Scale by $1/\sqrt{d_k}$ and add the causal mask (the same lower-triangular mask, broadcast across batch and heads).
3. Softmax over the last axis, so each of the $n$ rows sums to 1.
4. Blend values: multiply the $n \times n$ weights by the $n \times d_k$ values → back to $B \times h \times n \times d_k$.

Every head has done a full scaled, masked attention, all as batched matrix multiplies.

Fanning back in

Now reverse the reshape. Transpose the heads back beside the head-dimension and merge them:

This is the concatenation of the heads — again just a relabelling of axes. Finally apply the output projection $W_O$ (shape $d \times d$), mixing the heads' contributions, to produce the layer's output at shape $B \times n \times d$ — exactly the shape we started with, ready to be added back to the residual stream.

The round trip in one view

The shape you start with is the shape you end with — the layer enriches each token's vector without changing its size, which is exactly what lets you stack these blocks as deep as you like. A count worth carrying: four projections ($W_Q, W_K, W_V, W_O$), each $d \times d$, make up the parameters of a multi-head attention layer.