A GPT is just this block, repeated.
Multi-head attention, a feed-forward network, LayerNorm, residual connections, and positional encoding -- assembled into the block that stacks into a GPT.
A GPT is just a stack of identical Transformer blocks, one on top of the next. You can picture each block as a short team meeting with two phases:
- multi-head self-attention, where everyone shares information (each token gathers context from the others),
- a feed-forward network, where everyone goes back to their desk and thinks it over alone.
Two supports keep the meeting orderly: residual connections (keep the original notes and just add edits on top) and LayerNorm (keep everyone's numbers on a similar scale). And because attention on its own has no sense of word order, we first stamp each token with its position. We will build every piece in NumPy and run one full block's forward pass.
One block = communicate, then compute
A GPT is just this one block, stacked. Watch a token ride the residual stream down the middle: attention lets tokens mix, a feed-forward lets each token think, and each result is added back onto the stream rather than replacing it.
import numpy as np
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)
B, T, C = 4, 16, 64 # batch, time, embedding dim
n_head = 4Line by line: what each line does
import numpy as np; np.random.seed(1337): the array toolbox, randomness pinned.def softmax(x, axis=-1): ...: the usual scores-to-probabilities helper.B, T, C = 4, 16, 64: 4 sequences, 16 tokens each, 64 numbers per token.n_head = 4: we'll split those 64 numbers across 4 attention heads (16 each).
Positional encoding
Attention is order-blind. It is just a weighted sum, so if you shuffled the tokens it would return the same outputs (shuffled the same way); it has no idea which token came first. But order carries meaning, since "dog bites man" and "man bites dog" are very different. So before anything else we stamp each token with its position, adding a position signature onto its embedding (the embedding being the list of numbers that represents the token).
The classic method uses sinusoids, which are simply waves of different speeds. Think of a car's mileage counter: a fast-spinning ones wheel, a slower tens wheel, and a slower-still hundreds wheel. Read together, the mix of fast and slow wheels gives every position a unique fingerprint, and nearby positions get similar fingerprints. (A GPT actually learns its own table of position signatures instead of using fixed waves, but both achieve the same thing.)
Every position gets a fingerprint
Attention is order-blind, so before the blocks we stamp each token with where it sits. The classic way mixes sine waves of many speeds — like a car's mileage wheels, a fast ones-wheel and slower tens- and hundreds-wheels. Read together they give every position a unique signature, and nearby positions similar ones.
def positional_encoding(T, C):
pos = np.arange(T)[:, None]
i = np.arange(C)[None, :]
angle = pos / np.power(10000, (2 * (i // 2)) / C)
pe = np.where(i % 2 == 0, np.sin(angle), np.cos(angle))
return pe # (T, C)
pe = positional_encoding(T, C)
import matplotlib.pyplot as plt
plt.figure(figsize=(6,3)); plt.imshow(pe, aspect="auto", cmap="RdBu")
plt.xlabel("embedding dim"); plt.ylabel("position"); plt.title("sinusoidal positional encoding"); plt.colorbar(); plt.show()Line by line: what each line does
def positional_encoding(T, C):: build a unique "where am I" fingerprint for each of the T positions.pos = np.arange(T)[:, None]: positions 0..15 as a column (the[:, None]adds an axis so it can broadcast against a row).i = np.arange(C)[None, :]: the embedding-slot numbers 0..63 as a row.angle = pos / np.power(10000, (2 * (i // 2)) / C): the wave math.i // 2is whole-number division, pairing slots (0,0,1,1,...); each pair gets its own wave speed -- smalli→ divisor near 1 → fast wave, largei→ divisor near 10000 → very slow wave. Fast and slow wheels, like an odometer.np.where(i % 2 == 0, np.sin(angle), np.cos(angle)): even slots take the sine of their angle, odd slots the cosine (%is remainder). Together every position gets a one-of-a-kind pattern of wave readings.plt.imshow(pe, ...): draw it: rows are positions, columns are slots; you can see fast stripes on the left, slow on the right.
LayerNorm
As signals pass through many layers, the numbers can drift, either ballooning very large or shrinking towards nothing, and training then falls apart. LayerNorm is the fix. For each token it takes that token's list of numbers and rescales it to have an average of 0 and a standard deviation of 1. (Standard deviation is just a measure of how spread out a set of numbers is; forcing it to 1 keeps the spread consistent, rather like grading on a curve.) It then applies two learned dials, gamma and beta, so the model can stretch or shift the result if that helps.
The cell shows it working: numbers that started with an average of 0.44 and a spread of 5.02 come out at an average of 0 and a spread of 1. This is done for each token on its own and has nothing to do with the rest of the batch, which is what separates it from an older method called BatchNorm, and why it behaves the same whether you process one sequence or a thousand.
def layer_norm(x, gamma, beta, eps=1e-5):
mu = x.mean(axis=-1, keepdims=True)
var = x.var(axis=-1, keepdims=True)
return gamma * (x - mu) / np.sqrt(var + eps) + beta
z = np.random.randn(2, 5) * 5 + 3
zn = layer_norm(z, 1.0, 0.0)
print("before: mean=%.2f std=%.2f" % (z.mean(), z.std()))
print("after : mean=%.2f std=%.2f (per-row normalized)" % (zn.mean(), zn.std()))Line by line: what each line does
def layer_norm(x, gamma, beta, eps=1e-5):: rescale each token's vector to a standard size so deep stacks stay stable.mu = x.mean(axis=-1, keepdims=True): each token-vector's average (across its own features).var = x.var(axis=-1, keepdims=True): its variance: how spread out those numbers are.(x - mu) / np.sqrt(var + eps): the recipe in words: shift so the average is 0, divide by the spread so it becomes 1.eps(a tiny 0.00001) is insurance against dividing by zero if a vector is perfectly flat.gamma * ... + beta: two learned dials to scale and shift the normalized result. We pass 1.0 and 0.0 ("leave it as is"); a real model learns them.- The test: numbers built with mean 3 and spread 5 come out at mean 0, spread 1 -- row by row.
Watch LayerNorm flatten a token
LayerNorm's job is one sentence: take a token's row of numbers and rescale it to mean 0, standard deviation 1. The cell above printed the overall stats; here it is per row, before and after, with row 0 drawn as bars. However spread out a row starts, it lands on the same scale.
# LayerNorm rescales EACH row to mean 0, std 1 -- per row, then row 0 as bars.
for r in range(z.shape[0]):
print(f"before row {r}: mean={z[r].mean():+.2f} std={z[r].std():.2f}")
for r in range(zn.shape[0]):
print(f"after row {r}: mean={zn[r].mean():+.2f} std={zn[r].std():.2f}")
fig, ax = plt.subplots(1, 2, figsize=(7, 2.6))
idx = np.arange(z.shape[1])
ax[0].bar(idx, z[0], color="#d9822b"); ax[0].axhline(0, color="k", lw=.5)
ax[0].set_title("row 0 before (mean %.1f, std %.1f)" % (z[0].mean(), z[0].std()))
ax[1].bar(idx, zn[0], color="#2b8dd9"); ax[1].axhline(0, color="k", lw=.5)
ax[1].set_title("row 0 after (mean 0, std 1)")
plt.tight_layout(); plt.show()Line by line: what each line does
- The two loops print each row's mean and standard deviation, first for the raw
z, then for the normalizedzn. - Row 1 starts wildly spread (std over 6); after LayerNorm every row reads mean 0.00, std 1.00.
- The left bars are row 0 raw -- off-centre and uneven; the right bars are the same row after -- centred on 0 with a consistent spread.
gammaandbeta(here 1 and 0) then let the model rescale and shift if it wants, but the starting point is always this clean, fixed scale.
Multi-head attention
A single attention head can only look for one kind of relationship at a time. Multi-head attention runs several heads in parallel so the model can look for several at once: perhaps one head links each verb to its subject, another tracks quotation marks, and another watches for the previous line break. It is like a panel of specialists, each reading the sentence for their own purpose and then pooling their notes.
Mechanically, we slice the embedding into n_head equal pieces, run a separate attention head (exactly the one from notebook 05) on each piece in parallel, glue the results back together (concatenate them), and pass them through one more projection Wo that lets the heads mix. The output has the same shape as the input, so it drops straight into the block.
Split into heads, attend in parallel
One head can track only one kind of relationship at a time. Multi-head attention cuts each token's vector into equal heads and runs a separate attention on each in parallel — a panel of specialists — then glues the results back together and mixes them.
def multi_head_attention(x, Wq, Wk, Wv, Wo, n_head):
B, T, C = x.shape
hs = C // n_head
def proj(W): return (x @ W).reshape(B, T, n_head, hs).transpose(0, 2, 1, 3) # (B,nh,T,hs)
q, k, v = proj(Wq), proj(Wk), proj(Wv)
scores = (q @ k.transpose(0, 1, 3, 2)) / hs**0.5 # (B,nh,T,T)
scores = np.where(np.tril(np.ones((T, T))) == 0, -np.inf, scores)
attn = softmax(scores, axis=-1)
out = attn @ v # (B,nh,T,hs)
out = out.transpose(0, 2, 1, 3).reshape(B, T, C) # concat heads
return out @ Wo # output projection
Wq = np.random.randn(C, C)*0.02; Wk = np.random.randn(C, C)*0.02
Wv = np.random.randn(C, C)*0.02; Wo = np.random.randn(C, C)*0.02
x = np.random.randn(B, T, C)
print("multi-head output:", multi_head_attention(x, Wq, Wk, Wv, Wo, n_head).shape)Line by line: what each line does
def multi_head_attention(x, Wq, Wk, Wv, Wo, n_head):: run several attention heads at once and combine them.hs = C // n_head: whole-number division: 64 features split across 4 heads = 16 each.def proj(W): return (x @ W).reshape(B, T, n_head, hs).transpose(0, 2, 1, 3): one helper used three times: project, thenreshapecarves the 64-wide result into 4 heads of 16, andtransposereorders the axes to (batch, head, time, hs) so every head can work independently.q, k, v = proj(Wq), proj(Wk), proj(Wv): queries, keys, values, all carved into heads.scores = (q @ k.transpose(0, 1, 3, 2)) / hs**0.5: notebook 05's scaled query-key scores, now with a head axis riding along: 4 separate (T,T) grids per sequence.scores = np.where(np.tril(...) == 0, -np.inf, scores): the causal mask, applied to every head at once.out = attn @ vthen.transpose(...).reshape(B, T, C): blend values per head, then undo the carving so the 4 heads' 16-number outputs sit side by side as 64 again ("concatenate").return out @ Wo: one final projection lets the heads' findings mix with each other.
How C splits into heads
The reshape inside multi_head_attention is the one line that trips people up. All it does is cut each token's projected vector into n_head equal slices, run attention on every slice in parallel, then glue the results back and mix them with Wo. Here are the shapes at each step, and a picture of the split.
# Trace the shapes through one multi-head pass, then draw the head split.
hs = C // n_head
print(f"C = {C} -> {n_head} heads of head_size = {hs}\n")
print("x @ Wq :", (x @ Wq).shape, " (B, T, C)")
qh = (x @ Wq).reshape(B, T, n_head, hs).transpose(0, 2, 1, 3)
print("reshape + move heads up :", qh.shape, " (B, n_head, T, head_size)")
print("per-head attention q@k.T:", (B, n_head, T, T), " one T x T map per head")
print("glue heads back :", (B, T, C), " then a final @ Wo mixes them")
fig, ax = plt.subplots(figsize=(7, 1.5))
colors = plt.cm.Set2(np.linspace(0, 1, n_head))
for hi in range(n_head):
ax.barh(0, hs, left=hi * hs, color=colors[hi], edgecolor="white")
ax.text(hi * hs + hs / 2, 0, f"head {hi}\n[{hi*hs}:{hi*hs+hs}]", ha="center", va="center", fontsize=8)
ax.set_xlim(0, C); ax.set_ylim(-0.5, 0.5); ax.set_yticks([])
ax.set_xlabel(f"the {C} projected dims per token, sliced into {n_head} heads of {hs}")
ax.set_title("multi-head = split the projection into parallel heads")
plt.tight_layout(); plt.show()Line by line: what each line does
x @ Wq: project every token, still(B, T, C)-- nothing split yet..reshape(B, T, n_head, hs).transpose(...): cut the C numbers inton_headslices ofhead_size, and move the head axis up so each head is an independent(T, head_size)block.- Each head then runs the exact attention from notebook 05 on its own slice, giving one
T x Tmap per head. - The outputs are glued back into
(B, T, C)and passed throughWo, which lets the heads' findings mix. - The coloured bar is the split: dims 0-15 are head 0, 16-31 head 1, and so on -- parallel heads, each watching for a different kind of relationship.
Feed-forward network
Attention let the tokens talk to each other; now each token thinks on its own. The feed-forward network is a small two-layer MLP (multi-layer perceptron, from notebook 04) applied to every token independently. It expands the vector to width 4*C, four times wider, giving it room to compute richer features; passes it through a nonlinearity (here GELU, a smooth on/off gate that roughly keeps positive values and suppresses negative ones); and then projects it back down to width C.
That nonlinearity is essential. Without a bend in it, stacking layers would collapse into a single straight-line layer that cannot learn curves or rules. This sub-layer holds most of the model's numbers, and much of what we would loosely call its "knowledge" lives here.
def gelu(x):
return 0.5 * x * (1 + np.tanh(np.sqrt(2/np.pi) * (x + 0.044715 * x**3)))
def feed_forward(x, W1, b1, W2, b2):
return gelu(x @ W1 + b1) @ W2 + b2
W1 = np.random.randn(C, 4*C)*0.02; b1 = np.zeros(4*C)
W2 = np.random.randn(4*C, C)*0.02; b2 = np.zeros(C)
print("feed-forward output:", feed_forward(x, W1, b1, W2, b2).shape)Line by line: what each line does
def gelu(x):: the gentle on/off gate. The exact formula is a standard fast approximation; what matters is the shape: very negative inputs → near 0 (off), positive inputs → near x (passed through), with a smooth bend instead of a hard corner.def feed_forward(x, W1, b1, W2, b2):: the per-token "think about it" network.gelu(x @ W1 + b1) @ W2 + b2: the whole thing in one line: widen (64 → 256), bend with GELU, narrow back (256 → 64).b1/b2are bias vectors added to every token by broadcasting.W1 = np.random.randn(C, 4*C)*0.02; b1 = np.zeros(4*C): weights start small-random, biases start at zero.
See the GELU gate, and the width change
The feed-forward network does two things worth seeing. Every token passes through GELU, the "smooth on/off gate" the text mentions -- so here it is, plotted against the blunt ReLU. And it briefly widens each token from C to 4*C and back, which is where its extra room to compute lives.
# GELU vs ReLU, then the expand-then-contract widths of the feed-forward net.
xs = np.linspace(-4, 4, 200)
plt.figure(figsize=(6, 3))
plt.plot(xs, gelu(xs), label="GELU (used here)", color="#5b5bd6", lw=2)
plt.plot(xs, np.maximum(xs, 0), label="ReLU (for comparison)", color="#aaa", ls="--")
plt.axhline(0, color="k", lw=.5); plt.axvline(0, color="k", lw=.5)
plt.legend(); plt.title("GELU keeps positives, softly suppresses negatives")
plt.xlabel("input"); plt.ylabel("output"); plt.show()
print(f"feed-forward width per token: {C} -> {4*C} (expand + GELU) -> {C} (contract)")
print(f"W1: {W1.shape} W2: {W2.shape}")Line by line: what each line does
xs = np.linspace(-4, 4, 200): a range of inputs to feed the activation.plt.plot(xs, gelu(xs)): GELU's smooth S-curve -- it passes large positives through almost unchanged and eases negatives toward zero, with a soft dip near -1 rather than ReLU's hard corner.- The dashed ReLU (
max(x, 0)) is drawn for contrast: same idea, blunt instead of smooth. - The prints show the width change: each token goes
64 -> 256 -> 64. The wide middle (four times C) is the room the network computes richer features in before shrinking back.
Assemble the block (pre-norm, with residuals)
Now we wire the two sub-layers together. Modern GPTs use the pre-norm arrangement: apply LayerNorm first, then the sub-layer, then add the result back to the input. That "add back" is the residual (or skip) connection, and it is worth pausing on:
x = x + MultiHeadAttention(LayerNorm(x))
x = x + FeedForward(LayerNorm(x))Read each line as "keep x, and add an improvement on top." Each sub-layer only has to learn a small edit to the running representation, rather than rebuilding it from scratch. Just as important, the x + gives the gradients a clean route straight back through the whole stack during the backward pass; without residuals, deep Transformers simply do not train, because the gradient fades to nothing before it reaches the early layers.
Stacking the block three times below returns the same shape it took in, which is the whole point: identical blocks can be stacked, and a real GPT simply stacks more of them.
def transformer_block(x, params):
p = params
# sub-layer 1: attention, with a residual connection
a = multi_head_attention(layer_norm(x, p['g1'], p['b1']), p['Wq'], p['Wk'], p['Wv'], p['Wo'], n_head)
x = x + a
# sub-layer 2: feed-forward, with a residual connection
f = feed_forward(layer_norm(x, p['g2'], p['b2']), p['W1'], p['ff_b1'], p['W2'], p['ff_b2'])
x = x + f
return x
params = dict(
g1=np.ones(C), b1=np.zeros(C), g2=np.ones(C), b2=np.zeros(C),
Wq=Wq, Wk=Wk, Wv=Wv, Wo=Wo,
W1=W1, ff_b1=b1, W2=W2, ff_b2=b2,
)
# token embeddings + positional encoding go in; same shape comes out
h = x + positional_encoding(T, C)
for _ in range(3): # stack 3 blocks
h = transformer_block(h, params)
print("after 3 blocks:", h.shape, " (shape preserved -> blocks are stackable)")
print("finite:", np.isfinite(h).all())Line by line: what each line does
def transformer_block(x, params):: wire the two sub-layers together;paramsis a dictionary holding every weight, sop['Wq']fetches the query projection, etc.a = multi_head_attention(layer_norm(x, ...), ...): read inside-out: normalize, then attend.x = x + a: add the attention result back onto x. Thatx +is the residual: keep the original, add an edit on top.f = feed_forward(layer_norm(x, ...), ...)thenx = x + f: the same normalize → think → add-back pattern for the second sub-layer.params = dict(g1=np.ones(C), b1=np.zeros(C), ...): bundle all the pieces: LayerNorm dials (set to "no change"), the four attention grids, the two feed-forward grids and biases.h = x + positional_encoding(T, C): stamp positions onto the tokens before the first block.for _ in range(3): h = transformer_block(h, params): stack three blocks. Same shape out as in, every time -- which is exactly what lets you stack them.np.isfinite(h).all(): a health check: no infinities or NaNs anywhere after 3 blocks.
Residuals: keep x, add a small correction
Both block lines read x = x + something. That "+ x" is the residual connection. Measure the size (root-mean-square) of each piece and you can see what it buys: each sub-layer adds only a small correction on top of x, and the block's output stays on the same scale as its input.
# Measure the size of each quantity a block produces (rms = typical magnitude).
def rms(t): return float(np.sqrt((t ** 2).mean()))
x0 = x + positional_encoding(T, C) # block input
a = multi_head_attention(layer_norm(x0, params["g1"], params["b1"]), Wq, Wk, Wv, Wo, n_head)
x1 = x0 + a # first residual add
f = feed_forward(layer_norm(x1, params["g2"], params["b2"]), W1, b1, W2, b2)
x2 = x1 + f # second residual add
print(f"block input x : rms = {rms(x0):.3f}")
print(f"attention correction a : rms = {rms(a):.3f} <- added: x = x + a")
print(f"feed-forward corr. f : rms = {rms(f):.3f} <- added: x = x + f")
print(f"block output x+a+f : rms = {rms(x2):.3f} (same scale as the input)")
print("\nThe corrections are small on purpose: small initial weights make each block")
print("start as ALMOST the identity (output ~ input). The residual is what preserves")
print("the scale, so you can stack many blocks and still train them.")Line by line: what each line does
rms(t): the root-mean-square, a single number for "how big are these values typically".x0: the block's input (token embeddings plus positions), rms about 1.2.aandf: the attention and feed-forward outputs -- both much smaller thanx0, sox = x + aandx = x + fnudge the signal rather than overwrite it.- The output stays on the same scale as the input, which is exactly why cell above could stack three blocks and stay finite.
- Small corrections at the start (from small weights) mean the block begins near the identity and training gently grows it -- the residual is what makes that possible.
The wall we have just hit, on purpose
We have built the entire forward pass of a Transformer by hand. Notice what we did not do: the backward pass. Differentiating multi-head attention, LayerNorm, the residuals, and the feed-forward stack by hand, the way we did for the bigram in notebook 03, would be hundreds of careful and error-prone lines, and a single wrong sign would break the whole thing silently.
This is exactly the moment autograd (notebook 04) earns its place. In Part 2 we rebuild all of this in PyTorch, where every operation we just wrote already knows how to compute its own gradient. The concepts are the same, the code is a fraction of the length, and it actually trains.
Next: notebook 07, PyTorch and autograd.
06_transformer_block.ipynb