Let every token look back over the whole context.
Built in four small steps, each a tiny change to the one before -- from "average the past" to full scaled, causal, query-key-value attention.
The bigram saw a single character. A Transformer lets every token look back over all the earlier tokens and pull in whatever is relevant. That mechanism is self-attention, and it is the one idea that makes language models work.
Here is why it matters. When you read "the trophy did not fit in the suitcase because it was too big," your eye glances back to work out what "it" refers to. The bigram cannot do that, because it sees only the character immediately before. Attention is exactly that glance backward, turned into arithmetic. We will build it in four small steps, each a small change to the one before.
import numpy as np
import matplotlib.pyplot as plt
np.random.seed(1337)
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)
# a toy batch: B sequences, T tokens each, C features per token
B, T, C = 4, 8, 32
x = np.random.randn(B, T, C)
print("input:", x.shape, " (batch, time, channels)")Line by line: what each line does
import numpy as np; np.random.seed(1337): the array toolbox, with randomness pinned for repeatable numbers.B, T, C = 4, 8, 32: three sizes: a batch of 4 sequences, T=8 tokens each, C=32 numbers describing each token.x = np.random.randn(B, T, C): a 4×8×32 block of random numbers standing in for token embeddings. They're random because here we only care about the mechanics of mixing them.def softmax(x, axis=-1): ...: the same probabilities helper as always, on one line.print("input:", x.shape, ...): confirm the (batch, time, channels) shape we'll carry through the whole notebook.
Step 1: the simplest way to gather context, averaging the past
(First, a note on shapes, because everything from here on uses them. B is how many sequences we handle at once, T is how many tokens each sequence has, and C is how many numbers describe each token. So x is a block of numbers with shape (B, T, C).)
What is the simplest way for token t to use its context? Simply average it with every token before it, giving a blurred summary of the story so far. One firm rule: a token may never look at the future. When the model is generating text, the future has not been written yet (predicting it is the model's whole job), so during training we forbid looking ahead as well. This "only look backward" rule is called causality, and we enforce it everywhere.
The double loop below makes the averaging obvious: token 0 is just itself, token 2 is the average of tokens 0, 1, and 2, and so on.
xbow = np.zeros_like(x) # "bag of words" = running average
for b in range(B):
for t in range(T):
xbow[b, t] = x[b, :t+1].mean(axis=0) # mean of tokens 0..t
print("token 0 is itself: ", np.allclose(xbow[0,0], x[0,0]))
print("token 2 is mean of 0,1,2:", np.allclose(xbow[0,2], x[0,:3].mean(0)))Line by line: what each line does
xbow = np.zeros_like(x): an all-zeros block the same shape asx, to fill in. ("bow" = bag-of-words, i.e. a running average.)for b in range(B): for t in range(T):: visit every sequenceb, every positiont.x[b, :t+1]: slice along time: tokens 0 through t (the+1because slices exclude their end). Only the past and the current token -- never the future..mean(axis=0): average those vectors across time, giving one blended 32-number summary for position t.np.allclose(...): a sanity check that token 0's summary is just itself, and token 2's is the mean of tokens 0,1,2.
Step 2: the same thing as a matrix multiply
That averaging loop is secretly a single matrix multiply, and seeing this is what unlocks attention. Stack the averaging weights into a (T, T) table called wei, where row t says how much token t draws from each earlier token. The whole gather then becomes wei @ x.
Look at the printed table: row 2 is [0.33, 0.33, 0.33, 0, 0, ...], an equal third from tokens 0, 1, and 2 and nothing from the future. That lower-triangular shape (zeros above the diagonal, so each row reaches only backward, never forward) is what enforces causality, at no extra cost.
The key reframing is this: "who attends to whom" is just a (T, T) grid of weights. Right now every earlier token gets an equal share. The entire rest of attention is one upgrade: stop averaging equally, and let the data decide those weights.
The shape that becomes attention
Before real attention, the simplest way for a token to use its context is to average every token up to it. Stack those averaging weights into a table and the whole gather is one matrix multiply, wei @ x. The lower-triangular shape — and nothing above the diagonal — is exactly what makes it causal. Click a row.
tril = np.tril(np.ones((T, T))) # 1s on/below the diagonal
wei = tril / tril.sum(axis=1, keepdims=True) # normalize each row to average
xbow2 = wei @ x # (T,T) @ (B,T,C) -> (B,T,C) by broadcasting
print("matches the loop:", np.allclose(xbow, xbow2))
print("\nweight matrix (row t = how token t mixes the past):")
print(np.round(wei, 2))Line by line: what each line does
tril = np.tril(np.ones((T, T))): an 8×8 grid of 1s with everything above the diagonal zeroed ("tril" = lower triangle). Row t has 1s only in columns 0..t -- i.e. "position t may look at 0..t."wei = tril / tril.sum(axis=1, keepdims=True): divide each row by how many 1s it has, turning it into equal shares that sum to 1.xbow2 = wei @ x: the punchline: multiplying by that (T,T) weight grid computes every position's average-of-the-past in a single matrix multiply. NumPy broadcasts it across the batch.np.allclose(xbow, xbow2): proves the matmul gives exactly what the double loop gave.np.round(wei, 2): print the weight grid so you can see row t mixing tokens 0..t equally, future columns at 0.
Look at the averaging matrix
The cell above printed wei as numbers. Seen as a picture it is clearer: a lower-triangular staircase where row t spreads an equal share across tokens 0..t and puts nothing on the future. This is the crude, fixed version of the weights that attention will soon learn instead.
# The (T, T) averaging weights as an annotated heatmap.
fig, ax = plt.subplots(figsize=(5.2, 4.4))
im = ax.imshow(wei, cmap="Blues")
for i in range(T):
for j in range(T):
if wei[i, j] > 0:
ax.text(j, i, f"{wei[i, j]:.2f}", ha="center", va="center",
fontsize=7, color="white" if wei[i, j] > 0.5 else "#333")
ax.set_xlabel("past token j"); ax.set_ylabel("current token t")
ax.set_title("step-2 averaging: an equal share of the past")
fig.colorbar(im); plt.show()Line by line: what each line does
ax.imshow(wei, cmap="Blues"): draw the 8×8 weight table as an image; darker means a larger share.- The double loop writes each nonzero weight into its cell, so you can read the exact numbers off the picture.
- Row 0 is a lone
1.00(token 0 has only itself); row 1 is0.50, 0.50; row 7 is eight equal0.12s. - The blank upper triangle is the causal rule: no token draws anything from the future. Attention keeps this exact triangular shape but replaces the equal shares with learned ones.
Step 3: let the data set the weights, which is attention itself
Equal averaging is crude: a pronoun should lean heavily on its noun, not on every random word. So we let the data set the weights. The clearest way to picture this is a search engine, in which each token produces three things:
- a query
q, meaning "what am I looking for?" (like the words you type into a search box), - a key
k, meaning "what do I contain?" (like the tags a document advertises), - a value
v, meaning "what I will pass on if you attend to me" (like the document's actual content).
These are three different jobs, so each gets its own learned projection (W_q, W_k, W_v, which are just matrices the model trains). To score how much token i should attend to token j, take the dot product q_i . k_j, the query-meets-key similarity from notebook 01. A high dot product means a good match, so the token pays more attention there.
The rest is machinery you already know: hide the future by setting those scores to negative infinity, apply softmax to each row so the weights add up to 1, and return the weighted sum of the values. The whole thing is one compact formula:
$$\text{attention}(Q,K,V) = \text{softmax}\!\left(\frac{Q K^{\top}}{\sqrt{d}}\right) V$$
In plain words, reading from the inside out: match every query against every key ($QK^{\top}$, where $K^{\top}$ just means the table of keys turned on its side so the multiplication lines up), tone the scores down by dividing by $\sqrt{d}$ (explained in the next cell), turn them into weights with softmax, and use those weights to blend the values ($\cdot\, V$). Unlike the bigram, the token itself decides where to look.
head_size = 16
# learned projections (random here; in the GPT these are trained)
W_q = np.random.randn(C, head_size) * 0.1
W_k = np.random.randn(C, head_size) * 0.1
W_v = np.random.randn(C, head_size) * 0.1
q = x @ W_q # (B,T,head_size) - queries
k = x @ W_k # (B,T,head_size) - keys
v = x @ W_v # (B,T,head_size) - values
# affinities: every query dotted with every key
scores = q @ k.transpose(0, 2, 1) # (B,T,T)
scores = scores / head_size**0.5 # scale (step 4, below)
mask = np.tril(np.ones((T, T))) # causal mask
scores = np.where(mask == 0, -np.inf, scores) # future -> -inf -> 0 after softmax
attn = softmax(scores, axis=-1) # (B,T,T) each row sums to 1
out = attn @ v # (B,T,head_size) weighted sum of values
print("attention output:", out.shape)
print("causal? upper triangle is zero:", np.allclose(np.triu(attn[0], k=1), 0))
print("rows sum to 1:", np.allclose(attn[0].sum(1), 1))Line by line: what each line does
head_size = 16: each attention head will boil a token down to a 16-number query/key/value.W_q = np.random.randn(C, head_size) * 0.1(and W_k, W_v) -- the three learned projections. Multiplying a token's 32 numbers by a (32,16) grid gives its 16-number query / key / value. (Random here; in the GPT these are trained.)q = x @ W_q(and k, v) -- apply each projection to every token in every sequence at once: shape (B, T, 16).scores = q @ k.transpose(0, 2, 1):transposeswaps the last two axes so the matmul lines up; the result is every query dotted with every key -- a (B, T, T) grid of match scores.scores = scores / head_size**0.5: divide by √16 = 4: the scaling trick (next cell shows why).mask = np.tril(np.ones((T, T))): the same lower-triangle mask: 1 where attending is allowed, 0 for the future.scores = np.where(mask == 0, -np.inf, scores): wherever the mask is 0, overwrite the score with negative infinity; everywhere else keep it. After softmax, -inf becomes exactly 0 weight.attn = softmax(scores, axis=-1): turn each row of scores into weights that sum to 1.out = attn @ v: blend the value vectors by those weights: each token's output is its personalized weighted sum of the past.np.triu(attn[0], k=1): the strictly-upper triangle (future positions); checking it's all zeros proves no information leaked backward in time.
Open up one row of the attention matrix
attn is a (B, T, T) block, which is a lot to take in at once. So follow a single query token all the way through -- token t = 4 in the first sequence -- and watch its score against every other token become a set of weights. The recomputed weights match row 4 of the full attn, and their weighted sum of values matches out for that token, so this one row is the whole mechanism in miniature.
# Follow token t=4: query . keys -> scale -> causal mask -> softmax -> weights.
t = 4
qt = q[0, t] # this token's query vector (head_size,)
raw = (k[0] @ qt) / head_size**0.5 # score against every token's key (T,)
masked = raw.copy(); masked[t+1:] = -np.inf # a token may not look at the future
w = softmax(masked) # weights over tokens 0..t (future = 0)
print(f"token t={t}: query shape {qt.shape}, scored against {T} keys")
print("raw scores q.k/sqrt(hs):", np.round(raw, 2))
print("after causal mask :", np.round(masked, 2))
print("after softmax (weights) :", np.round(w, 3), " sum =", round(w.sum(), 3))
print("matches row 4 of attn :", np.allclose(w, attn[0, t]))
print("blended value == out[0,4]:", np.allclose(w @ v[0], out[0, t]))
fig, ax = plt.subplots(figsize=(6, 2.6))
colors = ["#5b5bd6" if j <= t else "#d9d9e0" for j in range(T)]
ax.bar(range(T), w, color=colors)
ax.set_xticks(range(T)); ax.set_xlabel("token j (grey = future, masked out)")
ax.set_ylabel("attention weight"); ax.set_title(f"how token {t} splits its attention over the past")
plt.show()Line by line: what each line does
qt = q[0, t]: pull out token 4's query vector (16 numbers).raw = (k[0] @ qt) / head_size**0.5: dot that query against every token's key, then apply the step-4 scale -- one score per token.masked[t+1:] = -np.inf: blank out tokens 5, 6, 7 (the future); softmax turns-infinto exactly 0.w = softmax(masked): the five surviving scores become weights that sum to 1.- The two
allclosechecks prove this hand-trace equals row 4 of the fullattnmatrix, and that blending the values with these weights gives the sameoutthe batched code produced. - The bar chart shows where token 4 actually looks: taller bars are tokens it weights more, the greyed bars are the masked future.
Step 4: why divide by the square root of the head size?
One small detail keeps attention healthy at the start of training. A dot product adds up head_size separate little products (head_size is just the length of the query and key vectors), so with longer vectors the scores naturally grow larger. Feed large numbers into softmax and it becomes peaky, meaning almost all the weight piles onto a single token. That is a problem early on, because a peaky softmax produces very small gradients (it is already acting "certain"), so the model can barely learn, and what it is certain about is random nonsense, since nothing has been trained yet.
Dividing every score by the square root of head_size cancels that growth and keeps the scores moderate, so attention starts out soft, meaning spread out, and therefore trainable. The demo makes it concrete: with no scaling the top weight is already 0.87 (for head size 16) or 0.94 (for head size 256), which is dangerously peaky, while the scaled version sits at a gentle 0.32 or 0.24.
Who attends to whom
Causal self-attention for "The cat sat on the mat ." Each row is a token's query; each cell is how much it attends to an earlier token (rows sum to 100%, the future is masked grey). Hover a cell for the exact weight; drag temperature to sharpen or soften -- low temperature is the "un-scaled, peaky" regime the 1/sqrt(d) trick protects against.
# Dot products grow with head_size, which makes softmax collapse toward one-hot.
# Compare one query against 8 keys, for two head sizes, with and without the scale:
rng = np.random.default_rng(0)
for hs in [16, 256]:
qq = rng.standard_normal(hs); kk = rng.standard_normal((8, hs))
s = kk @ qq
print(f"head_size={hs:4d}: unscaled max weight = {softmax(s).max():.2f} "
f"scaled max weight = {softmax(s / hs**0.5).max():.2f}")
print("\nBigger head_size -> unscaled softmax gets peaky (toward 1.0, starving gradients);")
print("the 1/sqrt(head_size) scale keeps it soft and trainable.")Line by line: what each line does
rng = np.random.default_rng(0): a seeded random generator so the demo numbers are repeatable.for hs in [16, 256]:: try a small and a large head size.qq = rng.standard_normal(hs); kk = rng.standard_normal((8, hs)): one random query and 8 random keys of that size.s = kk @ qq: the 8 match scores. With more dimensions added up, the scores naturally swing bigger.softmax(s).max(): the largest attention weight without scaling. Near 1.0 means all the attention piled onto one token (peaky -- bad early in training).softmax(s / hs**0.5).max(): the same scores divided by √(head size) first: attention stays spread out and trainable.
# visualize who attends to whom (row = query token, col = key token)
fig, ax = plt.subplots(figsize=(5,4))
im = ax.imshow(attn[0], cmap="viridis")
ax.set_xlabel("key (past tokens)"); ax.set_ylabel("query (current token)")
ax.set_title("attention weights (lower-triangular = causal)")
fig.colorbar(im); plt.show()Line by line: what each line does
fig, ax = plt.subplots(figsize=(5,4)): make a single plot panel to draw the heatmap on.ax.imshow(attn[0], cmap="viridis"): draw the first sequence's (T,T) attention grid as colors: brighter = more attention.ax.set_xlabel(...); ax.set_ylabel(...): label the axes: columns are the past tokens being looked at, rows are the token doing the looking.fig.colorbar(im); plt.show(): add the color scale and render. The bright lower triangle and dark upper triangle is causality, made visible.
Package it as a function
Those few lines are one attention head: query, key, value, score, mask, softmax, blend. Here it is wrapped into a single function that returns exactly the same output. We will reuse this precise logic in PyTorch later, where it shrinks to about six lines and the gradients are handled for us.
def attention_head(x, W_q, W_k, W_v):
B, T, C = x.shape; hs = W_q.shape[1]
q, k, v = x @ W_q, x @ W_k, x @ W_v
scores = (q @ k.transpose(0,2,1)) / hs**0.5
scores = np.where(np.tril(np.ones((T,T))) == 0, -np.inf, scores)
return softmax(scores, -1) @ v
out2 = attention_head(x, W_q, W_k, W_v)
print("output:", out2.shape, " same as before:", np.allclose(out, out2))Line by line: what each line does
def attention_head(x, W_q, W_k, W_v):: the previous cells folded into one reusable function.B, T, C = x.shape: read the input's three sizes back out (tuple unpacking).hs = W_q.shape[1]: read the head size off the projection grid instead of hard-coding it.return softmax(scores, -1) @ v: the same project → score → scale → mask → softmax → blend, compressed.np.allclose(out, out2): the function reproduces the step-by-step result exactly.
Recap
Self-attention in one sentence: each token builds query, key, and value vectors, scores every earlier token by query-key similarity, hides the future, turns the scores into weights with softmax, and returns the weighted sum of the values. One head mixes information in one particular way, for example "pronouns, find your nouns."
Next, notebook 06 runs several heads in parallel (so the model can look for several kinds of relationship at once), adds position information and a small feed-forward layer, and wraps it all into the Transformer block that we stack to build a GPT.
05_attention.ipynb