Multi-Head Attention: Parallel Subspaces

Multi-Head Attention (MHA) runs h attention operations in parallel, each on a lower-dimensional subspace, then concatenates the results and projects them back to d_model. The intuition: instead of one attention pass that has to capture all relational patterns at once, give the model h shorter "specialists" — different heads can learn different patterns.

MultiHead(Q, K, V) = Concat(head_1, ..., head_h) · W_O
head_i = Attention(Q · W_i^Q, K · W_i^K, V · W_i^V)

Concretely: pick h (typically 8, 12, 32, or 64). Set d_head = d_model / h (typically 64 or 128). Each head has its own W^Q_i, W^K_i, W^V_i of shape (d_model, d_head). Concatenate the h outputs of shape (seq_len, d_head) into one tensor of shape (seq_len, d_model). Apply a final output projection W_O of shape (d_model, d_model) to mix information across heads.

Why this matters

Different heads, observed in trained models, often specialize:

Some heads attend only to the previous token ("position −1 head").
Some attend to the same token across recent context ("induction heads").
Some attend to syntactically related tokens (subject ↔ verb).
Some attend nearly uniformly to capture semantic averaging.

This emergent specialization is one of the more robust empirical findings about Transformers.

Code

Multi-head attention forward pass·python

class MultiHeadAttention(nn.Module):
    def __init__(self, d_model, n_heads):
        super().__init__()
        assert d_model % n_heads == 0
        self.n_heads, self.d_head = n_heads, d_model // n_heads
        self.W_q = nn.Linear(d_model, d_model, bias=False)
        self.W_k = nn.Linear(d_model, d_model, bias=False)
        self.W_v = nn.Linear(d_model, d_model, bias=False)
        self.W_o = nn.Linear(d_model, d_model, bias=False)
    def forward(self, x, mask=None):
        B, L, d = x.shape
        H, dh = self.n_heads, self.d_head
        Q = self.W_q(x).view(B, L, H, dh).transpose(1, 2)   # (B, H, L, dh)
        K = self.W_k(x).view(B, L, H, dh).transpose(1, 2)
        V = self.W_v(x).view(B, L, H, dh).transpose(1, 2)
        out = F.scaled_dot_product_attention(Q, K, V, attn_mask=mask)
        out = out.transpose(1, 2).contiguous().view(B, L, d)
        return self.W_o(out)

Multi-Head Attention: Parallel Subspaces

Why this matters

Code

External links

Exercise

Progress

Comments 0