What eight heads cost that one head does not

Split, attend, combine

def split_heads(X, n_heads):
    n, d = X.shape
    d_head = d // n_heads
    return X.reshape(n, n_heads, d_head).transpose(1, 0, 2)

def combine_heads(H):
    h, n, d_head = H.shape
    return H.transpose(1, 0, 2).reshape(n, h * d_head)

Project X to Q, K and V, split each into heads, run scaled dot-product attention per head, stack the heads back and multiply by the output matrix W_o. The GPU sees the per-head scores as a single batched matmul of shape (heads, N, d_head) times (heads, d_head, N). Adding heads adds batch entries, not weights.

Why the output projection matters

The heads never talk during the attention step itself; they work on different subspaces. W_o, of shape (d_model, d_model), is where their outputs mix. Without it the heads could only recombine through later layers and the block would be artificially constrained.

What heads learn

Probing work from 2019 to 2024 found distinct roles: heads that attend to the previous token, copy heads, heads that track named entities, and induction heads, the pairs that detect an earlier occurrence of the current token and copy what followed it, which underlie in-context learning. One head would have to smear all of these into one softmax.

Picking the head width

The head width lands at 64 or 128 almost every time. It is the unit of how much one head can see.

The production implementation

One linear layer maps D to 3D and is sliced into Q, K and V. Three separate layers compute the same thing, because three matmuls by (D, D) weights equal one matmul by a stacked (3D, D) weight, but the fused layer launches one kernel instead of three.