Challenge 1
- Machine Learning course
- Expert
- Lesson 41/44
- 8 min
Apply the attention mechanism.
Challenge: causal multi-head attention
In this challenge you'll write the layer at the heart of GPT from scratch: causal multi-head attention. Put the pieces together:
- Q, K, V projections and the (n, d) → (h, n, d/h) split
- Scores scaled by √(d/h)
- A −∞ mask above the diagonal
- Concatenating the heads and the W_o output projection
def mha(x, Wq, Wk, Wv, Wo, h, nedensel=True):
n, d = x.shape
bol = lambda M: M.reshape(n, h, d // h).transpose(1, 0, 2)
...
return Y @ Wo, ATest its correctness two ways: the weight rows must sum to 1 and the past must not change when the future changes.
Practice Causal MHA
Write the mha(x, Wq, Wk, Wv, Wo, h, nedensel=True) function; have it return the output and the (h, n, n) weights. For the given x (6×12) and h = 3, print the output and weight shapes on one line. Then print these three tests on one line: do the rows sum to 1 (np.allclose), is the first position's weight on the others 0 ((A[:, 0, 1:] == 0).all()), are the first five outputs the same when the last row changes.
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 = 6, 12, 3
x = rng.normal(size=(n, d))
Wq, Wk, Wv, Wo = rng.normal(size=(4, d, d)) * 0.3
x2 = x.copy()
x2[-1] += 2
# 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 = 6, 12, 3
x = rng.normal(size=(n, d))
Wq, Wk, Wv, Wo = rng.normal(size=(4, d, d)) * 0.3
x2 = x.copy()
x2[-1] += 2
def mha(x, Wq, Wk, Wv, Wo, h, nedensel=True):
n, d = x.shape
bol = lambda M: M.reshape(n, h, d // h).transpose(1, 0, 2)
Qh, Kh, Vh = bol(x @ Wq), bol(x @ Wk), bol(x @ Wv)
skor = Qh @ Kh.transpose(0, 2, 1) / np.sqrt(d // h)
if nedensel:
skor = np.where(np.triu(np.ones((n, n), dtype=bool), k=1), -np.inf, skor)
A = softmax(skor)
Y = (A @ Vh).transpose(1, 0, 2).reshape(n, d)
return Y @ Wo, A
Y, A = mha(x, Wq, Wk, Wv, Wo, h)
Y2, _ = mha(x2, Wq, Wk, Wv, Wo, h)
print(Y.shape, A.shape)
print(np.allclose(A.sum(-1), 1), (A[:, 0, 1:] == 0).all(), np.allclose(Y[:5], Y2[:5]))Expected output
(6, 12) (3, 6, 6) True True True
Practice A soft dictionary lookup
Show that attention works like a lookup table. For the given 8 keys (unit vectors multiplied by a scale s) and 8 values, use each key itself as a query; find the nearest value (argmin distance) to the output of softmax(Q @ K.T) @ V. For s = 0.5, 2 and 8, print the hit rate and the average weight on the correct value (the diagonal mean) with round(..., 2) on one line each.
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(1)
birim = np.eye(8)
V = rng.normal(size=(8, 3))
# 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(1)
birim = np.eye(8)
V = rng.normal(size=(8, 3))
for s in [0.5, 2, 8]:
K = birim * s
A = softmax(K @ K.T)
Y = A @ V
uzak = ((Y[:, None, :] - V[None, :, :]) ** 2).sum(-1)
dogru = (uzak.argmin(axis=1) == np.arange(8)).mean()
print(s, round(dogru, 2), round(np.diag(A).mean(), 2))Expected output
0.5 0.12 0.16 2 1.0 0.89 8 1.0 1.0
Questions and answers
Review the lesson questions together with their answers. To test yourself, solve the lesson interactively.
1 With d = 12, h = 3, by which number are the scores divided?
- 2 Correct answer
- √12 ≈ 3.46
- 4
- 3
Explanation: Each head is d/h = 4 dimensional; √4 = 2.
2 Concatenate the heads back into shape (n, d).
Y = (A @ Vh).transpose(1, 0, 2).____(n, d)Y = (A @ Vh).transpose(1, 0, 2).reshape(n, d)Explanation: First put it in (n, h, d/h) order, then merge.
3 Match the tests.
| Term | Meaning |
|---|---|
A.sum(-1) ≈ 1 |
Valid distribution |
the past doesn't change |
Causality |
A[:, 0, 1:] = 0 |
The first position only looks at itself |
Y.shape = (n, d) |
Stackability |