Multi-head
- Machine Learning course
- Expert
- Lesson 5/44
- 5 min
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.
Practice 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.
Starter code
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))
# buraya yazSample 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))Expected output
(2, 5, 4) (2, 5, 5) (5, 8) True
Questions and answers
Review the lesson questions together with their answers. To test yourself, solve the lesson interactively.
1 If d_model = 768 with 12 heads, what is each head's size?
- 9216
- 64 Correct answer
- 768
- 12
Explanation: 768 / 12 = 64 (the BERT-base setting).
2 Split the representation into heads.
Qh = Q.reshape(n, h, d // h).____(1, 0, 2)Qh = Q.reshape(n, h, d // h).transpose(1, 0, 2)Explanation: In NumPy, transpose takes the new order of the axes.
3 Match the shapes (n=5, d=8, h=2).
| Term | Meaning |
|---|---|
Q |
(5, 8) |
Qh |
(2, 5, 4) |
score |
(2, 5, 5) |
d // h |
4 |