{
 "nbformat": 4,
 "nbformat_minor": 5,
 "metadata": {
  "kernelspec": {
   "name": "python3",
   "display_name": "Python 3",
   "language": "python"
  },
  "language_info": {
   "name": "python"
  },
  "colab": {
   "name": "attention.ipynb"
  }
 },
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": "# attention \u2014 Python demo\n\nNumerical companion to the entry [attention](https://dictionaryofml.org/terms/attention.html) of the [Dictionary of Applied Machine Learning](https://dictionaryofml.org/): it recomputes what the entry states and prints one line per check.\n\nIllustrate attention as a learned *associative memory*. Each token turns its own embedding into a query and searches, via inner products against the keys of the other tokens, for the most relevant tokens; the softmax-weighted values are then read out. Training makes queries and keys align so that a token can be reconstructed from the tokens it is associated with.\n\nRequires NumPy and Matplotlib only, and uses fixed seeds, so the printed numbers reproduce exactly. Generated from [`pythondemos/attention.py`](https://dictionaryofml.org/terms/attention.py); CC BY 4.0."
  },
  {
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": "# Notebook shim: the script resolves output paths relative to __file__,\n# which a notebook kernel does not define; everything lands in the\n# working directory instead.\nimport os\n__file__ = os.path.join(os.getcwd(), \"attention.py\")\nos.makedirs(\"pythondemos\", exist_ok=True)"
  },
  {
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": "\"\"\"\nattention.py \u2014 train a single self-attention layer on sentences from the\nUniversal Declaration of Human Rights.\n\nPurpose\n-------\nIllustrate attention as a learned *associative memory*.  Each token turns its\nown embedding into a query and searches, via inner products against the keys of\nthe other tokens, for the most relevant tokens; the softmax-weighted values are\nthen read out.  Training makes queries and keys align so that a token can be\nreconstructed from the tokens it is associated with.\n\nCorpus\n------\nSentences and clauses excerpted from the Universal Declaration of Human Rights\n(UN General Assembly resolution 217 A, 1948; public domain), stored verbatim\nbelow so the demo is self-contained and reproducible (no network access, fixed\nseed).  The display sentence is the first sentence of Article 1, which is also\nthe example sentence of the scaled-dot-product figure in the `attention`\nglossary entry.\n\nModel\n-----\nOne bidirectional self-attention head with learned matrices W_Q, W_K, W_V\nand a linear read-out W_O to one score per distinct token.  Embeddings E are learned too.\nThe self-supervised objective is a cloze task, associative memory put to\nwork: for every\nposition i the diagonal of the attention matrix is masked out (a token may not\nattend to itself), and the read-out must predict the held-out token at i from a\nsoftmax-weighted combination of the *other* tokens' values.  This forces the\nattention weights to encode which tokens are mutually predictive \u2014 the essence\nof content-addressable memory.\n\nManual forward/backward (numpy only):\n    Q = X W_Q,  K = X W_K,  V = X W_V                       (per sentence)\n    S = Q K^T / sqrt(d_k),  S_ii = -inf                     (mask self)\n    A = softmax(S, axis=1)                                  (rows sum to 1)\n    Z = A V,  scores = Z W_O\n    loss = mean negative log probability assigned to token i at position i\n\nThe whole model is implemented by hand in numpy \u2014 no autograd \u2014 so every line\nof the forward pass has an explicit matching line in the backward pass below.\nThat is the point of the demo: to expose the arithmetic that a deep-learning\nframework would otherwise hide.\n\nBlocks\n------\n[B-corpus]   the excerpted sentences of the Universal Declaration of Human\n             Rights, stored verbatim so the demo needs no network\n[B-tokens]   tokenization into lowercase word tokens, and the set of\n             distinct tokens\n[B-model]    the learned parameters: embeddings E and the matrices\n             W_Q, W_K, W_V, W_O, plus the Adam state\n[B-forward]  one sentence forward and backward by hand, every backward line\n             matching a forward one\n[B-train]    4000 full-corpus updates; the loss falls well below its\n             starting value\n[B-display]  the attention matrix of the display sentence, the same sentence\n             the entry's scaled-dot-product figure uses\n[B-csv]      the two CSVs the entry's pgfplots figures read\n[B-preview]  the matplotlib preview: training curve and attention heat map\n\nOutputs\n-------\n  pythondemos/attention_loss.csv    \u2014 columns: iter, loss (training curve)\n  pythondemos/attention_weights.csv \u2014 columns: q, k, w  (learned attention\n        matrix of one display sentence; q = query position, k = key position)\n  pythondemos/attention.png         \u2014 matplotlib preview (loss + heat map)\n\nThe display sentence and its token order are printed and asserted so the tick\nlabels hard-coded in the `attention` glossary entry stay in sync.\n\"\"\"\n\nimport re\nimport numpy as np\nimport matplotlib.pyplot as plt\nfrom pathlib import Path\n\nOUT_DIR = Path(__file__).parent\n# Fix the RNG so the learned embeddings, matrices, and hence the heat map\n# committed to the repo are reproducible bit-for-bit on every run.\nnp.random.seed(0)"
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": "**[B-corpus]** the excerpted sentences of the Universal Declaration of Human Rights, stored verbatim so the demo needs no network"
  },
  {
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": "# One sentence/clause per line (Articles 1, 3, 4, 5, 6, 7, 9, 13, 17,\n# 18, 19, 24, 26).  Short, thematically overlapping sentences share many\n# tokens (\"everyone\", \"right\", \"no one shall\", \"freedom\"), which gives the\n# single head recurring associations to latch onto during training.\nCORPUS = \"\"\"\nAll human beings are born free and equal in dignity and rights.\nThey are endowed with reason and conscience and should act towards one another in a spirit of brotherhood.\nEveryone has the right to life, liberty and security of person.\nNo one shall be held in slavery or servitude.\nSlavery and the slave trade shall be prohibited in all their forms.\nNo one shall be subjected to torture or to cruel, inhuman or degrading treatment or punishment.\nEveryone has the right to recognition everywhere as a person before the law.\nAll are equal before the law and are entitled without any discrimination to equal protection of the law.\nNo one shall be subjected to arbitrary arrest, detention or exile.\nEveryone has the right to freedom of movement and residence within the borders of each state.\nEveryone has the right to own property alone as well as in association with others.\nNo one shall be arbitrarily deprived of his property.\nEveryone has the right to freedom of thought, conscience and religion.\nEveryone has the right to freedom of opinion and expression.\nEveryone has the right to rest and leisure, including reasonable limitation of working hours and periodic holidays with pay.\nEveryone has the right to education.\n\"\"\""
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": "**[B-tokens]** tokenization into lowercase word tokens, and the set of distinct tokens"
  },
  {
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": "# Keep only runs of letters; this drops punctuation and case so that, e.g.,\n# \"rights.\" and \"rights\" map to the same token.\ndef tokenize(sentence):\n    return re.findall(r\"[a-z]+\", sentence.lower())\n\nsentences = [tokenize(s) for s in CORPUS.strip().split(\"\\n\")]\n# A self-attention head needs at least two tokens (a token attends to the\n# *others*), so discard any degenerate one-token line.\nsentences = [s for s in sentences if len(s) >= 2]\n\n# The distinct tokens are collected and sorted; stoi maps each token to\n# its integer id (its row in the embedding matrix E and its column in the\n# read-out scores).  Sorting makes the id assignment deterministic.\nvocab = sorted({tok for s in sentences for tok in s})\nstoi = {w: i for i, w in enumerate(vocab)}\nV = len(vocab)\nprint(f\"[B-tokens] {len(sentences)} sentences, {V} distinct tokens\")"
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": "**[B-model]** the learned parameters: embeddings E and the matrices W_Q, W_K, W_V, W_O, plus the Adam state"
  },
  {
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": "# The query, key, and value vectors all have the embedding width here for simplicity; in a real\n# transformer d_k and d_v are typically smaller than d (per head).\nd = 24        # embedding dimension\nd_k = 24      # query/key dimension\nd_v = 24      # value dimension\n\ndef randn(*shape):\n    # small random init, scaled by fan-in for stable gradients\n    return 0.1 * np.random.randn(*shape)\n\n# The learnable parameters.  E is the input embedding table; W_Q, W_K, W_V are\n# the matrices that turn an embedding into a query, key, and value;\n# W_O reads the attention output back out to a score for every distinct token.\nE = randn(V, d)          # token embeddings (learned)\nW_Q = randn(d, d_k)\nW_K = randn(d, d_k)\nW_V = randn(d, d_v)\nW_O = randn(d_v, V)      # read-out to one score per distinct token\nparams = {\"E\": E, \"W_Q\": W_Q, \"W_K\": W_K, \"W_V\": W_V, \"W_O\": W_O}\n\n# Adam state: a first-moment (m) and second-moment (v_) running\n# average per parameter tensor.  Adam adapts the update magnitude per\n# coordinate,\n# which makes this tiny hand-written model train in a few thousand updates.\nm = {k: np.zeros_like(v) for k, v in params.items()}\nv_ = {k: np.zeros_like(v) for k, v in params.items()}\nbeta1, beta2, eps, lr = 0.9, 0.999, 1e-8, 0.02"
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": "**[B-forward]** one sentence forward and backward by hand, every backward line matching a forward one"
  },
  {
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": "def softmax_rows(S):\n    # Numerically stable row-wise softmax: subtracting the row max before\n    # exponentiating avoids overflow and leaves the result unchanged.\n    S = S - S.max(axis=1, keepdims=True)\n    P = np.exp(S)\n    return P / P.sum(axis=1, keepdims=True)\n\ndef forward_backward(idx, grads):\n    \"\"\"One sentence: accumulate gradients into `grads`, return (loss, n_pred).\n\n    `idx` is the list of token ids for one sentence.  The forward pass computes\n    the attention output and the prediction loss; the backward pass then walks\n    the same operations in reverse, accumulating the gradient of the loss w.r.t.\n    every parameter into the shared `grads` dictionary (summed over sentences).\n    \"\"\"\n    n = len(idx)\n    # \u2500\u2500 forward \u2500\u2500\n    X = E[idx]                                   # (n, d)  embeddings of this sentence\n    Q = X @ W_Q                                  # (n, d_k) queries: what each token asks for\n    K = X @ W_K                                  # (n, d_k) keys: what each token advertises\n    Vv = X @ W_V                                 # (n, d_v) values: what each token contributes\n    S = (Q @ K.T) / np.sqrt(d_k)                 # (n, n) scaled query-key match scores\n    mask = np.eye(n, dtype=bool)                 # forbid self-attention\n    # Setting the diagonal to a large negative number makes softmax assign it\n    # ~zero weight: token i must reconstruct itself from the *other* tokens,\n    # forcing retrieval from the other tokens rather than trivial copying.\n    S = np.where(mask, -1e9, S)\n    A = softmax_rows(S)                          # (n, n), rows sum to 1: retrieval weights\n    Z = A @ Vv                                   # (n, d_v) retrieved content per query token\n    scores = Z @ W_O                             # (n, V) one score per distinct token\n    P = softmax_rows(scores)                     # predicted token distribution\n\n    # loss: negative log probability of the (held-out) token at each position\n    # targets[i] is the true token id at position i; the loss is small when\n    # probability mass sits on it.  The +1e-12 guards log(0).\n    targets = np.array(idx)\n    loss = -np.log(P[np.arange(n), targets] + 1e-12).sum()\n\n    # \u2500\u2500 backward \u2500\u2500\n    # Each block below differentiates the matching forward line, applied\n    # in reverse order (chain rule).  Shapes are annotated to make the matrix\n    # multiplications self-checking.\n    # d(loss)/d(scores) for softmax + negative log probability is simply P - onehot(target).\n    dscores = P.copy()\n    dscores[np.arange(n), targets] -= 1.0        # (n, V)\n    # scores = Z @ W_O  ->  gradients w.r.t. W_O and Z.\n    grads[\"W_O\"] += Z.T @ dscores\n    dZ = dscores @ W_O.T                          # (n, d_v)\n\n    # Z = A @ Vv  ->  split the gradient between the attention weights A and\n    # the values Vv.\n    dA = dZ @ Vv.T                                # (n, n)\n    dVv = A.T @ dZ                                # (n, d_v)\n    # softmax backward per row: Jacobian of a row-softmax applied to dA.\n    dS = A * (dA - (dA * A).sum(axis=1, keepdims=True))\n    # The masked diagonal receives no gradient (its score was a constant).\n    dS = np.where(mask, 0.0, dS)\n    # S = (Q @ K.T)/sqrt(d_k)  ->  gradients w.r.t. the queries and keys.\n    dQ = (dS @ K) / np.sqrt(d_k)\n    dK = (dS.T @ Q) / np.sqrt(d_k)\n\n    # Q = X @ W_Q, K = X @ W_K, Vv = X @ W_V  ->  gradients of W_Q, W_K, W_V.\n    grads[\"W_Q\"] += X.T @ dQ\n    grads[\"W_K\"] += X.T @ dK\n    grads[\"W_V\"] += X.T @ dVv\n    # X feeds all three matrices, so its gradient is the sum of the three\n    # paths back through W_Q, W_K, W_V.\n    dX = dQ @ W_Q.T + dK @ W_K.T + dVv @ W_V.T    # (n, d)\n    # Scatter-add into the embedding table: a token may occur several times in\n    # one sentence (e.g. \"and\", \"for\", \"the\"), so all of its positions\n    # accumulate into the single shared embedding row.  Plain E[idx] += dX would\n    # drop the duplicates; np.add.at sums them correctly.\n    np.add.at(grads[\"E\"], idx, dX)\n    return loss, n\n\n# Pre-convert every sentence to its list of token ids once.\nsent_idx = [[stoi[w] for w in s] for s in sentences]"
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": "**[B-train]** 4000 full-corpus updates; the loss falls well below its starting value"
  },
  {
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": "# The corpus is tiny, so each update sums the gradient\n# over all sentences, then take a single Adam update.\nEPOCHS = 4000\nloss_curve = []\nfor epoch in range(1, EPOCHS + 1):\n    grads = {k: np.zeros_like(v) for k, v in params.items()}\n    total_loss, total_pred = 0.0, 0\n    for idx in sent_idx:\n        l, n = forward_backward(idx, grads)\n        total_loss += l\n        total_pred += n\n    # Report the loss per predicted token so the curve is comparable across\n    # sentences of different lengths.\n    mean_loss = total_loss / total_pred\n\n    # Adam update: bias-corrected first/second moments give a per-coordinate\n    # adaptive step.  The gradient is averaged over sentences (\u00f7 len) so the\n    # update magnitude does not scale with corpus size.\n    for k in params:\n        g = grads[k] / len(sent_idx)\n        m[k] = beta1 * m[k] + (1 - beta1) * g\n        v_[k] = beta2 * v_[k] + (1 - beta2) * (g * g)\n        mhat = m[k] / (1 - beta1 ** epoch)\n        vhat = v_[k] / (1 - beta2 ** epoch)\n        params[k] -= lr * mhat / (np.sqrt(vhat) + eps)\n\n    # Subsample the curve (every 20th update) to keep the committed CSV small.\n    if epoch == 1 or epoch % 20 == 0:\n        loss_curve.append((epoch, mean_loss))\n\nprint(f\"[B-train] initial loss = {loss_curve[0][1]:.4f}, final loss = {loss_curve[-1][1]:.4f}\")\n# Guard against a silently broken gradient: a correct implementation drives the\n# loss well below its starting value.\nassert loss_curve[-1][1] < loss_curve[0][1] - 0.3, \"training did not reduce the loss\""
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": "**[B-display]** the attention matrix of the display sentence, the same sentence the entry's scaled-dot-product figure uses"
  },
  {
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": "# First sentence of Article 1 \u2014 the same sentence used in the\n# scaled-dot-product figure of the `attention` glossary entry.\nDISPLAY = [\"all\", \"human\", \"beings\", \"are\", \"born\", \"free\",\n           \"and\", \"equal\", \"in\", \"dignity\", \"and\", \"rights\"]\n# This sentence must appear verbatim (as a token list) in the corpus, so the\n# heat map shows weights the head actually trained on (not an out-of-sample one).\nassert DISPLAY in sentences, \"display sentence not found in the corpus\"\nprint(\"[B-display] display tokens:\", \" \".join(DISPLAY))\n\n# Recompute the forward attention weights for the display sentence only (no\n# gradient needed here), reusing the trained embeddings and matrices.\nd_idx = [stoi[w] for w in DISPLAY]\nX = E[d_idx]\nS = (X @ W_Q) @ (X @ W_K).T / np.sqrt(d_k)\nS = np.where(np.eye(len(d_idx), dtype=bool), -1e9, S)\nA = softmax_rows(S)"
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": "**[B-csv]** the two CSVs the entry's pgfplots figures read"
  },
  {
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": "# Training curve: consumed by the loss plot below (and available to the entry).\nnp.savetxt(OUT_DIR / \"attention_loss.csv\", np.array(loss_curve),\n           delimiter=\",\", header=\"iter,loss\", comments=\"\", fmt=\"%.6f\")\nprint(f\"[B-csv] saved {OUT_DIR / 'attention_loss.csv'}\")\n\n# Attention matrix in long/tidy form (one row per (query, key) cell) so that\n# pgfplots can read it directly with `matrix plot*` in the glossary figure.\nrows = []\nfor q in range(A.shape[0]):\n    for k in range(A.shape[1]):\n        rows.append((q, k, A[q, k]))\nnp.savetxt(OUT_DIR / \"attention_weights.csv\", np.array(rows),\n           delimiter=\",\", header=\"q,k,w\", comments=\"\", fmt=\"%.6f\")\nprint(f\"[B-csv] saved {OUT_DIR / 'attention_weights.csv'}\")"
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": "**[B-preview]** the matplotlib preview: training curve and attention heat map"
  },
  {
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": "# A local sanity-check PDF (not committed as the paper figure): the training\n# curve on the left, the learned attention heat map on the right.\nfig, (ax1, ax2) = plt.subplots(1, 2, figsize=(8.5, 3.6))\nep, ls = zip(*loss_curve)\nax1.plot(ep, ls, \"-\", color=\"tab:blue\")\nax1.set_xlabel(\"update\")\nax1.set_ylabel(\"loss per predicted token\")\nax1.set_title(\"training loss\")\n\n# Rows = query token (searching), columns = key token (searched); a dark cell\n# means the query in that row retrieves strongly from the key in that column.\nim = ax2.imshow(A, cmap=\"Blues\", vmin=0.0)\nax2.set_xticks(range(len(DISPLAY)))\nax2.set_yticks(range(len(DISPLAY)))\nax2.set_xticklabels(DISPLAY, rotation=90, fontsize=7)\nax2.set_yticklabels(DISPLAY, fontsize=7)\nax2.set_xlabel(\"key token (searched)\")\nax2.set_ylabel(\"query token (searching)\")\nax2.set_title(\"learned attention weights\")\nfig.colorbar(im, ax=ax2, fraction=0.046)\n\nfig.tight_layout()\nout = OUT_DIR / \"attention.png\"\nfig.savefig(out, bbox_inches=\"tight\", dpi=110)\nprint(f\"[B-preview] saved {out}\")\nplt.close()"
  }
 ]
}