One prompt through every layer of a four-dimensional model, ending in a prediction, a loss and one step of learning.
The notebook carries the source, the install step and every run below, so nothing has to be on your machine. Open it, then choose Runtime › Run all.
Colab opens it read only. Click Copy to Drive to keep
your changes. You can also download
tiny_transformer.ipynb, or run it
locally: python3 tiny_transformer.py.
python3 tiny_transformer.py
All sixteen steps in order, from the token ids to the parameter count.
python3 tiny_transformer.py --attention
Both heads: projections, scaled scores, the causal mask and the softmax.
python3 tiny_transformer.py --loss
Cross-entropy on the true next token, and the gradient that is simply p minus the target.
Download tiny_transformer.py
· served verbatim at https://nlp.jcrlabz.com/code/tiny_transformer.py
"""A whole transformer, small enough to check by hand.
One prompt through every layer of a decoder-only model with d_model = 4, two
heads and one block. Every number the chapter prints comes from here.
The weights are chosen, not trained. That is the point: each one is set so you
can see what it does, and the chapter says what a trained model would put
there instead. The last section takes one real gradient step, so the arithmetic
of learning is here too, not only the arithmetic of prediction.
python3 tiny_transformer.py every stage, in order
python3 tiny_transformer.py --embed tokens, embeddings, positions
python3 tiny_transformer.py --attention Q, K, V, scores, mask, softmax
python3 tiny_transformer.py --ffn the feed-forward half of the block
python3 tiny_transformer.py --logits unembedding and the prediction
python3 tiny_transformer.py --loss cross-entropy and one SGD step
python3 tiny_transformer.py --params the parameter count, and scaling
"""
import argparse
import math
import numpy as np
np.set_printoptions(precision=4, suppress=True)
# ---------------------------------------------------------------- the model
VOCAB = ["a", "chuck", "how", "much", "wood", "woodchuck", "would", "."]
V = len(VOCAB)
D = 4 # d_model
H = 2 # heads
DK = D // H # 2 per head
DFF = 8 # feed-forward width
EPS = 1e-5
PROMPT = ["a", "woodchuck", "would", "chuck"]
TARGET = "wood"
# The four dimensions are given meanings so the arithmetic can be read:
# 0 thing 1 action 2 quantifier 3 wood
EMB = np.array([
[0.0, 0.0, 1.0, 0.0], # a
[0.0, 1.0, 0.0, 0.5], # chuck
[0.0, 0.0, 0.5, 0.0], # how
[0.0, 0.0, 0.5, 0.0], # much
[1.0, 0.0, 0.0, 1.0], # wood
[1.0, 0.0, 0.0, 0.5], # woodchuck
[0.0, 1.0, 0.0, 0.0], # would
[0.0, 0.0, 0.0, 0.0], # .
])
# Head 1 asks "who is the thing I act on": actions query, things answer.
WQ1 = np.array([[0.0, 0.0], [1.0, 0.0], [0.0, 0.0], [0.0, 1.0]])
WK1 = np.array([[1.0, 0.0], [0.0, 0.0], [0.0, 0.0], [0.0, 1.0]])
WV1 = np.array([[1.0, 0.0], [0.0, 0.0], [0.0, 0.0], [0.0, 1.0]])
# Head 2 asks "what else here is about wood".
WQ2 = np.array([[0.0, 0.0], [0.0, 0.0], [0.0, 1.0], [1.0, 0.0]])
WK2 = np.array([[0.0, 0.0], [0.0, 0.0], [0.0, 1.0], [1.0, 0.0]])
WV2 = np.array([[0.0, 1.0], [0.0, 0.0], [0.0, 0.0], [1.0, 0.0]])
# Concatenated heads, projected back to d_model.
WO = np.array([
[1.0, 0.0, 0.0, 0.5],
[0.0, 0.0, 0.0, 1.0],
[0.0, 0.0, 0.0, 0.5],
[0.0, 0.0, 0.0, 1.0],
])
W1 = np.array([
[1.0, 0.0, 0.0, 0.0, 1.0, 0.0, 0.0, 0.0],
[0.0, 1.0, 0.0, 0.0, 0.0, 1.0, 0.0, 0.0],
[0.0, 0.0, 1.0, 0.0, 0.0, 0.0, 1.0, 0.0],
[0.0, 0.0, 0.0, 1.0, 0.0, 0.0, 0.0, 1.0],
])
B1 = np.array([0.0, 0.0, 0.0, 0.0, -0.5, -0.5, -0.5, -0.5])
W2 = np.array([
[0.5, 0.0, 0.0, 0.0],
[0.0, 0.5, 0.0, 0.0],
[0.0, 0.0, 0.5, 0.0],
[0.0, 0.0, 0.0, 0.5],
[0.5, 0.0, 0.0, 0.5],
[0.0, 0.5, 0.0, 0.0],
[0.0, 0.0, 0.5, 0.0],
[0.0, 0.0, 0.0, 1.0],
])
B2 = np.zeros(D)
# ---------------------------------------------------------------- utilities
def table(rows, head, title, note=None):
print(f"\n{title}\n")
widths = [max(len(str(h)), *(len(str(r[i])) for r in rows)) for i, h in enumerate(head)]
line = " " + " ".join(str(h).rjust(w) for h, w in zip(head, widths))
print(line)
print(" " + " ".join("-" * w for w in widths))
for r in rows:
print(" " + " ".join(str(c).rjust(w) for c, w in zip(r, widths)))
if note:
print(f"\n {note}")
def fmt(x, n=4):
return f"{x:.{n}f}"
def vec(v, n=4):
return [fmt(x, n) for x in v]
def softmax(x):
e = np.exp(x - np.max(x))
return e / e.sum()
def layer_norm(x):
"""One vector at a time, gain 1 and bias 0, so the equation is visible."""
mu = x.mean()
var = x.var()
return (x - mu) / math.sqrt(var + EPS)
def positional(pos, d=D):
"""The sinusoids of Vaswani et al., computed rather than looked up."""
out = np.zeros(d)
for i in range(d // 2):
w = 1.0 / (10000 ** (2 * i / d))
out[2 * i] = math.sin(pos * w)
out[2 * i + 1] = math.cos(pos * w)
return out
# ---------------------------------------------------------------- the passes
def stage_embed(show=True):
ids = [VOCAB.index(w) for w in PROMPT]
tok = EMB[ids]
pos = np.array([positional(p) for p in range(len(PROMPT))])
x = tok + pos
if show:
table([[i, w, VOCAB.index(w)] for i, w in enumerate(PROMPT)],
["position", "token", "id"],
"STEP 1: the prompt becomes integers",
"The tokenizer is the one from the subword chapter. Here every word is already in the vocabulary.")
table([[PROMPT[i], *vec(tok[i], 1)] for i in range(len(PROMPT))],
["token", "thing", "action", "quantifier", "wood"],
"STEP 2: each id indexes one row of E",
f"E is {V} by {D}. Nothing is computed here: the row IS the lookup.")
table([[p, *vec(pos[p])] for p in range(len(PROMPT))],
["position", "d0 sin", "d1 cos", "d2 sin", "d3 cos"],
"STEP 3: positions, as sinusoids",
"Attention sees a set, not a sequence. Without this the model cannot tell "
"'a woodchuck' from 'woodchuck a'.")
table([[PROMPT[i], *vec(x[i])] for i in range(len(PROMPT))],
["token", "x0", "x1", "x2", "x3"],
"STEP 4: the block's input is the sum",
"Added, not concatenated. The model has to share the four dimensions "
"between meaning and position.")
return x
def stage_attention(x, show=True):
xn = np.array([layer_norm(r) for r in x])
if show:
table([[PROMPT[i], *vec(xn[i])] for i in range(len(PROMPT))],
["token", "n0", "n1", "n2", "n3"],
"STEP 5: layer norm, before the sublayer",
"Each row is centred and scaled on its own: mean 0, variance 1. "
"Pre-norm is what keeps a deep stack trainable.")
heads, weights = [], []
for h, (wq, wk, wv) in enumerate([(WQ1, WK1, WV1), (WQ2, WK2, WV2)], start=1):
Q, K, Vm = xn @ wq, xn @ wk, xn @ wv
scores = Q @ K.T / math.sqrt(DK)
mask = np.triu(np.ones_like(scores), k=1).astype(bool)
masked = np.where(mask, -np.inf, scores)
A = np.array([softmax(r) for r in masked])
heads.append(A @ Vm)
weights.append(A)
if show:
table([[PROMPT[i], *vec(Q[i]), *vec(K[i]), *vec(Vm[i])] for i in range(len(PROMPT))],
["token", "q0", "q1", "k0", "k1", "v0", "v1"],
f"STEP 6.{h}: head {h} projects each row three ways",
"Q, K and V are the same input seen through three different matrices.")
table([[PROMPT[i]] + [fmt(scores[i][j]) for j in range(len(PROMPT))]
for i in range(len(PROMPT))],
["query \\ key", *PROMPT],
f"STEP 7.{h}: scores, already divided by sqrt(d_k) = {fmt(math.sqrt(DK), 4)}",
"Without the square root these grow with d_k and the softmax saturates.")
table([[PROMPT[i]] + [("-" if mask[i][j] else fmt(A[i][j]))
for j in range(len(PROMPT))]
for i in range(len(PROMPT))],
["query \\ key", *PROMPT],
f"STEP 8.{h}: causal mask, then softmax along each row",
"A dash is a position the token is not allowed to see. Each row sums to 1.")
concat = np.hstack(heads)
attn = concat @ WO
out = x + attn
if show:
table([[PROMPT[i], *vec(concat[i]), *vec(attn[i])] for i in range(len(PROMPT))],
["token", "h1_0", "h1_1", "h2_0", "h2_1", "o0", "o1", "o2", "o3"],
"STEP 9: heads concatenated, then projected by W_O",
"Two heads of width 2 make one vector of width 4 again.")
table([[PROMPT[i], *vec(x[i]), *vec(attn[i]), *vec(out[i])] for i in range(len(PROMPT))],
["token", "in0", "in1", "in2", "in3", "a0", "a1", "a2", "a3",
"r0", "r1", "r2", "r3"],
"STEP 10: the residual add",
"The sublayer proposes a change. The residual keeps the original and adds it.")
return out, weights
def stage_ffn(x, show=True):
xn = np.array([layer_norm(r) for r in x])
hid = np.maximum(0.0, xn @ W1 + B1)
ff = hid @ W2 + B2
out = x + ff
if show:
table([[PROMPT[i], *vec(hid[i], 3)] for i in range(len(PROMPT))],
["token"] + [f"u{j}" for j in range(DFF)],
"STEP 11: the feed-forward layer widens, then rectifies",
f"{D} in, {DFF} hidden, {D} out. Every position goes through the same "
"weights, independently: no mixing happens here.")
table([[PROMPT[i], *vec(ff[i]), *vec(out[i])] for i in range(len(PROMPT))],
["token", "f0", "f1", "f2", "f3", "y0", "y1", "y2", "y3"],
"STEP 12: the second residual, and the block is done",
"Attention moves information between positions. The feed-forward layer "
"thinks about each position on its own.")
return out
def stage_logits(y, show=True):
last = layer_norm(y[-1])
logits = last @ EMB.T # weights tied to the embedding table
probs = softmax(logits)
order = np.argsort(-probs)
if show:
table([[fmt(v) for v in last]],
["z0", "z1", "z2", "z3"],
"STEP 13: the final norm, on the last position only",
"Only the last row can predict the next token. The others were needed "
"to build it.")
table([[VOCAB[i], fmt(logits[i]), fmt(probs[i]), f"{probs[i]*100:.1f}%"]
for i in order],
["token", "logit", "probability", ""],
"STEP 14: logits by tied unembedding, then softmax",
"The unembedding is E transposed: the same table, read the other way. "
"One matrix, two jobs, and V times d_model parameters saved.")
return logits, probs
def stage_loss(logits, probs, show=True):
t = VOCAB.index(TARGET)
loss = -math.log(probs[t])
grad = probs.copy()
grad[t] -= 1.0 # d loss / d logits, the whole derivative
lr = 0.5
new_logits = logits - lr * grad
new_probs = softmax(new_logits)
new_loss = -math.log(new_probs[t])
if show:
table([[VOCAB[i], fmt(probs[i]), "1" if i == t else "0", fmt(grad[i]),
fmt(new_probs[i])] for i in range(V)],
["token", "p", "target", "dL/dlogit", "p after"],
"STEP 15: cross-entropy, and one step downhill",
f"Loss is -log p({TARGET}) = {fmt(loss)} nats. The gradient of the loss "
"with respect to the logits is p minus the one-hot target: no chain rule needed.")
table([[fmt(loss), fmt(new_loss), fmt(loss - new_loss)]],
["loss before", "loss after", "improvement"],
f"One step of size {lr} on the logits",
"Every weight in the model is reached by pushing this same gradient "
"backwards through the steps above.")
return loss
def stage_params(show=True):
per_head = 3 * D * DK
attn = H * per_head + D * D
ffn = D * DFF + DFF + DFF * D + D
emb = V * D
total = emb + attn + ffn
if show:
table([["embedding E (tied)", f"{V} x {D}", emb],
["attention, per head", f"3 x {D} x {DK}", per_head],
[f"attention, {H} heads", "", H * per_head],
["output projection W_O", f"{D} x {D}", D * D],
["feed-forward W1, b1", f"{D} x {DFF} + {DFF}", D * DFF + DFF],
["feed-forward W2, b2", f"{DFF} x {D} + {D}", DFF * D + D],
["TOTAL", "", total]],
["part", "shape", "parameters"],
"STEP 16: the whole model, counted",
"Tying the unembedding to E saves another " f"{V * D} parameters.")
rows = []
for name, d, layers, heads, vocab, dff in [
("this chapter", D, 1, H, V, DFF),
("GPT-2 small", 768, 12, 12, 50257, 4 * 768),
("GPT-2 medium", 1024, 24, 16, 50257, 4 * 1024),
("GPT-3", 12288, 96, 96, 50257, 4 * 12288),
]:
block = 4 * d * d + 2 * d * dff + dff + d
n = vocab * d + layers * block
rows.append([name, d, layers, heads, f"{n:,}"])
table(rows, ["model", "d_model", "layers", "heads", "parameters"],
"The same equations, at four sizes",
"Nothing in the arithmetic changed. Only d_model, the number of blocks "
"and the vocabulary did.")
def main():
ap = argparse.ArgumentParser(description=__doc__)
ap.add_argument("--embed", action="store_true")
ap.add_argument("--attention", action="store_true")
ap.add_argument("--ffn", action="store_true")
ap.add_argument("--logits", action="store_true")
ap.add_argument("--loss", action="store_true")
ap.add_argument("--params", action="store_true")
a = ap.parse_args()
picked = any([a.embed, a.attention, a.ffn, a.logits, a.loss, a.params])
every = not picked
print(f"\nPROMPT {' '.join(PROMPT)}")
print(f"TARGET {TARGET}")
print(f"MODEL d_model {D}, {H} heads of {DK}, feed-forward {DFF}, "
f"1 block, vocabulary {V}")
x = stage_embed(show=every or a.embed)
y, _ = stage_attention(x, show=every or a.attention)
z = stage_ffn(y, show=every or a.ffn)
logits, probs = stage_logits(z, show=every or a.logits or a.loss)
stage_loss(logits, probs, show=every or a.loss)
if every or a.params:
stage_params()
print()
if __name__ == "__main__":
main()