{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "tf_setup",
   "metadata": {},
   "outputs": [],
   "source": [
    "from platform import python_version\n",
    "print(python_version())\n",
    "import numpy as np, matplotlib.pyplot as plt, seaborn as sns\n",
    "import torch\n",
    "sns.set_style(\"whitegrid\"); sns.set_palette(\"colorblind\"); palette = sns.color_palette()\n",
    "figsize = (15, 8); legend_fontsize = 16"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "rnn0",
   "metadata": {},
   "source": [
    "# Recurrent neural networks"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "rnn1",
   "metadata": {},
   "source": [
    "## A linear recurrent layer is a dynamical system"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "rnn2",
   "metadata": {},
   "outputs": [],
   "source": [
    "import math, torch, torch.nn as nn\n",
    "sns.set_style(\"whitegrid\"); palette = sns.color_palette()\n",
    "torch.manual_seed(0)\n",
    "# A linear RNN  h_t = W h_{t-1}  is just a dynamical system: with W = rho*Q (Q orthogonal) we get\n",
    "# ||h_t|| = rho^t, so the spectral radius rho = max|eig(W)| decides -- the signal forgets, holds, or explodes.\n",
    "N, T = 200, 60\n",
    "Q, _ = torch.linalg.qr(torch.randn(N, N)); h0 = torch.randn(N); h0 = h0 / h0.norm()\n",
    "fig, axd = plt.subplot_mosaic([[\"L\",\"a\"],[\"L\",\"b\"],[\"L\",\"c\"]], figsize=(15, 6))\n",
    "for rho in [0.85, 0.95, 1.0, 1.05, 1.15]:\n",
    "    h = h0.clone(); norms = [1.0]\n",
    "    for _ in range(T): h = (rho*Q) @ h; norms.append(h.norm().item())\n",
    "    axd[\"L\"].semilogy(norms, lw=2.3, label=r\"$\\rho=%.2f$\" % rho)\n",
    "axd[\"L\"].axhline(1, color=\"0.4\", ls=\":\", lw=1); axd[\"L\"].legend(fontsize=12, loc=\"center left\")\n",
    "axd[\"L\"].set_xlabel(\"time step $t$\"); axd[\"L\"].set_ylabel(r\"$\\|h_t\\|$  (log scale)\")\n",
    "axd[\"L\"].set_title(r\"forward signal $h_t = W h_{t-1}$: the spectral radius $\\rho$ decides\", fontsize=13)\n",
    "# a 2-D rotation makes the \"echo\" of one input pulse visible: damped / sustained / growing\n",
    "ang = 2*math.pi/11\n",
    "for key, rho, lab, ci in [(\"a\",0.9,r\"$\\rho=0.9$: forgetting\",0),(\"b\",1.0,r\"$\\rho=1.0$: stable memory\",2),(\"c\",1.1,r\"$\\rho=1.1$: exploding\",4)]:\n",
    "    R = rho*torch.tensor([[math.cos(ang),-math.sin(ang)],[math.sin(ang),math.cos(ang)]])\n",
    "    h = torch.tensor([1.,0.]); xs = [1.]\n",
    "    for _ in range(T): h = R @ h; xs.append(h[0].item())\n",
    "    axd[key].plot(xs, lw=2.2, color=palette[ci]); axd[key].axhline(0, color=\"0.6\", lw=.8)\n",
    "    axd[key].margins(x=0); axd[key].set_title(lab, fontsize=11, loc=\"left\")\n",
    "    if key != \"c\": axd[key].set_xticklabels([])\n",
    "axd[\"c\"].set_xlabel(\"time step $t$ after a single input pulse\")\n",
    "plt.tight_layout(); plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "rnn3",
   "metadata": {},
   "source": [
    "## A long-memory task: the adding problem"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "rnn4",
   "metadata": {},
   "outputs": [],
   "source": [
    "# The adding problem (Hochreiter & Schmidhuber's classic LSTM benchmark): two input channels --\n",
    "# a stream of random values in [0,1] and a marker channel with exactly two 1's; the target is the\n",
    "# sum of the two marked values. To solve it the net must carry the first marked value across the gap.\n",
    "def adding_batch(B, T, seed=None):\n",
    "    g = torch.Generator().manual_seed(seed) if seed is not None else None\n",
    "    vals = torch.rand(B, T, 1, generator=g); mark = torch.zeros(B, T, 1)\n",
    "    i = torch.randint(0, T//2, (B,), generator=g); j = torch.randint(T//2, T, (B,), generator=g)\n",
    "    mark[torch.arange(B), i, 0] = 1; mark[torch.arange(B), j, 0] = 1\n",
    "    y = (vals[torch.arange(B), i, 0] + vals[torch.arange(B), j, 0]).unsqueeze(1)\n",
    "    return torch.cat([vals, mark], -1), y\n",
    "xb, yb = adding_batch(1, 40, seed=3); v = xb[0,:,0].numpy(); m = xb[0,:,1].numpy().astype(bool)\n",
    "fig, ax = plt.subplots(figsize=(14, 4))\n",
    "ax.stem(np.arange(40), v, linefmt=\"0.7\", markerfmt=\"o\", basefmt=\" \")\n",
    "ax.plot(np.where(m)[0], v[m], \"*\", ms=22, color=palette[3], label=\"marked  (add these two)\")\n",
    "ax.set_title(\"the adding problem: output the sum of the two marked values   (target = %.2f)\" % yb.item())\n",
    "ax.set_xlabel(\"time step\"); ax.set_ylabel(\"value channel\"); ax.legend(fontsize=13); plt.tight_layout(); plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "rnn5",
   "metadata": {},
   "source": [
    "## Vanilla RNN vs. LSTM vs. GRU"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "rnn6",
   "metadata": {},
   "outputs": [],
   "source": [
    "# One linear readout on top of three different recurrent cells; train each on the adding problem (T=40).\n",
    "class Seq(nn.Module):\n",
    "    def __init__(self, kind, hidden=64):\n",
    "        super().__init__()\n",
    "        self.rnn = {\"RNN\": nn.RNN, \"LSTM\": nn.LSTM, \"GRU\": nn.GRU}[kind](2, hidden, batch_first=True)\n",
    "        self.fc = nn.Linear(hidden, 1)\n",
    "    def forward(self, x): out, _ = self.rnn(x); return self.fc(out[:, -1])\n",
    "def train(kind, T=40, iters=600, B=128, lr=5e-3, seed=0):\n",
    "    torch.manual_seed(seed); net = Seq(kind); opt = torch.optim.Adam(net.parameters(), lr); lf = nn.MSELoss()\n",
    "    xte, yte = adding_batch(512, T, seed=999); curve = []\n",
    "    for it in range(iters):\n",
    "        x, y = adding_batch(B, T); opt.zero_grad(); lf(net(x), y).backward()\n",
    "        nn.utils.clip_grad_norm_(net.parameters(), 1.0); opt.step()   # clip: vanilla RNNs explode otherwise\n",
    "        if it % 20 == 0:\n",
    "            with torch.no_grad(): curve.append((it, lf(net(xte), yte).item()))\n",
    "    return net, curve\n",
    "models = {}; fig, ax = plt.subplots(figsize=(11, 6))\n",
    "for kind in [\"RNN\", \"LSTM\", \"GRU\"]:\n",
    "    models[kind], curve = train(kind); ax.semilogy(*zip(*curve), lw=2.4, label=\"%s  (final %.3f)\" % (kind, curve[-1][1]))\n",
    "ax.axhline(1/6, color=\"0.4\", ls=\"--\", lw=1.5, label=\"predict-the-mean baseline (0.167)\")\n",
    "ax.set_xlabel(\"iteration\"); ax.set_ylabel(\"test MSE\"); ax.set_title(\"adding problem, sequence length T=40\")\n",
    "ax.legend(fontsize=12); plt.tight_layout(); plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "rnn7",
   "metadata": {},
   "source": [
    "## What the trained LSTM actually does"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "rnn8",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Open up the trained LSTM: run it on one sequence and watch the gates and the cell state.\n",
    "net = models[\"LSTM\"].eval(); H = 64\n",
    "torch.manual_seed(4); xe = torch.zeros(1, 40, 2); xe[0,:,0] = torch.rand(40); t1, t2 = 7, 27\n",
    "xe[0,t1,1] = 1; xe[0,t2,1] = 1; tgt = (xe[0,t1,0] + xe[0,t2,0]).item()\n",
    "with torch.no_grad():\n",
    "    out, _ = net.rnn(xe); guess = net.fc(out[0]).squeeze(-1).numpy()         # running prediction fc(h_t)\n",
    "Wi, Wh = net.rnn.weight_ih_l0, net.rnn.weight_hh_l0; bi, bh = net.rnn.bias_ih_l0, net.rnn.bias_hh_l0\n",
    "h = torch.zeros(H); c = torch.zeros(H); wr = []; cn = []                     # replay the recurrence to read gates (order i,f,g,o)\n",
    "with torch.no_grad():\n",
    "    for t in range(40):\n",
    "        g = Wi @ xe[0,t] + bi + Wh @ h + bh\n",
    "        i, f, cg, o = torch.sigmoid(g[:H]), torch.sigmoid(g[H:2*H]), torch.tanh(g[2*H:3*H]), torch.sigmoid(g[3*H:])\n",
    "        c = f*c + i*cg; h = o*torch.tanh(c); wr.append((i*cg).norm().item()); cn.append(c.norm().item())\n",
    "v = xe[0,:,0].numpy(); run = np.cumsum([v[t] if t in (t1, t2) else 0 for t in range(40)])\n",
    "fig, ax = plt.subplots(3, 1, figsize=(13, 8), sharex=True)\n",
    "ax[0].stem(range(40), v, linefmt=\"0.75\", markerfmt=\"o\", basefmt=\" \"); ax[0].plot([t1,t2],[v[t1],v[t2]],\"*\",ms=22,color=palette[3])\n",
    "ax[0].set_ylabel(\"input value\"); ax[0].set_title(\"inside the LSTM: markers at t=%d, %d   (target = %.2f)\" % (t1,t2,tgt))\n",
    "ax[1].plot(wr, lw=2.4, color=palette[0], label=r\"info written to cell  $\\|i_t\\odot\\tilde c_t\\|$\")\n",
    "axb = ax[1].twinx(); axb.plot(cn, lw=2.4, color=palette[4], label=r\"cell memory  $\\|c_t\\|$\"); axb.grid(False)\n",
    "ax[1].set_ylabel(\"write\", color=palette[0]); axb.set_ylabel(\"memory\", color=palette[4])\n",
    "ax[1].legend(fontsize=11, loc=\"upper left\"); axb.legend(fontsize=11, loc=\"lower right\")\n",
    "ax[2].plot(guess, lw=2.6, color=palette[2], label=r\"LSTM running guess  $\\mathrm{fc}(h_t)$\")\n",
    "ax[2].step(range(40), run, where=\"post\", color=\"0.3\", ls=\"--\", lw=1.8, label=\"true running sum\")\n",
    "ax[2].axhline(tgt, color=palette[3], lw=1, alpha=.5); ax[2].set_ylabel(\"predicted sum\"); ax[2].set_xlabel(\"time step\")\n",
    "ax[2].legend(fontsize=11, loc=\"upper left\")\n",
    "for a in (ax[1], ax[2]):\n",
    "    for tt in (t1, t2): a.axvline(tt, color=palette[3], ls=\":\", lw=1.5)\n",
    "ax[2].margins(x=0); plt.tight_layout(); plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "at0",
   "metadata": {},
   "source": [
    "# Recurrent attention"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "at1",
   "metadata": {},
   "source": [
    "## Seq2seq attention learns to align"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "at2",
   "metadata": {},
   "outputs": [],
   "source": [
    "import math, torch, torch.nn as nn, torch.nn.functional as F\n",
    "sns.set_style(\"whitegrid\"); palette = sns.color_palette()\n",
    "# A GRU encoder-decoder with Bahdanau (additive) attention, trained to REVERSE a sequence.\n",
    "# The attention weights show the decoder learning to align each output to the right input.\n",
    "V, Ls, BOS = 12, 10, 12\n",
    "def rev_batch(B, distinct=False):\n",
    "    src = torch.stack([torch.randperm(V)[:Ls] for _ in range(B)]) if distinct else torch.randint(0, V, (B, Ls))\n",
    "    return src, torch.flip(src, [1])\n",
    "class Attn(nn.Module):                                       # additive attention: score = v^T tanh(Wk h_enc + Wq h_dec)\n",
    "    def __init__(s, h, a=64): super().__init__(); s.Wk=nn.Linear(h,a); s.Wq=nn.Linear(h,a); s.v=nn.Linear(a,1)\n",
    "    def forward(s, dh, enc):\n",
    "        e = s.v(torch.tanh(s.Wk(enc) + s.Wq(dh).unsqueeze(1))).squeeze(-1)\n",
    "        a = torch.softmax(e, -1); return (a.unsqueeze(-1)*enc).sum(1), a\n",
    "class Seq2Seq(nn.Module):\n",
    "    def __init__(s, h=64, e=32):\n",
    "        super().__init__(); s.se=nn.Embedding(V,e); s.te=nn.Embedding(V+1,e)\n",
    "        s.enc=nn.GRU(e,h,batch_first=True); s.dec=nn.GRUCell(e+h,h); s.attn=Attn(h); s.out=nn.Linear(h+h,V)\n",
    "    def forward(s, src, tgt):\n",
    "        enc, hn = s.enc(s.se(src)); dh = hn[0]; B = src.size(0)\n",
    "        prev = torch.full((B,), BOS, dtype=torch.long); logits, atts = [], []\n",
    "        for t in range(Ls):                                  # teacher-forced decoding, recording attention each step\n",
    "            ctx, a = s.attn(dh, enc); dh = s.dec(torch.cat([s.te(prev), ctx], -1), dh)\n",
    "            logits.append(s.out(torch.cat([dh, ctx], -1))); atts.append(a); prev = tgt[:, t]\n",
    "        return torch.stack(logits, 1), torch.stack(atts, 1)\n",
    "torch.manual_seed(1); m = Seq2Seq(); opt = torch.optim.Adam(m.parameters(), 3e-3)\n",
    "for it in range(900):\n",
    "    src, tgt = rev_batch(64); lo, _ = m(src, tgt)\n",
    "    opt.zero_grad(); F.cross_entropy(lo.reshape(-1, V), tgt.reshape(-1)).backward(); opt.step()\n",
    "torch.manual_seed(7); s1, t1 = rev_batch(1, distinct=True)   # distinct symbols -> razor-sharp alignment, no ambiguity\n",
    "with torch.no_grad(): lo, att = m(s1, t1)\n",
    "A = att[0].numpy(); pred = lo[0].argmax(-1).tolist()\n",
    "print(\"input   (source)  :\", s1[0].tolist())\n",
    "print(\"target  (reversed):\", t1[0].tolist())\n",
    "print(\"output  (predicted):\", pred)\n",
    "fig, ax = plt.subplots(figsize=(7.2, 6.6)); ax.grid(False)\n",
    "im = ax.imshow(A, cmap=\"viridis\", vmin=0, vmax=1)\n",
    "ax.set_xticks(range(Ls)); ax.set_xticklabels(s1[0].tolist()); ax.set_yticks(range(Ls)); ax.set_yticklabels(pred)\n",
    "ax.set_xlabel(\"source position (input symbol)\"); ax.set_ylabel(\"decoder step (output symbol)\")\n",
    "ax.set_title(\"seq2seq attention trained to reverse a sequence:\\neach output attends to its mirror-image input (anti-diagonal)\")\n",
    "plt.colorbar(im, fraction=0.046); plt.tight_layout(); plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "at3",
   "metadata": {},
   "source": [
    "## A harder alignment: learning to sort\n",
    "\n",
    "The reversal alignment is purely *positional* -- output $t$ always wants input $L\\!-\\!1\\!-\\!t$. Attention can also align by **content**: trained to emit the *sorted* sequence, the decoder learns to point, at step $t$, to wherever the $t$-th smallest value happens to sit. That is a different, data-dependent permutation for every input, so the attention matrix is no longer a fixed diagonal."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "at4",
   "metadata": {},
   "outputs": [],
   "source": [
    "import math, torch, torch.nn as nn, torch.nn.functional as F\n",
    "# The same attention seq2seq, now trained to SORT the input. The alignment is no longer a fixed anti-diagonal:\n",
    "# step t must point to wherever the t-th smallest value sits -- a different permutation for every input.\n",
    "V, Ls, BOS = 12, 8, 12\n",
    "def sort_batch(B):\n",
    "    src = torch.stack([torch.randperm(V)[:Ls] for _ in range(B)]); tgt, _ = torch.sort(src, 1); return src, tgt\n",
    "class Attn(nn.Module):\n",
    "    def __init__(s, h, a=64): super().__init__(); s.Wk=nn.Linear(h,a); s.Wq=nn.Linear(h,a); s.v=nn.Linear(a,1)\n",
    "    def forward(s, dh, enc):\n",
    "        e = s.v(torch.tanh(s.Wk(enc) + s.Wq(dh).unsqueeze(1))).squeeze(-1)\n",
    "        a = torch.softmax(e, -1); return (a.unsqueeze(-1)*enc).sum(1), a\n",
    "class Seq2Seq(nn.Module):                                    # bidirectional encoder lets each state see the whole sequence\n",
    "    def __init__(s, h=64, e=32):\n",
    "        super().__init__(); s.se=nn.Embedding(V,e); s.te=nn.Embedding(V+1,e)\n",
    "        s.enc=nn.GRU(e,h,batch_first=True,bidirectional=True); s.dec=nn.GRUCell(e+2*h,2*h); s.attn=Attn(2*h); s.out=nn.Linear(4*h,V)\n",
    "    def forward(s, src, tgt):\n",
    "        enc, hn = s.enc(s.se(src)); dh = torch.cat([hn[0], hn[1]], -1); B = src.size(0)\n",
    "        prev = torch.full((B,), BOS, dtype=torch.long); logits, atts = [], []\n",
    "        for t in range(Ls):\n",
    "            ctx, a = s.attn(dh, enc); dh = s.dec(torch.cat([s.te(prev), ctx], -1), dh)\n",
    "            logits.append(s.out(torch.cat([dh, ctx], -1))); atts.append(a); prev = tgt[:, t]\n",
    "        return torch.stack(logits, 1), torch.stack(atts, 1)\n",
    "torch.manual_seed(1); m = Seq2Seq(); opt = torch.optim.Adam(m.parameters(), 2e-3)\n",
    "for it in range(1500):\n",
    "    src, tgt = sort_batch(64); lo, _ = m(src, tgt)\n",
    "    opt.zero_grad(); F.cross_entropy(lo.reshape(-1, V), tgt.reshape(-1)).backward(); opt.step()\n",
    "torch.manual_seed(3); s1, t1 = sort_batch(1)\n",
    "with torch.no_grad(): lo, att = m(s1, t1)\n",
    "A = att[0].numpy(); pred = lo[0].argmax(-1).tolist()\n",
    "print(\"input   (source) :\", s1[0].tolist())\n",
    "print(\"target  (sorted) :\", t1[0].tolist())\n",
    "print(\"output  (predicted):\", pred)\n",
    "fig, ax = plt.subplots(figsize=(7.2, 6.6)); ax.grid(False)\n",
    "im = ax.imshow(A, cmap=\"viridis\", vmin=0, vmax=1)\n",
    "ax.set_xticks(range(Ls)); ax.set_xticklabels(s1[0].tolist()); ax.set_yticks(range(Ls)); ax.set_yticklabels(pred)\n",
    "ax.set_xlabel(\"source position (input value)\"); ax.set_ylabel(\"decoder step (output = sorted)\")\n",
    "ax.set_title(\"seq2seq attention trained to SORT:\\nstep t attends to wherever the t-th smallest value sits (content-driven)\")\n",
    "plt.colorbar(im, fraction=0.046); plt.tight_layout(); plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "tf0",
   "metadata": {},
   "source": [
    "# Transformers"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "tf1",
   "metadata": {},
   "source": [
    "## Self-attention is content-based lookup\n",
    "\n",
    "Every token emits a **query**, a **key**, and a **value**. Attention compares each query to every key (a scaled dot product), turns the scores into weights with a softmax, and returns the weighted sum of values. Below, each item token carries *(key + value)* and the final token carries only the *query key*; the network is trained just to output the right value.\n",
    "\n",
    "**What to look for:** the y-axis labels each row by the token doing the attending. The bottom row is the query -- it lights up almost entirely on the one item whose key matches (a near one-hot row, boxed), while every other row just keeps each item attending to itself. The right-hand bar chart is that query row on its own: a single tall bar on the matching item. The value is then read off from the attended item, giving the prediction printed above."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "tf2",
   "metadata": {},
   "outputs": [],
   "source": [
    "import math, torch, torch.nn as nn, torch.nn.functional as F, matplotlib.patches as mpatches\n",
    "# One layer, one head of self-attention. A sequence of (key,value) items plus a query key; the model must\n",
    "# output the value whose key matches the query -- i.e. self-attention is a soft, content-addressed lookup.\n",
    "NK, NV, Li = 8, 8, 6\n",
    "def lookup_batch(B):\n",
    "    keys = torch.stack([torch.randperm(NK)[:Li] for _ in range(B)]); vals = torch.randint(0, NV, (B, Li))\n",
    "    qi = torch.randint(0, Li, (B,)); ar = torch.arange(B)\n",
    "    return keys, vals, keys[ar, qi], vals[ar, qi]\n",
    "class Lookup(nn.Module):\n",
    "    def __init__(s, d=64):\n",
    "        super().__init__(); s.k=nn.Embedding(NK,d); s.val=nn.Embedding(NV,d); s.qm=nn.Parameter(torch.randn(d)*.1)\n",
    "        s.pos=nn.Parameter(torch.randn(Li+1,d)*.1)\n",
    "        s.Wq=nn.Linear(d,d); s.Wk=nn.Linear(d,d); s.Wv=nn.Linear(d,d); s.head=nn.Linear(d,NV); s.d=d\n",
    "    def forward(s, keys, vals, qkey):\n",
    "        items = s.k(keys) + s.val(vals); query = (s.k(qkey) + s.qm).unsqueeze(1)   # items carry key+value, query carries the key\n",
    "        x = torch.cat([items, query], 1) + s.pos\n",
    "        att = torch.softmax(s.Wq(x) @ s.Wk(x).transpose(1, 2) / math.sqrt(s.d), -1) # scaled dot-product attention\n",
    "        return s.head((att @ s.Wv(x))[:, -1]), att\n",
    "torch.manual_seed(2); ml = Lookup(); opt = torch.optim.Adam(ml.parameters(), 3e-3)\n",
    "for it in range(700):\n",
    "    keys, vals, qkey, ans = lookup_batch(128); lo, _ = ml(keys, vals, qkey)\n",
    "    opt.zero_grad(); F.cross_entropy(lo, ans).backward(); opt.step()\n",
    "torch.manual_seed(5); keys, vals, qkey, ans = lookup_batch(1)\n",
    "with torch.no_grad(): lo, att = ml(keys, vals, qkey)\n",
    "A2 = att[0].numpy(); match = int((keys[0] == qkey[0]).float().argmax())\n",
    "print(\"input  items:\", [\"k%d=v%d\" % (keys[0,j], vals[0,j]) for j in range(Li)])\n",
    "print(\"input  query: k%d\" % qkey[0])\n",
    "print(\"output: attends to item #%d (k%d=v%d)  ->  predicted value %d   (true %d)\"\n",
    "      % (match, keys[0,match], vals[0,match], lo.argmax(-1).item(), ans.item()))\n",
    "lab = [\"k%d=v%d\" % (keys[0,j], vals[0,j]) for j in range(Li)] + [\"query k%d\" % qkey[0]]\n",
    "fig, ax = plt.subplots(1, 2, figsize=(14, 5.6))\n",
    "ax[0].grid(False); im = ax[0].imshow(A2, cmap=\"viridis\", vmin=0, vmax=1)\n",
    "ax[0].add_patch(mpatches.Rectangle((-.5, Li-.5), Li+1, 1, fill=False, edgecolor=palette[3], lw=2.5))  # boxed: the query row\n",
    "ax[0].set_xticks(range(Li+1)); ax[0].set_xticklabels(lab, rotation=45, ha=\"right\"); ax[0].set_yticks(range(Li+1)); ax[0].set_yticklabels(lab)\n",
    "ax[0].set_xlabel(\"attended item\"); ax[0].set_ylabel(\"query (token doing the attending)\")\n",
    "ax[0].set_title(\"full self-attention matrix\"); plt.colorbar(im, ax=ax[0], fraction=0.046)\n",
    "ax[1].bar(range(Li+1), A2[-1], color=[palette[3] if j == match else palette[0] for j in range(Li+1)])\n",
    "ax[1].set_xticks(range(Li+1)); ax[1].set_xticklabels(lab, rotation=45, ha=\"right\"); ax[1].set_ylabel(\"attention weight\")\n",
    "ax[1].set_title(\"the query row alone: one tall bar on the matching item  (predicted %d, true %d)\" % (lo.argmax(-1).item(), ans.item()))\n",
    "plt.tight_layout(); plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "tf7",
   "metadata": {},
   "source": [
    "## Self-attention can pool, not just pick\n",
    "\n",
    "The lookup produced a one-hot attention row. The very same mechanism can instead **average over many positions**: trained so that every token must output the mean value of its key-group, attention spreads its weight *uniformly* over all tokens that share its key -- a block-structured, \"soft group-by\"."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "tf8",
   "metadata": {},
   "outputs": [],
   "source": [
    "import math, torch, torch.nn as nn, torch.nn.functional as F\n",
    "# Self-attention need not pick a single position. Trained so each token must output the MEAN value of all\n",
    "# tokens sharing its key, attention spreads uniformly over the whole key-group -- a soft \"group-by\".\n",
    "K, N = 4, 10\n",
    "def grp_batch(B):\n",
    "    keys = torch.randint(0, K, (B, N)); vals = torch.rand(B, N)\n",
    "    oh = F.one_hot(keys, K).float()\n",
    "    means = (oh.transpose(1, 2) @ vals.unsqueeze(-1)).squeeze(-1) / oh.sum(1).clamp(min=1)\n",
    "    return keys, vals, torch.gather(means, 1, keys)\n",
    "class GroupAttn(nn.Module):\n",
    "    def __init__(s, d=48):\n",
    "        super().__init__(); s.ke=nn.Embedding(K,d); s.vp=nn.Linear(1,d); s.pos=nn.Parameter(torch.randn(N,d)*.1)\n",
    "        s.Wq=nn.Linear(d,d); s.Wk=nn.Linear(d,d); s.Wv=nn.Linear(d,d); s.head=nn.Linear(d,1); s.d=d\n",
    "    def forward(s, keys, vals):\n",
    "        x = s.ke(keys) + s.vp(vals.unsqueeze(-1)) + s.pos\n",
    "        att = torch.softmax(s.Wq(x) @ s.Wk(x).transpose(1, 2)/math.sqrt(s.d), -1)\n",
    "        return s.head(att @ s.Wv(x)).squeeze(-1), att\n",
    "torch.manual_seed(2); g = GroupAttn(); opt = torch.optim.Adam(g.parameters(), 3e-3)\n",
    "for it in range(1200):\n",
    "    keys, vals, tgt = grp_batch(128); pr, _ = g(keys, vals)\n",
    "    opt.zero_grad(); F.mse_loss(pr, tgt).backward(); opt.step()\n",
    "torch.manual_seed(8); keys, vals, tgt = grp_batch(1); order = torch.argsort(keys[0])   # sort by key so groups are contiguous\n",
    "keys, vals, tgt = keys[:, order], vals[:, order], tgt[:, order]\n",
    "with torch.no_grad(): pr, att = g(keys, vals)\n",
    "print(\"input  keys  :\", keys[0].tolist())\n",
    "print(\"input  values:\", [round(v, 2) for v in vals[0].tolist()])\n",
    "print(\"output (per-key mean):\", [round(v, 2) for v in pr[0].tolist()])\n",
    "A2 = att[0].numpy(); lab = [\"k%d\\n%.2f\" % (keys[0,j], vals[0,j]) for j in range(N)]\n",
    "fig, ax = plt.subplots(1, 2, figsize=(14, 5.8))\n",
    "ax[0].grid(False); im = ax[0].imshow(A2, cmap=\"viridis\", vmin=0, vmax=A2.max())\n",
    "ax[0].set_xticks(range(N)); ax[0].set_xticklabels(lab, fontsize=8); ax[0].set_yticks(range(N)); ax[0].set_yticklabels(lab, fontsize=8)\n",
    "ax[0].set_xlabel(\"attended position (key, value)\"); ax[0].set_ylabel(\"query position\")\n",
    "ax[0].set_title(\"each token attends uniformly to all tokens sharing its key\\n(block structure = soft group-by)\"); plt.colorbar(im, ax=ax[0], fraction=0.046)\n",
    "ax[1].plot(tgt[0].numpy(), \"o-\", lw=2, label=\"true group mean\"); ax[1].plot(pr[0].detach().numpy(), \"x--\", lw=2, label=\"attention output\")\n",
    "ax[1].set_xlabel(\"position\"); ax[1].set_ylabel(\"value\"); ax[1].set_title(\"the attention output equals the per-key average value\"); ax[1].legend()\n",
    "plt.tight_layout(); plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "tf3",
   "metadata": {},
   "source": [
    "## Positional encodings"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "tf4",
   "metadata": {},
   "outputs": [],
   "source": [
    "import math, torch\n",
    "# Self-attention is permutation-invariant, so a Transformer ADDS a positional signal to every token.\n",
    "d_model, L = 64, 60\n",
    "pos = torch.arange(L).unsqueeze(1); div = torch.exp(torch.arange(0, d_model, 2) * (-math.log(10000.0)/d_model))\n",
    "PE = torch.zeros(L, d_model); PE[:, 0::2] = torch.sin(pos*div); PE[:, 1::2] = torch.cos(pos*div)\n",
    "fig, ax = plt.subplots(1, 2, figsize=(15, 5))\n",
    "for a in ax: a.grid(False)\n",
    "im0 = ax[0].imshow(PE.T, aspect=\"auto\", cmap=\"RdBu_r\", vmin=-1, vmax=1)\n",
    "ax[0].set_xlabel(\"position in sequence\"); ax[0].set_ylabel(\"encoding dimension\"); ax[0].set_title(\"sinusoidal positional encoding\")\n",
    "plt.colorbar(im0, ax=ax[0], fraction=0.025)\n",
    "im1 = ax[1].imshow(PE @ PE.T, cmap=\"magma\")\n",
    "ax[1].set_xlabel(\"position $j$\"); ax[1].set_ylabel(\"position $i$\"); ax[1].set_title(r\"$PE_i\\cdot PE_j$ depends only on the offset $|i-j|$\")\n",
    "plt.colorbar(im1, ax=ax[1], fraction=0.046); plt.tight_layout(); plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "tf9",
   "metadata": {},
   "source": [
    "## Rotary positional encodings (RoPE)\n",
    "\n",
    "Instead of *adding* a positional vector (previous cell), RoPE *rotates* each pair of dimensions of the query and key by an angle proportional to the position -- fast for low dimensions, slow for high ones. The elegant consequence: the attention score between positions $m$ and $n$ ends up depending only on their **relative** offset $m-n$, not on the absolute positions."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "tf10",
   "metadata": {},
   "outputs": [],
   "source": [
    "import math, numpy as np, torch\n",
    "# RoPE: instead of ADDING a position vector, rotate each pair of dims of q and k by an angle proportional to\n",
    "# the position. The attention score <R_m q, R_n k> then depends only on the relative offset m - n.\n",
    "def rope(x, pos, base=10000.0):\n",
    "    d = x.shape[-1]; half = d//2\n",
    "    freqs = base ** (-torch.arange(half).float()/half)\n",
    "    ang = pos.unsqueeze(-1)*freqs if torch.is_tensor(pos) else pos*freqs\n",
    "    c, s = torch.cos(ang), torch.sin(ang)\n",
    "    x1, x2 = x[..., 0::2], x[..., 1::2]\n",
    "    y = torch.empty_like(x); y[..., 0::2] = x1*c - x2*s; y[..., 1::2] = x1*s + x2*c; return y\n",
    "d, L = 32, 32; half = d//2; freqs = (10000.0 ** (-torch.arange(half).float()/half)); pos = torch.arange(L).float()\n",
    "torch.manual_seed(0); q = torch.randn(d); k = torch.randn(d)\n",
    "Rq = rope(q.expand(L, d).clone(), pos); Rk = rope(k.expand(L, d).clone(), pos); score = (Rq @ Rk.T).numpy()\n",
    "fig, ax = plt.subplots(1, 3, figsize=(16, 5))\n",
    "ax[0].grid(False); th = np.linspace(0, 2*np.pi, 200); ax[0].plot(np.cos(th), np.sin(th), color=\"0.85\", lw=1)\n",
    "for bi, nm, mk in [(1, \"dim-pair 1 (fast)\", \"o\"), (4, \"dim-pair 4 (medium)\", \"s\"), (10, \"dim-pair 10 (slow)\", \"^\")]:\n",
    "    a = (pos*freqs[bi]).numpy(); sc = ax[0].scatter(np.cos(a), np.sin(a), c=pos.numpy(), cmap=\"viridis\", s=30, marker=mk, label=nm)\n",
    "ax[0].set_aspect(\"equal\"); ax[0].set_xlim(-1.25, 1.25); ax[0].set_ylim(-1.25, 1.25); ax[0].legend(fontsize=9, loc=\"lower left\")\n",
    "ax[0].set_title(\"RoPE rotates each dim-pair by angle = position x frequency\\n(low dims spin fast, high dims slow)\", fontsize=11)\n",
    "cb = plt.colorbar(sc, ax=ax[0], fraction=0.046); cb.set_label(\"position\")\n",
    "ax[1].grid(False); im = ax[1].imshow(score, cmap=\"magma\")\n",
    "ax[1].set_xlabel(\"key position $n$\"); ax[1].set_ylabel(\"query position $m$\")\n",
    "ax[1].set_title(r\"$\\langle R_m q,\\,R_n k\\rangle$ for fixed $q,k$:\" \"\\n\" r\"constant along each diagonal $\\Rightarrow$ depends only on $m-n$\", fontsize=11)\n",
    "plt.colorbar(im, ax=ax[1], fraction=0.046)\n",
    "ax[2].grid(True)\n",
    "for shift, col, st, lw in [(0, palette[0], \"-\", 4.5), (8, palette[3], \"--\", 2.0)]:\n",
    "    offs = np.arange(0, L-8); row = [score[shift, shift+o] for o in offs]\n",
    "    ax[2].plot(offs, row, st, color=col, lw=lw, label=\"query at position %d\" % shift)\n",
    "ax[2].set_xlabel(\"relative offset  $n-m$\"); ax[2].set_ylabel(r\"$\\langle R_m q, R_n k\\rangle$\")\n",
    "ax[2].set_title(\"same content at different absolute positions:\\nidentical score vs. relative offset\", fontsize=11); ax[2].legend(fontsize=10)\n",
    "plt.tight_layout(); plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "tf5",
   "metadata": {},
   "source": [
    "## Why attention is scaled by 1/√d"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "tf6",
   "metadata": {},
   "outputs": [],
   "source": [
    "import math, torch\n",
    "# Dot products of d-dimensional vectors have variance ~ d, so logits grow with d; the 1/sqrt(d) factor keeps\n",
    "# the softmax from saturating into a one-hot (which would kill gradients). We measure attention entropy vs d.\n",
    "dims = [2, 4, 8, 16, 32, 64, 128, 256, 512]; nkeys, nsamp = 16, 4000\n",
    "ent_s, ent_u = [], []\n",
    "for d in dims:\n",
    "    q = torch.randn(nsamp, d); K = torch.randn(nsamp, nkeys, d); logit = (K @ q.unsqueeze(-1)).squeeze(-1)\n",
    "    for scale, store in [(1/math.sqrt(d), ent_s), (1.0, ent_u)]:\n",
    "        p = torch.softmax(logit*scale, -1); store.append((-(p*torch.log(p+1e-9)).sum(-1).mean()/math.log(nkeys)).item())\n",
    "fig, ax = plt.subplots(figsize=(9, 5.5))\n",
    "ax.semilogx(dims, ent_s, \"-o\", lw=2.4, color=palette[0], label=r\"scaled by $1/\\sqrt{d}$\")\n",
    "ax.semilogx(dims, ent_u, \"-o\", lw=2.4, color=palette[3], label=\"unscaled\")\n",
    "ax.set_xlabel(\"key / query dimension $d$\"); ax.set_ylabel(\"attention entropy  (1 = uniform, 0 = one-hot)\")\n",
    "ax.set_ylim(-0.03, 1.03); ax.set_title(r\"without the $1/\\sqrt{d}$ scale, attention collapses to one-hot as $d$ grows\")\n",
    "ax.legend(fontsize=13); plt.tight_layout(); plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ae88b68a-7948-4378-9ade-bbb6d2d394fa",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}