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