{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "c001",
   "metadata": {},
   "source": [
    "# Step 4 \u2014 A Causal Self-Attention Head\n",
    "\n",
    "**Solution notebook.** No training this week. The deliverables are\n",
    "Lecture 4's propositions **executed as assertions**, plus pictures of\n",
    "attention matrices on real text. Runs in seconds.\n",
    "\n",
    "Every assertion below is annotated with the result it verifies."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c002",
   "metadata": {},
   "source": [
    "## 4.1 The head\n",
    "\n",
    "$$\\operatorname{Attn}(X)=\\operatorname{softmax}\\!\\Bigl(\\operatorname{mask}\\bigl(\\tfrac{XW_Q(XW_K)^\\top}{\\sqrt{d_k}}\\bigr)\\Bigr)XW_V$$\n",
    "\n",
    "`register_buffer` holds the causal mask as non-parameter state, so it\n",
    "moves with `.to(device)` but is not trained. We stash `A` for plotting."
   ]
  },
  {
   "cell_type": "code",
   "id": "c003",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "import torch, torch.nn as nn, torch.nn.functional as F\n",
    "import matplotlib.pyplot as plt\n",
    "torch.manual_seed(0)\n",
    "\n",
    "class Head(nn.Module):\n",
    "    def __init__(self, d, d_head, T_max, causal=True):\n",
    "        super().__init__()\n",
    "        self.key   = nn.Linear(d, d_head, bias=False)\n",
    "        self.query = nn.Linear(d, d_head, bias=False)\n",
    "        self.value = nn.Linear(d, d_head, bias=False)\n",
    "        self.causal = causal\n",
    "        self.register_buffer('tril', torch.tril(torch.ones(T_max, T_max)))\n",
    "\n",
    "    def forward(self, x):                      # x: (B, T, d)\n",
    "        B, T, d = x.shape\n",
    "        q, k, v = self.query(x), self.key(x), self.value(x)\n",
    "        scores = q @ k.transpose(-2, -1) / k.shape[-1]**0.5    # (B,T,T)\n",
    "        if self.causal:\n",
    "            scores = scores.masked_fill(self.tril[:T, :T] == 0, float('-inf'))\n",
    "        A = F.softmax(scores, dim=-1)\n",
    "        self.A = A.detach()\n",
    "        return A @ v                                            # (B,T,d_head)\n",
    "\n",
    "d, dh, T, B = 32, 16, 8, 4\n",
    "head = Head(d, dh, T)\n",
    "x = torch.randn(B, T, d)\n",
    "out = head(x)\n",
    "print('out', out.shape, '| A', head.A.shape)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c004",
   "metadata": {},
   "source": [
    "## 4.2 The proofs, as assertions"
   ]
  },
  {
   "cell_type": "code",
   "id": "c005",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "A = head.A\n",
    "\n",
    "# Prop 3.1 \u2014 rows of A lie in the simplex\n",
    "assert (A >= 0).all()\n",
    "assert torch.allclose(A.sum(-1), torch.ones(B, T), atol=1e-6)\n",
    "print(f'Prop 3.1  simplex rows        OK  (max |rowsum-1| = '\n",
    "      f'{(A.sum(-1)-1).abs().max():.1e})')\n",
    "\n",
    "# Cor 3.2 \u2014 output lies in the convex hull of the value vectors\n",
    "v = head.value(x)\n",
    "lo, hi = v.min(dim=1).values, v.max(dim=1).values\n",
    "assert ((out >= lo[:, None, :] - 1e-6) & (out <= hi[:, None, :] + 1e-6)).all()\n",
    "print('Cor 3.2   output in conv hull  OK')\n",
    "\n",
    "# Cor 3.3 \u2014 every allowed weight is strictly positive\n",
    "allowed = torch.tril(torch.ones(T, T)).bool()\n",
    "assert (A[0][allowed] > 0).all()\n",
    "print(f'Cor 3.3   allowed weights > 0  OK  (min = {A[0][allowed].min():.2e})')\n",
    "\n",
    "# Prop 5.1 \u2014 the mask makes A lower-triangular\n",
    "assert (A.triu(diagonal=1) == 0).all()\n",
    "print('Prop 5.1  A lower-triangular   OK')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c006",
   "metadata": {},
   "source": [
    "### Causality, two ways\n",
    "\n",
    "By perturbation (positions before the change must be **bitwise**\n",
    "identical, not merely close) and by gradient \u2014 the derivative form of\n",
    "the same statement."
   ]
  },
  {
   "cell_type": "code",
   "id": "c007",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "x2 = x.clone(); x2[:, 5, :] = torch.randn(B, d)\n",
    "out2 = head(x2)\n",
    "assert torch.equal(out[:, :5, :], out2[:, :5, :])\n",
    "assert not torch.allclose(out[:, 5:, :], out2[:, 5:, :])\n",
    "print('Prop 5.1  perturbation test    OK  (positions 0-4 bitwise identical)')\n",
    "\n",
    "xg = x.clone().requires_grad_(True)\n",
    "head(xg)[0, 3].sum().backward()\n",
    "assert (xg.grad[0, 4:] == 0).all()\n",
    "assert (xg.grad[0, :4].abs().sum(-1) > 0).all()\n",
    "print('Prop 5.1  gradient test        OK  (d out_3 / d x_j = 0 for j > 3)')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c008",
   "metadata": {},
   "source": [
    "### Lemma 4.5 \u2014 why the $\\sqrt{d_k}$\n",
    "\n",
    "Raw scores have standard deviation $\\sqrt{d_k}$, so without the scaling\n",
    "the softmax saturates to one-hot as $d_k$ grows \u2014 its Hessian\n",
    "$\\operatorname{diag}(p)-pp^\\top$ goes numerically to zero and no\n",
    "gradient flows through the attention weights."
   ]
  },
  {
   "cell_type": "code",
   "id": "c009",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "print(f\"{'d_k':>5} {'raw std':>9} {'sqrt(d_k)':>10} {'scaled std':>11}\"\n",
    "      f\" {'max wt raw':>11} {'max wt scaled':>14}\")\n",
    "for dk in (16, 64, 256):\n",
    "    q, k = torch.randn(1, T, dk), torch.randn(1, T, dk)\n",
    "    raw = q @ k.transpose(-2, -1)\n",
    "    sc = raw / dk**0.5\n",
    "    print(f'{dk:>5} {raw.std():>9.2f} {dk**0.5:>10.1f} {sc.std():>11.2f}'\n",
    "          f' {F.softmax(raw,-1).max():>11.3f} {F.softmax(sc,-1).max():>14.3f}')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c010",
   "metadata": {},
   "source": [
    "Reference: raw std 4.12 / 7.95 / 18.00 against $\\sqrt{d_k}$ = 4 / 8 /\n",
    "16, and the unscaled max weight is **1.000 already at $d_k=16$** \u2014\n",
    "fully saturated to three decimals."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c011",
   "metadata": {},
   "source": [
    "### Exercise 4 \u2014 masking before vs after softmax\n",
    "\n",
    "Zeroing $A$ *after* the softmax (instead of $-\\infty$ before) destroys\n",
    "row-stochasticity, and with it Corollary 3.3's guarantee."
   ]
  },
  {
   "cell_type": "code",
   "id": "c012",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "s = torch.randn(1, T, T)\n",
    "right = F.softmax(s.masked_fill(torch.tril(torch.ones(T,T))==0, float('-inf')), -1)\n",
    "wrong = F.softmax(s, -1) * torch.tril(torch.ones(T, T))\n",
    "print(f'-inf before softmax : row sums {right.sum(-1)[0].min():.3f} '\n",
    "      f'.. {right.sum(-1)[0].max():.3f}')\n",
    "print(f'zeroing after       : row sums {wrong.sum(-1)[0].min():.3f} '\n",
    "      f'.. {wrong.sum(-1)[0].max():.3f}   <- not stochastic')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c013",
   "metadata": {},
   "source": [
    "## 4.3 Equivariance \u2014 and the trap\n",
    "\n",
    "**Theorem 6.1 is about $\\operatorname{Attn}$ as a function of $X$**, and\n",
    "it stays true however $X$ was built \u2014 including after adding positional\n",
    "embeddings. Permuting the *rows of $X$* therefore will **not** show you\n",
    "the symmetry breaking, and most people set this experiment up that way\n",
    "the first time.\n",
    "\n",
    "Equivariance is lost one step earlier, in the embedding map\n",
    "$\\iota(x)_t = E_{x_t} + p_t$, because the positional term does not\n",
    "travel with the permuted token. You must permute **tokens**."
   ]
  },
  {
   "cell_type": "code",
   "id": "c014",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "hn = Head(d, dh, T, causal=False)      # no mask, so S_T acts\n",
    "perm = torch.randperm(T)\n",
    "\n",
    "# (a) the layer is equivariant in X ...\n",
    "assert torch.allclose(hn(x[:, perm, :]), hn(x)[:, perm, :], atol=1e-5)\n",
    "print('Thm 6.1   equivariant in X                     OK')\n",
    "\n",
    "# (b) ... and STAYS equivariant when X contains positional vectors\n",
    "xp = x + nn.Embedding(T, d)(torch.arange(T))[None]\n",
    "assert torch.allclose(hn(xp[:, perm, :]), hn(xp)[:, perm, :], atol=1e-5)\n",
    "print('          still equivariant with pos-emb in X   OK  <- the trap')\n",
    "\n",
    "# (c) the composite TOKEN -> output: equivariant without positions ...\n",
    "tok, pos = nn.Embedding(V := 65, d), nn.Embedding(T, d)\n",
    "idx = torch.randint(0, V, (1, T))\n",
    "P = pos(torch.arange(T))[None]\n",
    "a1, b1 = hn(tok(idx[:, perm])), hn(tok(idx))[:, perm, :]\n",
    "assert torch.allclose(a1, b1, atol=1e-5)\n",
    "print(f'Thm 6.1   token-perm, NO pos-emb               OK  '\n",
    "      f'(diff {(a1-b1).abs().max():.1e})')\n",
    "\n",
    "# (d) ... and BROKEN with them. This is Corollary 6.2.\n",
    "a2, b2 = hn(tok(idx[:, perm]) + P), hn(tok(idx) + P)[:, perm, :]\n",
    "assert not torch.allclose(a2, b2, atol=1e-3)\n",
    "print(f'Cor 6.2   token-perm, WITH pos-emb        BROKEN  '\n",
    "      f'(diff {(a2-b2).abs().max():.1e})')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c015",
   "metadata": {},
   "source": [
    "## 4.4 Attention patterns on real text\n",
    "\n",
    "Untrained weights, five seeds. The structure you see is a *prior*: row\n",
    "1 has only one position available so $A_{11}=1$ by the simplex\n",
    "constraint, row 2 splits between two, and so on \u2014 the mask alone forces\n",
    "early positions to carry weight. In Step 6 you will plot these same\n",
    "pictures after training and find actual algorithms in them."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c016",
   "metadata": {},
   "source": [
    "## Setup: data and tokenizer (Step 1's solution, reproduced)\n",
    "\n",
    "Every notebook in this series is self-contained: it re-creates what it\n",
    "needs from earlier steps in one compact cell, so you can run it top to\n",
    "bottom without opening the others. On Colab, uncomment the download."
   ]
  },
  {
   "cell_type": "code",
   "id": "c017",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "# !wget -q https://raw.githubusercontent.com/karpathy/char-rnn/master/data/tinyshakespeare/input.txt\n",
    "import torch\n",
    "\n",
    "with open('input.txt') as f:\n",
    "    text = f.read()\n",
    "vocab = sorted(set(text))\n",
    "V = len(vocab)\n",
    "stoi = {ch: i for i, ch in enumerate(vocab)}\n",
    "itos = {i: ch for i, ch in enumerate(vocab)}\n",
    "encode = lambda s: [stoi[c] for c in s]\n",
    "decode = lambda ids: ''.join(itos[i] for i in ids)\n",
    "ids = torch.tensor(encode(text), dtype=torch.long)\n",
    "\n",
    "assert len(text) == 1_115_394 and V == 65\n",
    "print(f'{len(text):,} characters, vocab {V}')"
   ]
  },
  {
   "cell_type": "code",
   "id": "c018",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "Tc = 48\n",
    "tok_emb = nn.Embedding(V, 32)\n",
    "pos_emb = nn.Embedding(Tc, 32)\n",
    "idx = ids[:Tc][None, :]\n",
    "chars = [decode([i]) for i in idx[0].tolist()]\n",
    "labels = [{'\\n': '\\\\n', ' ': '\u2423'}.get(c, c) for c in chars]\n",
    "\n",
    "fig, axes = plt.subplots(1, 5, figsize=(16, 3.4))\n",
    "for s, ax in enumerate(axes):\n",
    "    torch.manual_seed(s)\n",
    "    h = Head(32, 16, Tc)\n",
    "    xin = tok_emb(idx) + pos_emb(torch.arange(Tc))[None]\n",
    "    h(xin)\n",
    "    ax.imshow(h.A[0], cmap='Blues')\n",
    "    ax.set_title(f'seed {s}', fontsize=10); ax.set_xticks([]); ax.set_yticks([])\n",
    "plt.suptitle('untrained causal attention, five seeds '\n",
    "             '(strictly lower-triangular, rows sum to 1)')\n",
    "plt.tight_layout()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c019",
   "metadata": {},
   "source": [
    "## 4.5 Attention as kernel smoothing\n",
    "\n",
    "Set $W_Q = W_K = 0$ and attention becomes a running mean; hardwire\n",
    "scores by *position*, $S_{ij} = -(i-j)^2/2s^2$, and it becomes a\n",
    "Gaussian smoother of bandwidth $s$ \u2014 i.e. a **convolution**. This is\n",
    "Nadaraya\u2013Watson along the sequence. Content-dependence is the only\n",
    "thing real attention adds on top, which is exactly why it is powerful:\n",
    "the 'bandwidth' and 'location' can vary per token and per input."
   ]
  },
  {
   "cell_type": "code",
   "id": "c020",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "i = torch.arange(Tc).float()\n",
    "mask = torch.tril(torch.ones(Tc, Tc))\n",
    "fig, axes = plt.subplots(1, 4, figsize=(14, 3.4))\n",
    "for ax, s in zip(axes, [0.5, 2.0, 8.0, float('inf')]):\n",
    "    if s == float('inf'):\n",
    "        S = torch.zeros(Tc, Tc); title = 'W_Q=W_K=0 (running mean)'\n",
    "    else:\n",
    "        S = -(i[:, None] - i[None, :])**2 / (2 * s**2); title = f'bandwidth s={s}'\n",
    "    Ak = F.softmax(S.masked_fill(mask == 0, float('-inf')), dim=-1)\n",
    "    ax.imshow(Ak, cmap='Blues'); ax.set_title(title, fontsize=10)\n",
    "    ax.set_xticks([]); ax.set_yticks([])\n",
    "plt.suptitle('positional-kernel attention: the degenerate, content-blind case')\n",
    "plt.tight_layout()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c021",
   "metadata": {},
   "source": [
    "## Exercise 7 \u2014 rank of the score matrix\n",
    "\n",
    "$S = QK^\\top/\\sqrt{d_k}$ factors through $\\mathbb R^{d_k}$, so\n",
    "$\\operatorname{rank} S \\le d_k$ regardless of $T$. A single head can\n",
    "only produce a low-rank family of patterns \u2014 which is the reason\n",
    "Lecture 5 runs several in parallel."
   ]
  },
  {
   "cell_type": "code",
   "id": "c022",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "big = Head(64, 8, 32, causal=False)\n",
    "xb = torch.randn(1, 32, 64)\n",
    "big(xb)\n",
    "S = (big.query(xb) @ big.key(xb).transpose(-2, -1))[0]\n",
    "print(f'score matrix is {tuple(S.shape)} with d_k = 8; '\n",
    "      f'rank = {torch.linalg.matrix_rank(S).item()}')\n",
    "assert torch.linalg.matrix_rank(S).item() <= 8"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c023",
   "metadata": {},
   "source": [
    "\u2192 Continue with [Step 5](../project/step-5): multiply this head, wrap\n",
    "it in residuals and LayerNorm, and count every parameter."
   ]
  }
 ],
 "nbformat": 4,
 "nbformat_minor": 5,
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "name": "python",
   "version": "3.11"
  }
 }
}