Lesson 5 - Guided Project: A Multi-Head Attention Layer
Welcome to the Guided Project
Across this module you scaled a single attention head up to many. You saw why one head is a bottleneck, learned to split the model width into parallel heads of size , ran scaled dot-product attention in each head at once, concatenated the results, and mixed them back together with an output projection. In Lesson 4 you derived the backward pass through that whole pipeline — the split, the per-head attention, the concatenation, and the output projection. In this guided project you stop treating those steps as loose code and package them into one object: a MultiHeadAttention class with a forward and a backward method, holding its own four weight matrices and their gradients.
This class is the real thing, not a throwaway demo. It is the exact component you will drop into a transformer block in Module 5, wrap with a causal mask in Module 6, and train into a working char-level GPT in Module 7. Getting it clean and provably correct now means every later module inherits a block you already trust. So the project ends the way every serious layer should: a float64 gradient check comparing the analytic gradients to finite differences, and a short gradient-descent run that shows a real loss going down.
By the end of this project, you will be able to:
- Structure multi-head attention as a class that stores its four parameters (
Wq,Wk,Wv,Wo), a forward cache, and its gradients - Implement
forward(X)so it projects, splits into(B, h, T, hs), attends per head, concatenates, and projects withWo - Implement
backward(dout)so it returnsdXand storesself.dWq,self.dWk,self.dWv,self.dWo - Gradient-check the whole class in
float64, reporting max relative errors for all four weights andX - Run gradient-descent steps on the class and confirm a scalar loss decreases
This project assumes you have followed the earlier lessons in this module, especially the multi-head backward pass in Lesson 4. You only need numpy (version 2.x is fine) — no PyTorch, no GPU, no API key.
Stage 1: The Class
A layer needs three kinds of state: its parameters (the weights it learns), a cache (forward values the backward pass will reuse), and its gradients (filled in by backward). We set all three up in __init__, along with the three numbers that define the layer’s shape: the model width C, the number of heads h, and the head size hs = C // h.
Multi-head attention has four weight matrices, and this is the key difference from the single head of Module 2. Three of them — Wq, Wk, Wv — are now full (C, C) matrices rather than (C, hs): each projects a token’s whole C-dimensional vector to a whole C-dimensional query, key, or value, which we then slice into h heads. The fourth, Wo, is the (C, C) output projection that mixes the concatenated head outputs back together. We initialize all four small (* 0.02, the standard transformer init) and seed the draw so the class is reproducible.
import numpy as np
def softmax(x, axis=-1):
x = x - x.max(axis=axis, keepdims=True) # stable: subtract row max
e = np.exp(x)
return e / e.sum(axis=axis, keepdims=True)
class MultiHeadAttention:
"""Multi-head self-attention as one reusable NumPy layer."""
def __init__(self, C, h, seed=42):
assert C % h == 0, "C must be divisible by h"
self.C = C # model width
self.h = h # number of heads
self.hs = C // h # size of each head
rng = np.random.default_rng(seed)
self.Wq = rng.standard_normal((C, C)) * 0.02 # (C, C)
self.Wk = rng.standard_normal((C, C)) * 0.02 # (C, C)
self.Wv = rng.standard_normal((C, C)) * 0.02 # (C, C)
self.Wo = rng.standard_normal((C, C)) * 0.02 # (C, C) output projection
self.cache = None # filled by forward, read by backward
self.dWq = self.dWk = self.dWv = self.dWo = None # parameter gradients
layer = MultiHeadAttention(C=8, h=4)
print("C, h, hs:", layer.C, layer.h, layer.hs)
print("Wq/Wk/Wv:", layer.Wq.shape, layer.Wk.shape, layer.Wv.shape)
print("Wo :", layer.Wo.shape)
print("cache :", layer.cache)C, h, hs: 8 4 2
Wq/Wk/Wv: (8, 8) (8, 8) (8, 8)
Wo : (8, 8)
cache : NoneThe skeleton holds exactly what the layer must remember about itself: its shape (C = 8, h = 4, so hs = 2), its four (8, 8) weight matrices, and empty slots for the cache and gradients. Notice that __init__ takes only C and h — the batch size B and sequence length T are never baked in, so the same object works on a (2, 6, 8) input today and a (32, 32, 64) input when we scale up. The assert C % h == 0 guards the one hard requirement: the width must divide evenly into heads, or the split later would be ragged.
Why Wq/Wk/Wv are (C, C) now, not (C, hs)
In Module 2 a single head used (C, hs) projections. Multi-head attention instead projects to the full width C with one (C, C) matmul, then slices those C columns into h contiguous blocks of hs — head i is columns [i*hs : (i+1)*hs]. One wide matmul does the work of h small ones, and the new Wo at the end lets the heads’ outputs talk to each other before leaving the layer. Four matrices, one for each role: query, key, value, and mix.
Stage 2: The forward Method
Now we give the class a forward(X) method. It runs the full multi-head pipeline from this module, and — crucially — it caches every value the backward pass will need. Concretely, forward projects X to Q, K, V; reshapes each into h heads to get Qh, Kh, Vh of shape (B, h, T, hs); runs scaled dot-product attention in every head at once to get the weights A and per-head outputs Oh; concatenates the heads back to concat of shape (B, T, C); and applies the output projection out = concat @ Wo. It then stores (X, Qh, Kh, Vh, A, concat) on self.cache.
Why those six? The backward pass needs the input X (to form dWq, dWk, dWv), the per-head Vh and weights A (to differentiate Oh = A @ Vh), the per-head Qh and Kh (to differentiate the scores), and concat (to form dWo). The split is a reshape followed by transpose(0, 2, 1, 3), which moves the head axis up next to the batch so each head becomes an independent (T, hs) attention problem; concatenation is exactly that operation reversed.
import numpy as np
def softmax(x, axis=-1):
x = x - x.max(axis=axis, keepdims=True)
e = np.exp(x)
return e / e.sum(axis=axis, keepdims=True)
class MultiHeadAttention:
"""Multi-head self-attention as one reusable NumPy layer."""
def __init__(self, C, h, seed=42):
assert C % h == 0, "C must be divisible by h"
self.C = C
self.h = h
self.hs = C // h
rng = np.random.default_rng(seed)
self.Wq = rng.standard_normal((C, C)) * 0.02
self.Wk = rng.standard_normal((C, C)) * 0.02
self.Wv = rng.standard_normal((C, C)) * 0.02
self.Wo = rng.standard_normal((C, C)) * 0.02
self.cache = None
self.dWq = self.dWk = self.dWv = self.dWo = None
def forward(self, X):
B, T, C = X.shape
h, hs = self.h, self.hs
Q = X @ self.Wq # (B, T, C)
K = X @ self.Wk # (B, T, C)
V = X @ self.Wv # (B, T, C)
Qh = Q.reshape(B, T, h, hs).transpose(0, 2, 1, 3) # (B, h, T, hs)
Kh = K.reshape(B, T, h, hs).transpose(0, 2, 1, 3)
Vh = V.reshape(B, T, h, hs).transpose(0, 2, 1, 3)
scores = Qh @ Kh.transpose(0, 1, 3, 2) / np.sqrt(hs) # (B, h, T, T)
A = softmax(scores, axis=-1) # (B, h, T, T)
Oh = A @ Vh # (B, h, T, hs)
concat = Oh.transpose(0, 2, 1, 3).reshape(B, T, C) # (B, T, C)
out = concat @ self.Wo # (B, T, C)
self.cache = (X, Qh, Kh, Vh, A, concat) # all backward needs
return out
np.random.seed(42)
B, T, C, h = 2, 6, 8, 4
X = np.random.randn(B, T, C).astype(np.float64) # (2, 6, 8)
layer = MultiHeadAttention(C, h)
out = layer.forward(X)
Xc, Qh, Kh, Vh, A, concat = layer.cache
print("X.shape :", X.shape)
print("per-head Q/K/V:", Qh.shape, Kh.shape, Vh.shape)
print("A (weights) :", A.shape, "| rows sum to 1:", np.allclose(A.sum(-1), 1.0))
print("concat.shape :", concat.shape)
print("out.shape :", out.shape)X.shape : (2, 6, 8)
per-head Q/K/V: (2, 4, 6, 2) (2, 4, 6, 2) (2, 4, 6, 2)
A (weights) : (2, 4, 6, 6) | rows sum to 1: True
concat.shape : (2, 6, 8)
out.shape : (2, 6, 8)Every shape lands where the module’s conventions promise. The (2, 6, 8) input projects through the (8, 8) matrices to full-width Q, K, V, which split into (2, 4, 6, 2) = (B, h, T, hs): four heads, each holding six tokens of width two. The scores compare all six queries against all six keys within each head, giving (2, 4, 6, 6) = (B, h, T, T), and softmax over the last axis makes every row a distribution — confirmed by rows sum to 1: True. The per-head outputs concatenate back to (2, 6, 8), and Wo keeps them there. Unlike the single head of Module 2, the output is now the full model width C, ready to hand straight to the next layer. Written as math, the whole forward pass is:
where indexes the heads and lays them side by side back to width .
Stage 3: The backward Method
This is where the class earns its keep. The backward(dout) method takes the gradient of the loss with respect to the layer’s output and walks it back through every operation in the exact reverse order of forward, producing dX (which it returns) and the four parameter gradients self.dWq, self.dWk, self.dWv, self.dWo (which it stores). Each line is one step from the Lesson 4 derivation, reading its inputs from self.cache.
Walk it top to bottom and you are walking the forward pass in reverse:
out = concat @ WogivesdWo = concat^T doutanddconcat = dout @ Wo^T.- The concatenation is a reshape/transpose, so reversing it splits
dconcatback into per-headdOhof shape(B, h, T, hs). Oh = A @ VhgivesdA = dOh @ Vh^TanddValueH = A^T @ dOh.A = softmax(scores)gives the row-wise softmax gradientdS = A * (dA - sum(dA * A)).scores = Qh Kh^T / sqrt(hs)givesdscores, thendQueryH = dscores @ KhanddKeyH = dscores^T @ Qh.- Merging the heads turns
dQueryH,dKeyH,dValueHback into full-widthdQuery,dKey,dValue, which givedWq,dWk,dWvand the three contributions todX.
The suffix H marks a per-head (B, h, T, hs) gradient; dropping it after the merge marks the full-width (B, T, C) version.
import numpy as np
def softmax(x, axis=-1):
x = x - x.max(axis=axis, keepdims=True)
e = np.exp(x)
return e / e.sum(axis=axis, keepdims=True)
class MultiHeadAttention:
"""Multi-head self-attention as one reusable NumPy layer."""
def __init__(self, C, h, seed=42):
assert C % h == 0, "C must be divisible by h"
self.C = C
self.h = h
self.hs = C // h
rng = np.random.default_rng(seed)
self.Wq = rng.standard_normal((C, C)) * 0.02
self.Wk = rng.standard_normal((C, C)) * 0.02
self.Wv = rng.standard_normal((C, C)) * 0.02
self.Wo = rng.standard_normal((C, C)) * 0.02
self.cache = None
self.dWq = self.dWk = self.dWv = self.dWo = None
def forward(self, X):
B, T, C = X.shape
h, hs = self.h, self.hs
Q = X @ self.Wq
K = X @ self.Wk
V = X @ self.Wv
Qh = Q.reshape(B, T, h, hs).transpose(0, 2, 1, 3)
Kh = K.reshape(B, T, h, hs).transpose(0, 2, 1, 3)
Vh = V.reshape(B, T, h, hs).transpose(0, 2, 1, 3)
scores = Qh @ Kh.transpose(0, 1, 3, 2) / np.sqrt(hs)
A = softmax(scores, axis=-1)
Oh = A @ Vh
concat = Oh.transpose(0, 2, 1, 3).reshape(B, T, C)
out = concat @ self.Wo
self.cache = (X, Qh, Kh, Vh, A, concat)
return out
def backward(self, dout):
X, Qh, Kh, Vh, A, concat = self.cache
B, T, C = X.shape
h, hs = self.h, self.hs
# out = concat @ Wo
self.dWo = np.einsum('btc,btd->cd', concat, dout) # (C, C)
dconcat = dout @ self.Wo.T # (B, T, C)
# concat = merge(Oh): reverse the reshape/transpose
dOh = dconcat.reshape(B, T, h, hs).transpose(0, 2, 1, 3) # (B, h, T, hs)
# Oh = A @ Vh
dA = dOh @ Vh.transpose(0, 1, 3, 2) # (B, h, T, T)
dValueH = A.transpose(0, 1, 3, 2) @ dOh # (B, h, T, hs)
# A = softmax(scores), row-wise
dS = A * (dA - (dA * A).sum(axis=-1, keepdims=True)) # (B, h, T, T)
# scores = Qh Kh^T / sqrt(hs)
dscores = dS / np.sqrt(hs) # (B, h, T, T)
dQueryH = dscores @ Kh # (B, h, T, hs)
dKeyH = dscores.transpose(0, 1, 3, 2) @ Qh # (B, h, T, hs)
# merge heads back to full width (B, T, C)
dQuery = dQueryH.transpose(0, 2, 1, 3).reshape(B, T, C)
dKey = dKeyH.transpose(0, 2, 1, 3).reshape(B, T, C)
dValue = dValueH.transpose(0, 2, 1, 3).reshape(B, T, C)
# Q, K, V = X @ Wq, X @ Wk, X @ Wv
self.dWq = np.einsum('btc,btd->cd', X, dQuery) # (C, C)
self.dWk = np.einsum('btc,btd->cd', X, dKey)
self.dWv = np.einsum('btc,btd->cd', X, dValue)
dX = dQuery @ self.Wq.T + dKey @ self.Wk.T + dValue @ self.Wv.T
return dX
np.random.seed(42)
B, T, C, h = 2, 6, 8, 4
X = np.random.randn(B, T, C).astype(np.float64)
layer = MultiHeadAttention(C, h)
out = layer.forward(X)
dout = np.random.randn(*out.shape) # a pretend gradient from above
dX = layer.backward(dout)
print("dout.shape :", dout.shape)
print("dWq/dWk/dWv/dWo :", layer.dWq.shape, layer.dWk.shape, layer.dWv.shape, layer.dWo.shape)
print("dX.shape :", dX.shape)dout.shape : (2, 6, 8)
dWq/dWk/dWv/dWo : (8, 8) (8, 8) (8, 8) (8, 8)
dX.shape : (2, 6, 8)The shapes are the first sanity check, and they all match: a gradient always has the same shape as the thing it belongs to. Each of dWq, dWk, dWv, dWo is (8, 8) like its weight, and dX is (2, 6, 8) like the input. The two reshape/transpose pairs are worth reading slowly. The forward concat merged heads with transpose(0, 2, 1, 3).reshape(B, T, C); the backward dOh undoes it with the mirror reshape(B, T, h, hs).transpose(0, 2, 1, 3). Getting that mirror exactly right — same axis order, same reshape target — is the single most error-prone part of multi-head backprop, and shapes matching is necessary but not sufficient to know it is correct. A swapped transpose can still produce the right shape while scrambling which head owns which gradient. That is why the next stage proves correctness numerically.
The split and merge must be exact mirrors
Forward splits with reshape(B, T, h, hs).transpose(0, 2, 1, 3); backward merges with transpose(0, 2, 1, 3).reshape(B, T, C); and each is the other run in reverse. If you split one way in forward but merge a different way in backward, gradients from head 1 can silently land on head 2’s weights. The shapes still line up, so nothing crashes — the loss just never really drops. This class of bug is exactly what a gradient check catches, which is why we never ship a backward pass without one.
Stage 4: Verify the Class Is Correct
A backward pass you cannot verify is a backward pass you cannot trust. We prove this one two ways: a gradient check (do the analytic gradients match finite differences?) and a learning check (does a real loss go down when we descend those gradients?).
Gradient Check
The idea is the definition of a derivative. Nudge one entry of a parameter by , measure how much a scalar loss changes, and divide — that finite difference must match the analytic gradient:
We use a simple scalar loss for a fixed random matrix of shape (B, T, C), because then the analytic gradient flowing into backward is exactly dout = G. We run the check in float64 (finite differences are numerically delicate; single precision would drown the signal in rounding noise) and report the max relative error across every entry of all four weights and X. The snippet below assumes the MultiHeadAttention class from Stage 3 is already defined.
def scalar_loss(out, G):
return np.sum(out * G)
def rel_err(a, b):
return np.max(np.abs(a - b) / np.maximum(1e-12, np.abs(a) + np.abs(b)))
def numeric_grad(layer, X, P, G, eps=1e-6):
"""Central-difference gradient of scalar_loss w.r.t. every entry of array P."""
g = np.zeros_like(P)
it = np.nditer(P, flags=['multi_index'])
while not it.finished:
idx = it.multi_index
old = P[idx]
P[idx] = old + eps; lp = scalar_loss(layer.forward(X), G)
P[idx] = old - eps; lm = scalar_loss(layer.forward(X), G)
P[idx] = old
g[idx] = (lp - lm) / (2 * eps)
it.iternext()
return g
np.random.seed(42)
B, T, C, h = 2, 6, 8, 4
X = np.random.randn(B, T, C).astype(np.float64)
layer = MultiHeadAttention(C, h)
G = np.random.randn(B, T, C) # fixed upstream gradient
layer.forward(X)
dX = layer.backward(G) # analytic dWq/dWk/dWv/dWo and dX
num_Wq = numeric_grad(layer, X, layer.Wq, G)
num_Wk = numeric_grad(layer, X, layer.Wk, G)
num_Wv = numeric_grad(layer, X, layer.Wv, G)
num_Wo = numeric_grad(layer, X, layer.Wo, G)
num_X = numeric_grad(layer, X, X, G)
layer.forward(X); dX = layer.backward(G) # recompute analytic dX after perturbing X
print("max rel err Wq:", rel_err(layer.dWq, num_Wq))
print("max rel err Wk:", rel_err(layer.dWk, num_Wk))
print("max rel err Wv:", rel_err(layer.dWv, num_Wv))
print("max rel err Wo:", rel_err(layer.dWo, num_Wo))
print("max rel err X:", rel_err(dX, num_X))max rel err Wq: 2.51496403314037e-07
max rel err Wk: 7.932527961277073e-07
max rel err Wv: 3.2120264580953197e-10
max rel err Wo: 8.446283742081068e-09
max rel err X: 1.2461196832591731e-08Every max relative error is around or far smaller — the largest, on Wk, is 7.9e-07, comfortably under the 1e-5 bar for “the backward pass is correct.” Read the sizes and they tell the same story as the single head did: dWo (8.4e-09) and dWv (3.2e-10) are the most accurate because they flow through the fewest operations, while dWq and dWk accumulate a little more finite-difference noise because their gradients thread all the way back through the softmax in every head. All five are tiny, which is the numerical proof that backward — split, per-head attention, merge, and all — implements the true derivative of forward.
It Learns
Matching finite differences proves the gradients are correct; the final check proves they are useful. We pick a fixed random target Y of shape (B, T, C), define a mean-squared-error loss between the layer’s output and Y, and take ten gradient-descent steps — each a forward, a backward, and a subtraction of lr times each of the four parameter gradients. Because the class’s default * 0.02 init makes the doubly-projected output tiny (it passes through both an input projection and Wo), we reinitialize the weights at a larger * 0.3 scale here so there is real signal to descend on; the grad check above used the small default. If the gradients truly point downhill, the printed loss must fall.
np.random.seed(0)
B, T, C, h = 2, 6, 8, 4
X = np.random.randn(B, T, C)
Y = np.random.randn(B, T, C) # a fixed random target to fit
mha = MultiHeadAttention(C, h)
rng = np.random.default_rng(7) # larger init so the output has signal
mha.Wq = rng.standard_normal((C, C)) * 0.3
mha.Wk = rng.standard_normal((C, C)) * 0.3
mha.Wv = rng.standard_normal((C, C)) * 0.3
mha.Wo = rng.standard_normal((C, C)) * 0.3
lr = 1.0
for step in range(10):
o = mha.forward(X)
loss = 0.5 * np.mean((o - Y) ** 2)
dout = (o - Y) / o.size # gradient of the MSE loss
mha.backward(dout)
mha.Wq -= lr * mha.dWq
mha.Wk -= lr * mha.dWk
mha.Wv -= lr * mha.dWv
mha.Wo -= lr * mha.dWo
print(f"step {step:2d} loss {loss:.4f}")step 0 loss 0.6508
step 1 loss 0.5570
step 2 loss 0.5173
step 3 loss 0.4907
step 4 loss 0.4696
step 5 loss 0.4511
step 6 loss 0.4340
step 7 loss 0.4174
step 8 loss 0.4009
step 9 loss 0.3840The loss falls monotonically from 0.6508 to 0.3840 across ten steps — every single step lower than the last. That is the whole training loop of a real network in miniature: forward, compute a loss, backward, nudge all four parameters against their gradients, repeat. A small attention layer cannot perfectly fit a random target, so the loss keeps descending rather than snapping to zero, but the steady, gap-free descent is exactly the behavior correct gradients produce. Run this whole stage a second time: because every random draw is seeded, you get byte-for-byte identical errors and losses — this course has no nondeterminism to hide behind.
Practice Exercises
Try these before checking the hints. They reuse the MultiHeadAttention class and the setup you built above.
Exercise 1: Gradient-Check a Bigger Layer
Rerun the gradient check with a larger configuration — say B, T, C, h = 3, 5, 16, 8 (so hs = 2). Confirm the max relative errors for all four weights and X are still under 1e-5. A correct backward pass should not care about the sizes or the number of heads.
# Your code here: rebuild X, layer, G with the larger shapes and re-run numeric_gradHint
Set B, T, C, h = 3, 5, 16, 8, then X = np.random.randn(B, T, C).astype(np.float64), layer = MultiHeadAttention(C, h), and G = np.random.randn(B, T, C). Call layer.forward(X) and layer.backward(G), then reuse numeric_grad and rel_err exactly as in Stage 4 for Wq, Wk, Wv, Wo, and X. Because the class never bakes in B, T, or even h, every error should again come out around 1e-7 or smaller. If any error blows up, the bug is a shape or head-count assumption in backward, not in your new sizes.
Exercise 2: Show That Wo Does Real Work
In the “it learns” loop, comment out the update to Wo (keep updating Wq, Wk, Wv) and rerun. Does the loss still fall, and does it reach as low as 0.3840?
# Your code here: update only Wq, Wk, Wv each step; leave Wo frozenHint
Remove the line mha.Wo -= lr * mha.dWo so Wo stays at its initial value. The loss should still decrease — the three projections still reshape what each head computes — but it will not fall as far, because Wo is the only weight that can mix and rescale the concatenated heads to match the target’s full width. Freezing it removes the layer’s final degree of freedom. This is a concrete demonstration that the output projection is not decorative: it is what lets the heads’ separate answers be combined into one useful output.
Exercise 3: Confirm One Head Recovers Module 2
Set h = 1 and check that the layer still gradient-checks. With a single head, hs = C, there is no real splitting, and Wq, Wk, Wv are full (C, C) matrices — structurally the single-head attention of Module 2 plus an output projection Wo. Verify the max relative errors stay under 1e-5.
# Your code here: build MultiHeadAttention(C, h=1) and run the Stage 4 gradient checkHint
Use B, T, C, h = 2, 6, 8, 1, so hs = 8. The reshape(B, T, 1, hs).transpose(0, 2, 1, 3) still runs — it just inserts a length-1 head axis and does no real interleaving — so the same code path works untouched. Run the Stage 4 gradient check and every error should again be around 1e-7 or smaller. This confirms multi-head attention is a strict generalization: one head is simply the special case h = 1, and the machinery you built handles it without any special-casing.
Summary
You turned a module’s worth of loose steps into one reusable, provably correct component. Let’s review what you built.
Key Concepts
The Class Design
__init__(C, h)stores the shape (C,h,hs = C // h) and four(C, C)matrices:Wq,Wk,Wv, and the output projectionWo, initialized small (* 0.02) and seeded- The layer commits only to widths (
C,h), never toBorT, so the same object works on any batch size or sequence length - Three kinds of state live on
self: parameters, a forwardcache, and gradients
forward and backward
forward(X)projects to full-widthQ,K,V, splits into(B, h, T, hs), attends per head, concatenates to(B, T, C), and appliesWo, caching(X, Qh, Kh, Vh, A, concat)backward(dout)reverses each step:dWoanddconcat, the head split ofdconcatintodOh, per-headdValueH/softmax-backwarddS/dQueryH/dKeyH, the merge into full-widthdQuery/dKey/dValue, the parameter gradsdWq/dWk/dWv, and the returneddX- The forward split and the backward merge must be exact mirrors, or gradients land on the wrong head
Proving Correctness
- A
float64central-difference gradient check gave max relative errors from7.9e-07(Wk) down to3.2e-10(Wv) — all under the1e-5bar - A ten-step gradient-descent run drove a real MSE loss monotonically from 0.6508 to 0.3840, confirming the gradients point downhill
- Everything is seeded and byte-for-byte reproducible on a second run
Why This Matters
This class is the centerpiece you will genuinely reuse rather than rebuild. Module 5 drops one MultiHeadAttention inside a transformer block alongside a feed-forward network and layer norm; Module 6 adds a causal mask to forward so each position attends only to earlier ones; Module 7 stacks several of these and trains them into a working char-level GPT on the Lantern Bay corpus. Because you gradient-checked the whole layer here — split, per-head attention, concatenation, and output projection — none of those modules has to re-derive or re-verify the attention math. They inherit a block already known to be correct.
More broadly, the forward-caches / backward-reads-the-cache pattern is the shape of every layer in every deep-learning framework, and multi-head attention is the most intricate version of it you will build before the full model. You now know exactly what a real attention Module is hiding: four weight matrices, a reshape that fans work out across heads, its exact mirror that folds the gradients back, and a cache that ties the two together. That understanding is what lets you debug a real transformer instead of treating it as a black box.
Continue Building Your Skills
You now hold a single, trustworthy multi-head attention layer packaged as a class — the workhorse of every model ahead, and the exact object you will stack in Module 5. But there is a gap in it that no amount of attention math can fix: the layer treats a sequence as a set. Shuffle the tokens of the input and the attention weights permute along with them, because nothing you have built so far tells the model where each token sits. In the next module, Positional Encoding, you will close that gap — adding a signal that stamps each position with its place in the sequence so “the boat drifts near the rocks” and “the rocks drift near the boat” stop looking identical to the model. Keep this gradient-checked class close; positional encoding is the piece that finally lets it read order, and the two together are everything the transformer block needs.