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