Multi-Head Attention

A single attention operation produces one weighted average per query. Whatever the model wants from a position — the subject of the sentence, the matching bracket, the last occurrence of this token — it must be extracted by one softmax over one similarity function. Multi-head attention runs $h$ of them over disjoint $d_h = d/h$-dimensional subspaces and adds the results.

Written as a sum rather than a concatenation, the structure is easier to read. Each head has its own read matrices $W^Q_i, W^K_i, W^V_i \in \R^{d \times d_h}$ and its own write matrix $W^O_i \in \R^{d_h \times d}$, and it contributes to the residual stream independently of the others. Heads do not interact within a layer; they only ever meet by adding into the same stream.

The parameter count does not change

With $d_h = d/h$, the four projections total $4 d^2$ regardless of $h$. Splitting into more heads is free in parameters and nearly free in FLOPs — it buys parallel reads rather than capacity.

What does change is the attention matrices: $h$ of them, each $n \times n$. That term is why context length is expensive, and why it is the first thing FlashAttention refuses to materialise.

What a head is for

The circuits literature treats each head as a pair of low-rank operators: a $QK$ matrix deciding where to read, and an $OV$ matrix deciding what to write. The two are independent, which is why a head can be described by where it attends without reference to what it copies — and why induction heads can be identified by their attention pattern alone.

Pruning studies find heads are not equally load-bearing: many can be removed at test time with almost no loss, while a few are individually critical. The redundancy appears to be useful during training and largely surplus afterwards.

import torch
from torch import Tensor, nn


class MultiHeadAttention(nn.Module):
    def __init__(self, d_model: int, n_heads: int) -> None:
        super().__init__()
        assert d_model % n_heads == 0
        self.h = n_heads
        self.d_head = d_model // n_heads
        self.qkv = nn.Linear(d_model, 3 * d_model, bias=False)
        self.proj = nn.Linear(d_model, d_model, bias=False)

    def forward(self, x: Tensor, causal: bool = True) -> Tensor:
        b, n, _ = x.shape
        # One matmul for all three projections, then split heads out.
        q, k, v = self.qkv(x).chunk(3, dim=-1)
        q, k, v = (
            t.view(b, n, self.h, self.d_head).transpose(1, 2) for t in (q, k, v)
        )

        out = torch.nn.functional.scaled_dot_product_attention(
            q, k, v, is_causal=causal
        )
        # Concatenate the heads back into the model dimension.
        out = out.transpose(1, 2).reshape(b, n, self.h * self.d_head)
        return self.proj(out)

Fusing $W^Q, W^K, W^V$ into one Linear is worth doing: it turns three small matmuls into one large one, which matters more than it looks at batch sizes where the GEMMs are launch-bound. The transpose(1, 2) before and after is the only real cost of the head decomposition, and it is a view plus a contiguous copy, not arithmetic.