Challenge 1

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, A

Test its correctness two ways: the weight rows must sum to 1 and the past must not change when the future changes.

Exercise 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.

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 = 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]))

Exercise 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.

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(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))

Test yourself

With d = 12, h = 3, by which number are the scores divided?
- 2
- √12 ≈ 3.46
- 4
- 3
Show answer

2 — Each head is d/h = 4 dimensional; √4 = 2.

Concatenate the heads back into shape (n, d).

Y = (A @ Vh).transpose(1, 0, 2).____(n, d)

Show answer

Y = (A @ Vh).transpose(1, 0, 2).reshape(n, d)

Match the tests.

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

Next lesson: Challenge 2

We Are Enhancing Your Experience

Cookie Preferences

Strictly Necessary Cookies

Analytical Cookies