Machine Learning
Multi-head
Multi-head attention runs multiple attention heads in parallel.
Multi-head attention
A single attention head focuses on a single kind of relationship. Multi-head attention splits the d-dimensional representation into h parts; each head learns a different relationship (grammar, coreference, position) with its own Q, K, V, and the results are concatenated and passed through an output transformation:
# d_model = 512, h = 8 -> each head d_k = 64
# each head: attention(X W_q^i, X W_k^i, X W_v^i) (n, 64)
# concat: concat(head_1..head_8) (n, 512)
# output: concat @ W_o (n, 512)
The total computation is about the same as a single 512-dimensional head; but the model gains several points of view at once.
Splitting into heads
In practice the heads are computed not with separate matrices but by reshaping a single large projection: (n, d) → (h, n, d/h). That way all heads are processed with a single batched matrix multiplication:
n, d, h = 5, 8, 2
Q = X @ W_q # (5, 8)
Qh = Q.reshape(n, h, d // h).transpose(1, 0, 2) # (2, 5, 4)
skor = Qh @ Kh.transpose(0, 2, 1) / np.sqrt(d // h) # (2, 5, 5)
Yh = softmax(skor) @ Vh # (2, 5, 4)
Y = Yh.transpose(1, 0, 2).reshape(n, d) # (5, 8)
In PyTorch, nn.MultiheadAttention(embed_dim, num_heads) does all of this.
Exercise Multi-head attention
Compute 2-head attention with the given X (5×8) and weights: split Q, K, V into heads, apply scaled attention, concatenate. Print Qh.shape, the shape of the heads' weight matrix and the shape of the concatenated output on one line. Then print with np.allclose that every head's weight rows sum to 1.
Show example solution
import numpy as np
def softmax(z, eksen=-1):
e = np.exp(z - z.max(axis=eksen, keepdims=True))
return e / e.sum(axis=eksen, keepdims=True)
rng = np.random.default_rng(0)
n, d, h = 5, 8, 2
X = rng.normal(size=(n, d))
W_q, W_k, W_v = rng.normal(size=(3, d, d))
bol = lambda M: M.reshape(n, h, d // h).transpose(1, 0, 2)
Qh, Kh, Vh = bol(X @ W_q), bol(X @ W_k), bol(X @ W_v)
A = softmax(Qh @ Kh.transpose(0, 2, 1) / np.sqrt(d // h), eksen=-1)
Y = (A @ Vh).transpose(1, 0, 2).reshape(n, d)
print(Qh.shape, A.shape, Y.shape)
print(np.allclose(A.sum(axis=-1), 1))
Test yourself
-
If d_model = 768 with 12 heads, what is each head's size?
Show answer
64 — 768 / 12 = 64 (the BERT-base setting).
-
Split the representation into heads.
Qh = Q.reshape(n, h, d // h).____(1, 0, 2)
Show answer
Qh = Q.reshape(n, h, d // h).transpose(1, 0, 2)
-
Match the shapes (n=5, d=8, h=2).
- Q
- (5, 8)
- Qh
- (2, 5, 4)
- score
- (2, 5, 5)
- d // h
- 4
Next lesson: Summary - Attention
