Understanding Transformers · Part 8 of 16

Multi-Head Attention

Running several attention operations at once, so one representation can track several kinds of relationship.

The intuition: several relationships at once

A single attention operation has to squeeze every relationship worth noticing into one pattern of weights. But a sentence carries several kinds of relationship at the same time. Take: “The scientist thanked the engineer because she solved the problem.” Resolving “she” is one relationship; linking “solved” to “problem” is another; tracking the overall grammar is a third. Forcing all of these through one set of attention weights makes them compete.

Multi-head attention runs several attention operations in parallel (Vaswani et al., 2017), each with its own WQ, WK, and WV. Each of these heads can settle into a different kind of comparison without stepping on the others: one on nearby grammar, one on subject–verb links, one on long-range references.

flowchart TD
  X["X"] --> H1["head 1"]
  X --> H2["head 2"]
  X --> H3["head 3"]
  X --> H4["head h"]
  H1 --> CAT["concat the outputs"]
  H2 --> CAT
  H3 --> CAT
  H4 --> CAT
  CAT --> WO["W_O"]
  

In the original design each head works in a smaller Query/Key/Value space, dk = dmodel/h, so h heads cost about as much as one full-width head (Vaswani et al., 2017). Their outputs are concatenated and passed through one more learned matrix, WO (used even when there is only one head), which recombines the separate perspectives into the single representation the next layer expects.

import numpy as np

# same 3 tokens as Part 7
X = np.array([[2, 0], [1, 2], [2, 0]])

def head(Wq, Wk, Wv):
    Q, K, V = X @ Wq, X @ Wk, X @ Wv
    scores = Q @ K.T / np.sqrt(Q.shape[-1])
    weights = np.exp(scores) / np.exp(scores).sum(axis=1, keepdims=True)
    return weights @ V

# Two heads, each with its own Wq/Wk/Wv, so each makes a different comparison.
# For readability these heads keep all 2 dimensions instead of d_model/h = 1,
# which is why the concatenation below is wider than X.
head1 = head(np.array([[0,1],[1,0]]), np.eye(2), np.eye(2))
head2 = head(np.eye(2), np.array([[1,0],[0,-1]]), np.eye(2))

concat = np.concatenate([head1, head2], axis=-1)   # 3x4

# one more matrix, back down to 3x2 (random here, learned in training)
Wo = np.random.randn(4, 2)
output = concat @ Wo