Multi-head attention

A single attention head computes one weighted average over all value vectors. Different positions might need to attend to different aspects of the context simultaneously -- syntactic head agreements, coreference, and positional proximity are all different signals. Multi-head attention runs $h$ attention heads in parallel on different linear projections of the input, then concatenates their outputs[1].

The multi-head mechanism

Given input $X \in \mathbb{R}^{T \times d_{\text{model}}}$, multi-head attention with $h$ heads and head dimension $d_k = d_v = d_{\text{model}} / h$:

For each head $i = 1, \dots, h$:

where $W_i^Q, W_i^K \in \mathbb{R}^{d_{\text{model}} \times d_k}$ and $W_i^V \in \mathbb{R}^{d_{\text{model}} \times d_v}$ are per-head projection matrices.

Concatenate all heads and apply an output projection:

where $W^O \in \mathbb{R}^{h d_v \times d_{\text{model}}}$ projects back to $d_{\text{model}}$.

Parameter count

Each head has three projection matrices of size $d_{\text{model}} \times d_k$. With $h$ heads and $d_k = d_{\text{model}}/h$:

Same total as one head of size $d_{\text{model}}$, but split across $h$ subspaces. The output projection adds no parameters beyond what a single full-rank projection would cost.

Why multiple heads?

Each head learns different projection matrices, so it attends to a different subspace of the representation. Empirically:

• Some heads specialize in positional patterns (attending to the previous token or the first token).
• Some heads capture syntactic structure (subject-verb agreement, dependency relations).
• Some heads implement coreference (associating pronouns with their referents).

No head is forced into these roles -- they emerge from training. Multi-head attention gives the model enough capacity to represent all these patterns simultaneously.

Implementation (batched)

import numpy as np

def softmax(x, axis=-1):
    e = np.exp(x - x.max(axis=axis, keepdims=True))
    return e / e.sum(axis=axis, keepdims=True)

class MultiHeadAttention:
    def __init__(self, d_model, n_heads, seed=0):
        assert d_model % n_heads == 0
        rng = np.random.default_rng(seed)
        s = 0.02
        self.h  = n_heads
        self.dk = d_model // n_heads
        # Stack all h projections into one matrix for efficiency
        self.Wq = rng.normal(0, s, (d_model, d_model))   # (d, d)
        self.Wk = rng.normal(0, s, (d_model, d_model))
        self.Wv = rng.normal(0, s, (d_model, d_model))
        self.Wo = rng.normal(0, s, (d_model, d_model))

    def split_heads(self, x):
        """x: (T, d_model) -> (h, T, dk)"""
        T, d = x.shape
        return x.reshape(T, self.h, self.dk).transpose(1, 0, 2)

    def forward(self, x, mask=None):
        """x: (T, d_model). mask: (T, T) boolean."""
        T, d = x.shape
        Q = self.split_heads(x @ self.Wq)   # (h, T, dk)
        K = self.split_heads(x @ self.Wk)
        V = self.split_heads(x @ self.Wv)

        # Scaled dot-product attention for all heads simultaneously
        scores = Q @ K.transpose(0, 2, 1) / np.sqrt(self.dk)  # (h, T, T)
        if mask is not None:
            scores[:, mask] = -1e9
        attn = softmax(scores)                                  # (h, T, T)
        out  = attn @ V                                         # (h, T, dk)

        # Concatenate heads: (h, T, dk) -> (T, d_model)
        out = out.transpose(1, 0, 2).reshape(T, d)
        return out @ self.Wo, attn                              # (T, d_model), (h, T, T)

T, d_model, h = 6, 16, 4
mha = MultiHeadAttention(d_model, h)
x   = np.random.randn(T, d_model)
mask = np.triu(np.ones((T, T), dtype=bool), k=1)  # causal

out, attn = mha.forward(x, mask=mask)
print("Output shape:", out.shape)          # (T, d_model)
print("Attention shape:", attn.shape)      # (h, T, T)
print("Head 0 attention (first 3 tokens):")
print(np.round(attn[0, :3, :3], 3))

Efficient implementation note

In practice, the $h$ projections are fused into a single matrix multiply producing $(B, T, 3 d_{\text{model}})$ and then split into Q, K, V -- a factor-of-3 memory reduction and better GPU utilization vs. $3h$ separate matrix multiplies.

Exercises

Exercise 1. Take $d_{\text{model}} = 512$ and $h = 8$. Find the per-head dimension $d_k$, the parameter count of the four projection matrices, and confirm it matches a single full-rank attention with $d_k = d_{\text{model}}$.

Exercise 2. The split_heads operation reshapes $(T, d_{\text{model}})$ to $(h, T, d_k)$ and the final concat reverses it. Explain why reshape-then-transpose is the correct order, and why simply reshaping to $(h, T, d_k)$ without a transpose would be wrong.

Exercise 3 (code). Verify the split/concat roundtrip is exact, and that two heads of a randomly initialized layer produce different attention patterns (so the heads are not redundant copies).

import numpy as np

def softmax(x, axis=-1):
    e = np.exp(x - x.max(axis=axis, keepdims=True))
    return e / e.sum(axis=axis, keepdims=True)

rng = np.random.default_rng(0)
T, d_model, h = 6, 16, 4
dk = d_model // h

x = rng.normal(size=(T, d_model))
split = x.reshape(T, h, dk).transpose(1, 0, 2)   # (h, T, dk)
concat = split.transpose(1, 0, 2).reshape(T, d_model)
assert np.allclose(concat, x)                    # roundtrip is lossless

Wq = rng.normal(0, 0.5, (d_model, d_model))
Wk = rng.normal(0, 0.5, (d_model, d_model))
Q = (x @ Wq).reshape(T, h, dk).transpose(1, 0, 2)
K = (x @ Wk).reshape(T, h, dk).transpose(1, 0, 2)
attn = softmax(Q @ K.transpose(0, 2, 1) / np.sqrt(dk))   # (h, T, T)

assert not np.allclose(attn[0], attn[1])         # head 0 and head 1 differ
print("roundtrip exact; heads attend differently, max diff",
      round(np.abs(attn[0] - attn[1]).max(), 3))

The next lesson assembles multi-head attention, a position-wise feed-forward network, residual connections, and layer normalization into the full transformer block.