{
 "nbformat": 4,
 "nbformat_minor": 5,
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "name": "python",
   "version": "3.11"
  },
  "colab": {
   "name": "week-07-h2-attention.ipynb",
   "provenance": []
  }
 },
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": "# W07 H2: Why Attention, and How It Works\n\nMachine Learning in Engineering Science | NCKU + LAiMM | Fall 2026\n\n**Teaching date:** Monday 19 October 2026. **Guided notebook time:** 25 minutes within a 60-minute module.\n\n### What you will be able to explain\n- Explain the fixed-context bottleneck and query-dependent retrieval.\n- Compute scores, softmax weights, and a weighted sum by hand.\n- Distinguish masks, attention weights, and causal explanations.\n\n### Start here, even if this is your first notebook\n1. In Colab choose **Runtime > Change runtime type > CPU**. In Jupyter choose a Python 3 kernel.\n2. Run cells from top to bottom with **Shift+Enter**. Read the Markdown before each code cell.\n3. The setup installs only missing free packages. A network connection is needed for first installation; the data itself is generated locally.\n4. If a variable is undefined, restart the kernel and run all earlier cells. Do not skip a setup cell.\n5. Save a personal copy before changing parameters. Every core cell is complete: there are no hidden files, tokens, paid services, or unfinished code blocks.\n\n**Data:** synthetic teaching data, not engineering measurements. Record your actual output; a model is not required to outperform a baseline.\n**Practice policy:** in-class exercise and optional Homework extension, not a new weekly graded assignment.\n",
   "id": "cell-000"
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": "## Mechanism reference: compare, normalize, retrieve\n\nFor $Q$ of shape $(N_q,d_k)$, $K$ of shape $(N_k,d_k)$, and $V$ of shape $(N_k,d_v)$:\n\n$$S=QK^\\top/\\sqrt{d_k},\\quad A=\\operatorname{softmax}_{\\mathrm{keys}}(S+M),\\quad C=AV.$$\n\nMask $M$ is zero for allowed entries and negative infinity for blocked entries. Rows of $A$ are receiving queries and columns are source keys. Every valid row sums to one. $C$ has shape $(N_q,d_v)$. Keep key rows paired with their value rows.\n\nAttention addresses a final-state encoder's fixed-context bottleneck by retrieving earlier states for each output query. It can be attached to an LSTM or GRU. Self-attention derives Q,K,V from one sequence; cross-attention derives queries from one stream and keys/values from another.\n\n",
   "id": "cell-001"
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": "## 1. Set up a numerical experiment\n\nAttention can be understood without training a large model. First we use NumPy so each intermediate value is visible.\n",
   "id": "cell-002"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "metadata": {},
   "source": "import sys, subprocess, importlib.util, platform, json, random, hashlib\nfrom pathlib import Path\nrequired = {\"numpy\": \"numpy\", \"matplotlib\": \"matplotlib\", \"sklearn\": \"scikit-learn\"}\nmissing = [package for module, package in required.items() if importlib.util.find_spec(module) is None]\nif missing:\n    subprocess.check_call([sys.executable, \"-m\", \"pip\", \"install\", *missing])\nimport numpy as np\nimport matplotlib.pyplot as plt\nSEED = 2026\nrandom.seed(SEED)\nnp.random.seed(SEED)\nrng = np.random.default_rng(SEED)\nARTIFACTS = Path(\"artifacts\")\nARTIFACTS.mkdir(exist_ok=True)\nprint(\"Python:\", platform.python_version(), \"NumPy:\", np.__version__, \"Seed:\", SEED)\n",
   "id": "cell-003"
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": "**Checkpoint.** Run this notebook independently of H1. No saved H1 model is required.\n",
   "id": "cell-004"
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": "## 2. Start with a query and an addressable memory\n\nQ asks what to retrieve, K describes how each memory item matches the query, and V contains the information to combine. These numbers are hand-designed teaching vectors, not learned semantic embeddings. Q has shape (1,2), K is (3,2), and V is (3,2).\n",
   "id": "cell-005"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "metadata": {},
   "source": "Q = np.array([[1.,0.]])\nK = np.array([[1.,0.],[0.,1.],[1.,1.]])\nV = np.array([[10.,0.],[0.,10.],[10.,10.]])\nlabels = [\"memory A\",\"memory B\",\"memory C\"]\nprint(\"Q:\",Q,\"\\nK:\\n\",K,\"\\nV:\\n\",V)\nplt.scatter(K[:,0],K[:,1],s=120)\nfor name, point in zip(labels,K): plt.annotate(name,point,xytext=(5,7),textcoords=\"offset points\")\nplt.arrow(0,0,Q[0,0],Q[0,1],width=.01,color=\"tab:red\",length_includes_head=True)\nplt.xlim(-.2,1.4); plt.ylim(-.2,1.4); plt.xlabel(\"Key dimension 1\"); plt.ylabel(\"Key dimension 2\"); plt.show()\n",
   "id": "cell-006"
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": "**Checkpoint.** A and C have the same dot product with this query. A key is an address; its associated value need not equal the key.\n",
   "id": "cell-007"
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": "## 3. Calculate scaled scores\n\nFor each query, compute its dot product with every key, then divide by sqrt(d_k). This controls score scale as the key dimension grows under a unit-variance coordinate assumption. It does not normalize vector lengths.\n",
   "id": "cell-008"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "metadata": {},
   "source": "scores = Q @ K.T / np.sqrt(K.shape[-1])\nprint(\"Raw dot products:\",Q@K.T)\nprint(\"Scaled scores:\",scores)\nassert np.allclose(scores,[[1/np.sqrt(2),0,1/np.sqrt(2)]])\n",
   "id": "cell-009"
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": "**Checkpoint.** The scaled score row is approximately [0.7071, 0, 0.7071]. Scores are not yet probabilities and may be negative.\n",
   "id": "cell-010"
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": "## 4. Convert scores into a row of weights\n\nSoftmax subtracts the row maximum for numerical stability, exponentiates, and divides by the row sum. Subtracting a shared constant changes no probability. Each query row must sum to 1 over the available keys.\n",
   "id": "cell-011"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "metadata": {},
   "source": "def softmax(z):\n    shifted = z-z.max(axis=-1,keepdims=True)\n    exp = np.exp(shifted)\n    return exp/exp.sum(axis=-1,keepdims=True)\nweights = softmax(scores)\nprint(\"Weights:\",weights,\"Row sum:\",weights.sum(axis=-1))\nassert np.allclose(weights.sum(axis=-1),1)\nplt.bar(labels,weights[0],color=[\"#1e8798\",\"#7f2635\",\"#f3b23a\"])\nplt.ylim(0,1); plt.ylabel(\"Attention weight\"); plt.show()\n",
   "id": "cell-012"
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": "**Checkpoint.** The approximate weights are [0.4011, 0.1978, 0.4011]. B still has positive weight even though its raw score is zero.\n",
   "id": "cell-013"
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": "## 5. Retrieve a weighted combination of values\n\nThe output is weights @ V. Each output dimension combines the same weights with the corresponding value coordinates. There is no argmax selection: multiple values contribute.\n",
   "id": "cell-014"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "metadata": {},
   "source": "context = weights @ V\nprint(\"Weighted contributions:\\n\",weights[0,:,None]*V)\nprint(\"Context vector:\",context)\nassert context.shape == (1,2)\nassert np.allclose(context,(weights[0,:,None]*V).sum(axis=0,keepdims=True))\n",
   "id": "cell-015"
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": "**Checkpoint.** The context is approximately [8.0222, 5.9889]. Explain which memory items contribute to each coordinate before moving on.\n",
   "id": "cell-016"
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": "## 6. Change only the query\n\nThe memory stays fixed but the task changes. Try Q=[0,1] and Q=[-1,0]. Compute your predicted ranking before running. Temperature below 1 sharpens and above 1 softens the distribution; standard scaled dot-product attention uses temperature=1.\n",
   "id": "cell-017"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "metadata": {},
   "source": "queries = np.array([[1.,0.],[0.,1.],[-1.,0.]])\ntemperature = 1.0\nassert temperature > 0\nW = softmax((queries@K.T/np.sqrt(2))/temperature)\nprint(\"Contexts:\\n\",W@V)\nfig,ax=plt.subplots(figsize=(6,3))\nim=ax.imshow(W,cmap=\"YlGnBu\",vmin=0,vmax=1)\nax.set_xticks(range(3),labels); ax.set_yticks(range(3),[\"q=[1,0]\",\"q=[0,1]\",\"q=[-1,0]\"])\nfor row in range(3):\n    for col in range(3): ax.text(col,row,f\"{W[row,col]:.2f}\",ha=\"center\",va=\"center\",color=\"black\")\nfig.colorbar(im,ax=ax); plt.show()\n",
   "id": "cell-018"
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": "**Checkpoint.** A query with negative first coordinate reduces the weights of keys with positive first coordinate. Attention retrieval depends on the query, not on one universal ranking.\n",
   "id": "cell-019"
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": "## 7. Add causal masking before softmax\n\nFor a token at position t, a causal decoder may access positions <=t, including itself. Set future scores to negative infinity before softmax; then their exponentials are zero. Every row must retain at least one allowed key.\n",
   "id": "cell-020"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "metadata": {},
   "source": "X_demo = np.array([[1.,0.],[0.,1.],[1.,1.]])\nself_scores = X_demo@X_demo.T/np.sqrt(2)\nallowed = np.tril(np.ones((3,3),dtype=bool))\nmasked_scores = np.where(allowed,self_scores,-np.inf)\ncausal_weights = softmax(masked_scores)\nprint(\"Allowed (query rows, key columns):\\n\",allowed.astype(int))\nprint(\"Causal attention:\\n\",causal_weights)\nassert np.all(causal_weights[~allowed] == 0)\nassert np.allclose(causal_weights.sum(axis=-1),1)\nplt.imshow(causal_weights,cmap=\"YlGnBu\",vmin=0,vmax=1)\nplt.xlabel(\"Key/source position\"); plt.ylabel(\"Query/destination position\"); plt.colorbar(); plt.show()\n",
   "id": "cell-021"
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": "**Checkpoint.** The first row must be [1,0,0]. Do not apply softmax before masking, and do not confuse the direction of rows and columns.\n",
   "id": "cell-022"
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": "## 8. Separate padding from future masking\n\nA padding mask hides nonexistent tokens, while a causal mask hides future tokens. Combine the two permissions. For padded query rows, the output is normally excluded from the loss rather than interpreted as a real prediction.\n",
   "id": "cell-023"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "metadata": {},
   "source": "is_real_key = np.array([True,True,False])\ncombined_allowed = allowed & is_real_key[None,:]\nassert combined_allowed.any(axis=-1).all()\ncombined_weights = softmax(np.where(combined_allowed,self_scores,-np.inf))\nprint(combined_weights)\nassert np.all(combined_weights[:,2] == 0)\n",
   "id": "cell-024"
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": "**Checkpoint.** No query can retrieve the padded final key. An all-masked row would make softmax undefined; prevent that case rather than silently treating NaN as an answer.\n",
   "id": "cell-025"
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": "## 9. Verify the same calculation with PyTorch\n\nThe high-level SDPA API uses True for allowed entries in a Boolean mask. Other PyTorch APIs use different mask conventions; always read the called API. We use an explicit implementation for comparison and disable dropout.\n",
   "id": "cell-026"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "metadata": {},
   "source": "if importlib.util.find_spec(\"torch\") is None:\n    subprocess.check_call([sys.executable, \"-m\", \"pip\", \"install\", \"torch\"])\nimport torch\nfrom torch import nn\nimport copy, time\ntorch.set_num_threads(1)\ntorch.manual_seed(SEED)\ntorch.use_deterministic_algorithms(True)\ndevice = torch.device(\"cpu\")\nprint(\"PyTorch:\", torch.__version__, \"Device:\", device)\n\nimport torch.nn.functional as F\nq=torch.tensor(Q,dtype=torch.float32)[None,None,:,:]\nk=torch.tensor(K,dtype=torch.float32)[None,None,:,:]\nv=torch.tensor(V,dtype=torch.float32)[None,None,:,:]\ntorch_context = F.scaled_dot_product_attention(q,k,v,dropout_p=0.0)\nassert np.allclose(torch_context.numpy()[0,0],context,atol=1e-5)\nprint(\"NumPy and PyTorch agree:\",torch_context[0,0])\n",
   "id": "cell-027"
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": "**Checkpoint.** The parity assertion must pass. This is an implementation check for the chosen example, not evidence that attention has learned language.\n",
   "id": "cell-028"
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": "## 10. Test what an attention plot can explain\n\nWeights describe a computation, but causal importance requires an intervention. Hold Q and K fixed and change one V; weights remain unchanged while output changes. Then remove a memory item by masking it and observe that normalization also changes.\n",
   "id": "cell-029"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "metadata": {},
   "source": "V_changed=V.copy(); V_changed[1]=[100.,100.]\nchanged_context=weights@V_changed\nremoved_weights=softmax(np.where(np.array([[True,False,True]]),scores,-np.inf))\nprint(\"Original weights:\",weights)\nprint(\"Same weights, changed values:\",changed_context)\nprint(\"Remove B and renormalize:\",removed_weights@V)\nassert not np.allclose(changed_context,context)\nevidence={\"Q\":Q.tolist(),\"K\":K.tolist(),\"V\":V.tolist(),\"scores\":scores.tolist(),\n          \"weights\":weights.tolist(),\"context\":context.tolist(),\n          \"causal_weights\":causal_weights.tolist(),\"seed\":SEED,\n          \"interpretation\":\"Attention weights are routing weights, not proof of causal importance.\"}\n(ARTIFACTS/\"w07-h2-attention.json\").write_text(json.dumps(evidence,indent=2))\n",
   "id": "cell-030"
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": "**Checkpoint.** State exactly what the heatmap supports: where the chosen computation routes information. Do not claim it proves why a real-world outcome occurred.\n",
   "id": "cell-031"
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": "## 11. Independent practice and solutions\n1. Set temperature to 0.5 and 2.0 in Step 6. Predict and then compare the sharpness, keeping Q, K, V fixed.\n2. Replace all keys by the same vector. **Expected:** every available key gets equal weight, regardless of query; masked keys still get zero.\n3. Make one value very large. **Expected:** the context may change greatly with no change to attention weights.\n4. Explain why additive attention can be attached to an LSTM encoder-decoder. **Answer:** a decoder query can score all encoder states; attention does not require removing recurrence.\n5. Explain self-attention versus cross-attention. **Answer:** self-attention derives Q,K,V from one sequence; cross-attention derives Q from one stream and K,V from another.\n\n**Stop when:** your evidence JSON contains scores, weights, context, and a causal-mask assertion; explain all matrix axes to a partner.\n\n**Bridge to H3:** a Transformer combines attention with positional information, a tokenwise feed-forward network, residual paths, and normalization. Attention alone is not the complete architecture.\n\n### Sources\n- [Bahdanau et al. (2014/2015), joint alignment and translation](https://arxiv.org/abs/1409.0473)\n- [Vaswani et al. (2017), Attention Is All You Need](https://arxiv.org/abs/1706.03762)\n- [PyTorch scaled dot-product attention API](https://docs.pytorch.org/docs/stable/generated/torch.nn.functional.scaled_dot_product_attention.html)\n- [3Blue1Brown attention lesson](https://www.3blue1brown.com/lessons/attention/)\n",
   "id": "cell-032"
  }
 ]
}