{ "cells": [ { "cell_type": "markdown", "metadata": {}, "source": [ "# Self-attention, one query at a time\n", "\n", "The lecture's three-key example with every number, the scaling, and the causal mask.\n", "\n", "From chapter 16, [Self-Attention and the Transformer](https://nlp.jcrlabz.com/book/self-attention/), of the course notes.\n", "\n", "Source: `book/code/worked_examples/attention.py`" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## What you need\n", "\n", "`pip install numpy`. The next cell does it. Colab usually has numpy already, but nothing here assumes that." ] }, { "cell_type": "code", "metadata": {}, "execution_count": null, "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 live in the notebook." ] }, { "cell_type": "code", "metadata": {}, "execution_count": null, "outputs": [], "source": [ "\"\"\"Self-attention, worked as Chapter 16 works it.\n", "\n", "Five things, all reproducing the chapter's tables:\n", "\n", " one query the lecture's three-key example, every number\n", " when it selects a query that actually picks a key, and one that cannot\n", " why root d_k what the scaling stops the softmax from doing\n", " the causal mask mask then softmax, and why never the other way round\n", " the parameter count where a transformer layer's numbers sit\n", "\n", " python3 attention.py # all five\n", " python3 attention.py --one # the lecture example\n", " python3 attention.py --select # attention that discriminates\n", " python3 attention.py --scale # the square root, tested\n", " python3 attention.py --mask # causal masking\n", " python3 attention.py --params # one layer, counted\n", "\n", "Install: pip install numpy\n", "\"\"\"\n", "\n", "import argparse\n", "import math\n", "\n", "try:\n", " import numpy as np\n", "except ImportError:\n", " raise SystemExit(\"this one needs numpy: pip install numpy\")\n", "\n", "\n", "def softmax(z):\n", " e = np.exp(z - np.max(z, axis=-1, keepdims=True))\n", " return e / e.sum(axis=-1, keepdims=True)\n", "\n", "\n", "def entropy(p):\n", " \"\"\"In bits. Maximum is log2(n) for a uniform distribution.\"\"\"\n", " p = p[p > 0]\n", " return float(-(p * np.log2(p)).sum())\n", "\n", "\n", "# ---------------------------------------------------------------- one query\n", "\n", "def show_one():\n", " q = np.array([0.1, 0.2, 0.3])\n", " keys = {\"is\": [0.4, 0.5, 0.6], \"the\": [0.2, 0.1, 0.3],\n", " \"best\": [0.6, 0.7, 0.8]}\n", " K = np.array(list(keys.values()))\n", " V = K.copy() # values equal keys, as the lecture has it\n", " d = q.shape[0]\n", "\n", " print(\"\\nSTEP 1 one query against three keys\\n\")\n", " print(f\" query 'who' q = {q}\")\n", " for w, k in keys.items():\n", " label = f\"key '{w}'\"\n", " print(f\" {label:<13} k = {np.array(k)}\")\n", " print(f\"\\n Values equal keys here, to keep the arithmetic short.\")\n", " print(f\" d_k = {d}, so the scale is sqrt({d}) = {math.sqrt(d):.4f}\\n\")\n", "\n", " raw = K @ q\n", " scaled = raw / math.sqrt(d)\n", " ex = np.exp(scaled)\n", " alpha = ex / ex.sum()\n", "\n", " print(f\" {'key':<7}{'q . k':>9}{'/ sqrt(d)':>12}{'exp':>9}{'alpha':>9}\")\n", " print(\" \" + \"-\" * 48)\n", " for w, r, s, e, a in zip(keys, raw, scaled, ex, alpha):\n", " print(f\" {w:<7}{r:>9.4f}{s:>12.4f}{e:>9.4f}{a:>9.4f}\")\n", " print(\" \" + \"-\" * 48)\n", " print(f\" {'sum':<7}{'':>9}{'':>12}{ex.sum():>9.4f}{alpha.sum():>9.4f}\")\n", "\n", " z = alpha @ V\n", " print(f\"\\n output zeta = sum of alpha_i v_i\")\n", " for w, a, v in zip(keys, alpha, V):\n", " print(f\" {a:.4f} x {v}\")\n", " print(f\" = ({z[0]:.4f}, {z[1]:.4f}, {z[2]:.4f})\")\n", "\n", " print(f\"\\n Now look at the alpha column: {alpha[0]:.2f},\"\n", " f\" {alpha[1]:.2f}, {alpha[2]:.2f}.\")\n", " print(f\" That is almost uniform. Its entropy is {entropy(alpha):.4f} bits\")\n", " print(f\" against a maximum of {math.log2(3):.4f}.\")\n", " print(f\"\\n So this query attends to everything about equally, and the\")\n", " print(f\" output is close to the plain average of the values. The\")\n", " print(f\" mechanism ran correctly and selected nothing, because these\")\n", " print(f\" three keys all point the same way. Step 2 fixes that.\")\n", "\n", "\n", "# ------------------------------------------------------------ when it selects\n", "\n", "def show_select():\n", " print(\"\\n\\nSTEP 2 attention that actually discriminates\\n\")\n", " print(\" Same machinery, keys that disagree. Three dimensions standing\")\n", " print(\" for (animal, vehicle, verb).\\n\")\n", "\n", " keys = {\"dog\": [0.9, 0.0, 0.1], \"truck\": [0.0, 0.9, 0.1],\n", " \"barked\": [0.1, 0.0, 0.9]}\n", " K = np.array(list(keys.values()))\n", " V = np.array([[1, 0, 0], [0, 1, 0], [0, 0, 1]], dtype=float)\n", " d = 3\n", "\n", " for name, q in [(\"an animal query\", np.array([5.0, 0.0, 0.0])),\n", " (\"a verb query\", np.array([0.0, 0.0, 5.0])),\n", " (\"an undecided query\", np.array([2.0, 2.0, 2.0]))]:\n", " alpha = softmax(K @ q / math.sqrt(d))\n", " z = alpha @ V\n", " cells = \" \".join(f\"{w} {a:.4f}\" for w, a in zip(keys, alpha))\n", " print(f\" {name:<20}q = {q}\")\n", " print(f\" alpha {cells}\")\n", " print(f\" output ({z[0]:.4f}, {z[1]:.4f}, {z[2]:.4f})\"\n", " f\" entropy {entropy(alpha):.4f} bits\\n\")\n", "\n", " print(f\" The first two queries commit. The third cannot, and its\")\n", " print(f\" entropy sits at the {math.log2(3):.4f} bit maximum for three keys.\")\n", " print(\"\\n Low entropy means the head has made a choice. High entropy\")\n", " print(\" means it is averaging. Both are legitimate, and reading which\")\n", " print(\" one a head is doing is how attention maps get interpreted.\")\n", "\n", "\n", "# ------------------------------------------------------------- the scaling\n", "\n", "def show_scale():\n", " print(\"\\n\\nSTEP 3 why divide by the square root of d_k\\n\")\n", " print(\" Take random query and key vectors with unit variance entries.\")\n", " print(\" Their dot product is a sum of d terms, so its standard\")\n", " print(\" deviation grows like sqrt(d). Left alone, the scores get\")\n", " print(\" larger and larger as the model gets wider.\\n\")\n", "\n", " rng = np.random.default_rng(0)\n", " n = 8\n", " print(f\" {'d_k':>6}{'sd of q.k':>12}{'entropy raw':>14}\"\n", " f\"{'entropy scaled':>16}{'max alpha raw':>15}\")\n", " print(\" \" + \"-\" * 66)\n", " for d in (4, 16, 64, 256, 1024):\n", " sd, e_raw, e_scaled, top = [], [], [], []\n", " for _ in range(200):\n", " q = rng.normal(0, 1, d)\n", " K = rng.normal(0, 1, (n, d))\n", " s = K @ q\n", " sd.append(s.std())\n", " a_raw = softmax(s)\n", " a_sc = softmax(s / math.sqrt(d))\n", " e_raw.append(entropy(a_raw))\n", " e_scaled.append(entropy(a_sc))\n", " top.append(a_raw.max())\n", " print(f\" {d:>6}{np.mean(sd):>12.3f}{np.mean(e_raw):>14.4f}\"\n", " f\"{np.mean(e_scaled):>16.4f}{np.mean(top):>15.4f}\")\n", "\n", " print(f\"\\n Maximum possible entropy over {n} keys is\"\n", " f\" {math.log2(n):.4f} bits.\")\n", " print(\"\\n Read the raw column. By d = 1024 the distribution has\")\n", " print(\" collapsed onto a single key, before any training has happened.\")\n", " print(\" A softmax that saturated is a softmax with no gradient, so the\")\n", " print(\" head cannot learn its way out.\")\n", " print(\"\\n The scaled column barely moves. Dividing by sqrt(d_k) cancels\")\n", " print(\" exactly the growth the second column shows, which is why the\")\n", " print(\" scale is a square root and not something tuned.\")\n", "\n", "\n", "# ------------------------------------------------------------- causal mask\n", "\n", "def show_mask():\n", " print(\"\\n\\nSTEP 4 the causal mask, and the order of operations\\n\")\n", " print(\" A language model must not read the future. When scoring\")\n", " print(\" position i, only positions j <= i may contribute.\\n\")\n", " print(\" Set the blocked scores to minus infinity BEFORE the softmax:\\n\")\n", " print(\" S_ij = (Q K^T)_ij / sqrt(d_k) if j <= i, else -inf\\n\")\n", "\n", " S = np.array([[2.0, 1.0, 0.5, 1.5],\n", " [0.5, 2.0, 1.0, 0.5],\n", " [1.0, 0.5, 2.0, 1.0],\n", " [1.5, 1.0, 0.5, 2.0]])\n", " n = len(S)\n", " tokens = [\"the\", \"dog\", \"barked\", \"loudly\"]\n", "\n", " print(f\" raw scores S (rows are queries, columns are keys):\\n\")\n", " print(\" \" + \" \" * 10 + \"\".join(f\"{t:>10}\" for t in tokens))\n", " for i, t in enumerate(tokens):\n", " print(f\" {t:<10}\" + \"\".join(f\"{v:>10.2f}\" for v in S[i]))\n", "\n", " mask = np.tril(np.ones((n, n), dtype=bool))\n", " Sm = np.where(mask, S, -np.inf)\n", " A = softmax(Sm)\n", "\n", " print(f\"\\n after masking then softmax:\\n\")\n", " print(\" \" + \" \" * 10 + \"\".join(f\"{t:>10}\" for t in tokens) + f\"{'sum':>8}\")\n", " for i, t in enumerate(tokens):\n", " cells = \"\".join(\" .\" if not mask[i][j] else f\"{A[i][j]:>10.4f}\"\n", " for j in range(n))\n", " print(f\" {t:<10}{cells}{A[i].sum():>8.4f}\")\n", "\n", " print(\"\\n Every row still sums to 1. That is the whole reason the mask\")\n", " print(\" goes before the softmax rather than after.\\n\")\n", "\n", " wrong = softmax(S) * mask\n", " print(\" Compare softmax first, then zero the blocked cells:\\n\")\n", " print(\" \" + \" \" * 10 + \"\".join(f\"{t:>10}\" for t in tokens) + f\"{'sum':>8}\")\n", " for i, t in enumerate(tokens):\n", " cells = \"\".join(\" .\" if not mask[i][j] else f\"{wrong[i][j]:>10.4f}\"\n", " for j in range(n))\n", " print(f\" {t:<10}{cells}{wrong[i].sum():>8.4f}\")\n", "\n", " print(f\"\\n Row 1 now sums to {wrong[0].sum():.4f} instead of 1. The mass\")\n", " print(\" that belonged to the blocked positions has simply been thrown\")\n", " print(\" away rather than redistributed.\")\n", " print(\"\\n Masking first lets the surviving positions inherit that mass.\")\n", " print(\" Masking second silently scales the whole row down, and the\")\n", " print(\" earlier the token, the worse it gets.\")\n", "\n", "\n", "# --------------------------------------------------------- the parameter count\n", "\n", "def show_params(V=50000, D=512, L=6, h=8, H=2048):\n", " d_k = d_v = D // h\n", " print(f\"\\n\\nSTEP 5 a transformer layer, counted\\n\")\n", " print(f\" |V| = {V}, D = {D}, layers L = {L}, heads h = {h},\")\n", " print(f\" d_k = d_v = D/h = {d_k}, feed-forward H = {H}\\n\")\n", "\n", " attn = 4 * D * D # W^Q, W^K, W^V, W^O\n", " ffn = 2 * D * H + H + D\n", " ln = 4 * D # two layer norms, scale and shift\n", " layer = attn + ffn + ln\n", " emb = V * D\n", "\n", " print(f\" {'component':<26}{'formula':<22}{'parameters':>14}\")\n", " print(\" \" + \"-\" * 64)\n", " print(f\" {'attention Q,K,V,O':<26}{'4 D^2':<22}{attn:>14,}\")\n", " print(f\" {'feed-forward':<26}{'2 D H + H + D':<22}{ffn:>14,}\")\n", " print(f\" {'layer norms':<26}{'4 D':<22}{ln:>14,}\")\n", " print(\" \" + \"-\" * 64)\n", " print(f\" {'one layer':<26}{'':<22}{layer:>14,}\")\n", " print(f\" {'all ' + str(L) + ' layers':<26}{'L x layer':<22}\"\n", " f\"{L * layer:>14,}\")\n", " print(f\" {'embeddings':<26}{'V D':<22}{emb:>14,}\")\n", " print(\" \" + \"-\" * 64)\n", " print(f\" {'total':<48}{L * layer + emb:>14,}\")\n", "\n", " print(f\"\\n Two things to notice.\")\n", " print(f\"\\n The heads are free. Splitting D into {h} heads of {d_k} costs\")\n", " print(f\" nothing, because h x d_k = D. Multi-head attention is a\")\n", " print(f\" reshape, not an extra parameter budget.\")\n", " print(f\"\\n The feed-forward block is larger than the attention block,\")\n", " print(f\" {ffn:,} against {attn:,}. Attention gets the attention, and\")\n", " print(f\" most of the parameters sit next door.\")\n", "\n", " print(f\"\\n\\n WHAT THE SEQUENCE LENGTH COSTS\\n\")\n", " print(f\" No parameter count above mentions sequence length. The\")\n", " print(f\" computation does: every position attends to every position.\\n\")\n", " print(f\" {'tokens n':>10}{'n^2 scores per head':>22}{'vs n = 512':>13}\")\n", " print(\" \" + \"-\" * 46)\n", " for nt in (128, 512, 2048, 8192, 32768):\n", " print(f\" {nt:>10}{nt * nt:>22,}{nt * nt / 512 ** 2:>12.2f}x\")\n", " print(f\"\\n Quadruple the context and the attention cost goes up\")\n", " print(f\" sixteen-fold. That is the wall the efficiency chapter climbs.\")\n", "\n", "\n", "def main():\n", " ap = argparse.ArgumentParser(description=__doc__)\n", " for f in (\"one\", \"select\", \"scale\", \"mask\", \"params\"):\n", " ap.add_argument(f\"--{f}\", action=\"store_true\")\n", " a = ap.parse_args()\n", " picked = a.one or a.select or a.scale or a.mask or a.params\n", " if a.one or not picked:\n", " show_one()\n", " if a.select or not picked:\n", " show_select()\n", " if a.scale or not picked:\n", " show_scale()\n", " if a.mask or not picked:\n", " show_mask()\n", " if a.params or not picked:\n", " show_params()\n", " print()" ] }, { "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()`. The heading shows the equivalent shell command." ] }, { "cell_type": "code", "metadata": {}, "execution_count": null, "outputs": [], "source": [ "import sys\n", "\n", "def run(*args):\n", " sys.argv = [\"attention.py\", *args]\n", " main()" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "### `python3 attention.py`\n", "\n", "All five sections." ] }, { "cell_type": "code", "metadata": {}, "execution_count": null, "outputs": [], "source": [ "run()" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "### `python3 attention.py --one`\n", "\n", "One query worked end to end, every number the chapter prints." ] }, { "cell_type": "code", "metadata": {}, "execution_count": null, "outputs": [], "source": [ "run(\"--one\")" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "### `python3 attention.py --scale`\n", "\n", "Without the root d_k, entropy over 8 keys collapses from 1.78 bits at d = 4 to 0.13 bits at d = 1024, before any training." ] }, { "cell_type": "code", "metadata": {}, "execution_count": null, "outputs": [], "source": [ "run(\"--scale\")" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "### `python3 attention.py --mask`\n", "\n", "Mask then softmax and every row sums to 1. Mask after and row one sums to 0.4551, with the blocked mass thrown away instead of redistributed." ] }, { "cell_type": "code", "metadata": {}, "execution_count": null, "outputs": [], "source": [ "run(\"--mask\")" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "---\n", "\n", "Read the chapter this comes from: [Self-Attention and the Transformer](https://nlp.jcrlabz.com/book/self-attention/)." ] } ], "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 }