W07 H2: Why Attention, and How It Works

Machine Learning in Engineering Science | NCKU + LAiMM | Fall 2026

Teaching date: Monday 19 October 2026. Guided notebook time: 25 minutes within a 60-minute module.

What you will be able to explain

- Explain the fixed-context bottleneck and query-dependent retrieval.

- Compute scores, softmax weights, and a weighted sum by hand.

- Distinguish masks, attention weights, and causal explanations.

Start here, even if this is your first notebook

1. In Colab choose Runtime > Change runtime type > CPU. In Jupyter choose a Python 3 kernel.

2. Run cells from top to bottom with Shift+Enter. Read the Markdown before each code cell.

3. The setup installs only missing free packages. A network connection is needed for first installation; the data itself is generated locally.

4. If a variable is undefined, restart the kernel and run all earlier cells. Do not skip a setup cell.

5. Save a personal copy before changing parameters. Every core cell is complete: there are no hidden files, tokens, paid services, or unfinished code blocks.

Data: synthetic teaching data, not engineering measurements. Record your actual output; a model is not required to outperform a baseline.

Practice policy: in-class exercise and optional Homework extension, not a new weekly graded assignment.

Mechanism reference: compare, normalize, retrieve

For $Q$ of shape $(N_q,d_k)$, $K$ of shape $(N_k,d_k)$, and $V$ of shape $(N_k,d_v)$:

$$S=QK^\top/\sqrt{d_k},\quad A=\operatorname{softmax}_{\mathrm{keys}}(S+M),\quad C=AV.$$

Mask $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.

Attention 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.

1. Set up a numerical experiment

Attention can be understood without training a large model. First we use NumPy so each intermediate value is visible.

import sys, subprocess, importlib.util, platform, json, random, hashlib
from pathlib import Path
required = {"numpy": "numpy", "matplotlib": "matplotlib", "sklearn": "scikit-learn"}
missing = [package for module, package in required.items() if importlib.util.find_spec(module) is None]
if missing:
    subprocess.check_call([sys.executable, "-m", "pip", "install", *missing])
import numpy as np
import matplotlib.pyplot as plt
SEED = 2026
random.seed(SEED)
np.random.seed(SEED)
rng = np.random.default_rng(SEED)
ARTIFACTS = Path("artifacts")
ARTIFACTS.mkdir(exist_ok=True)
print("Python:", platform.python_version(), "NumPy:", np.__version__, "Seed:", SEED)

Checkpoint. Run this notebook independently of H1. No saved H1 model is required.

2. Start with a query and an addressable memory

Q 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).

Q = np.array([[1.,0.]])
K = np.array([[1.,0.],[0.,1.],[1.,1.]])
V = np.array([[10.,0.],[0.,10.],[10.,10.]])
labels = ["memory A","memory B","memory C"]
print("Q:",Q,"\nK:\n",K,"\nV:\n",V)
plt.scatter(K[:,0],K[:,1],s=120)
for name, point in zip(labels,K): plt.annotate(name,point,xytext=(5,7),textcoords="offset points")
plt.arrow(0,0,Q[0,0],Q[0,1],width=.01,color="tab:red",length_includes_head=True)
plt.xlim(-.2,1.4); plt.ylim(-.2,1.4); plt.xlabel("Key dimension 1"); plt.ylabel("Key dimension 2"); plt.show()

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.

3. Calculate scaled scores

For 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.

scores = Q @ K.T / np.sqrt(K.shape[-1])
print("Raw dot products:",Q@K.T)
print("Scaled scores:",scores)
assert np.allclose(scores,[[1/np.sqrt(2),0,1/np.sqrt(2)]])

Checkpoint. The scaled score row is approximately [0.7071, 0, 0.7071]. Scores are not yet probabilities and may be negative.

4. Convert scores into a row of weights

Softmax 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.

def softmax(z):
    shifted = z-z.max(axis=-1,keepdims=True)
    exp = np.exp(shifted)
    return exp/exp.sum(axis=-1,keepdims=True)
weights = softmax(scores)
print("Weights:",weights,"Row sum:",weights.sum(axis=-1))
assert np.allclose(weights.sum(axis=-1),1)
plt.bar(labels,weights[0],color=["#1e8798","#7f2635","#f3b23a"])
plt.ylim(0,1); plt.ylabel("Attention weight"); plt.show()

Checkpoint. The approximate weights are [0.4011, 0.1978, 0.4011]. B still has positive weight even though its raw score is zero.

5. Retrieve a weighted combination of values

The output is weights @ V. Each output dimension combines the same weights with the corresponding value coordinates. There is no argmax selection: multiple values contribute.

context = weights @ V
print("Weighted contributions:\n",weights[0,:,None]*V)
print("Context vector:",context)
assert context.shape == (1,2)
assert np.allclose(context,(weights[0,:,None]*V).sum(axis=0,keepdims=True))

Checkpoint. The context is approximately [8.0222, 5.9889]. Explain which memory items contribute to each coordinate before moving on.

6. Change only the query

The 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.

queries = np.array([[1.,0.],[0.,1.],[-1.,0.]])
temperature = 1.0
assert temperature > 0
W = softmax((queries@K.T/np.sqrt(2))/temperature)
print("Contexts:\n",W@V)
fig,ax=plt.subplots(figsize=(6,3))
im=ax.imshow(W,cmap="YlGnBu",vmin=0,vmax=1)
ax.set_xticks(range(3),labels); ax.set_yticks(range(3),["q=[1,0]","q=[0,1]","q=[-1,0]"])
for row in range(3):
    for col in range(3): ax.text(col,row,f"{W[row,col]:.2f}",ha="center",va="center",color="black")
fig.colorbar(im,ax=ax); plt.show()

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.

7. Add causal masking before softmax

For 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.

X_demo = np.array([[1.,0.],[0.,1.],[1.,1.]])
self_scores = X_demo@X_demo.T/np.sqrt(2)
allowed = np.tril(np.ones((3,3),dtype=bool))
masked_scores = np.where(allowed,self_scores,-np.inf)
causal_weights = softmax(masked_scores)
print("Allowed (query rows, key columns):\n",allowed.astype(int))
print("Causal attention:\n",causal_weights)
assert np.all(causal_weights[~allowed] == 0)
assert np.allclose(causal_weights.sum(axis=-1),1)
plt.imshow(causal_weights,cmap="YlGnBu",vmin=0,vmax=1)
plt.xlabel("Key/source position"); plt.ylabel("Query/destination position"); plt.colorbar(); plt.show()

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.

8. Separate padding from future masking

A 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.

is_real_key = np.array([True,True,False])
combined_allowed = allowed & is_real_key[None,:]
assert combined_allowed.any(axis=-1).all()
combined_weights = softmax(np.where(combined_allowed,self_scores,-np.inf))
print(combined_weights)
assert np.all(combined_weights[:,2] == 0)

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.

9. Verify the same calculation with PyTorch

The 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.

if importlib.util.find_spec("torch") is None:
    subprocess.check_call([sys.executable, "-m", "pip", "install", "torch"])
import torch
from torch import nn
import copy, time
torch.set_num_threads(1)
torch.manual_seed(SEED)
torch.use_deterministic_algorithms(True)
device = torch.device("cpu")
print("PyTorch:", torch.__version__, "Device:", device)

import torch.nn.functional as F
q=torch.tensor(Q,dtype=torch.float32)[None,None,:,:]
k=torch.tensor(K,dtype=torch.float32)[None,None,:,:]
v=torch.tensor(V,dtype=torch.float32)[None,None,:,:]
torch_context = F.scaled_dot_product_attention(q,k,v,dropout_p=0.0)
assert np.allclose(torch_context.numpy()[0,0],context,atol=1e-5)
print("NumPy and PyTorch agree:",torch_context[0,0])

Checkpoint. The parity assertion must pass. This is an implementation check for the chosen example, not evidence that attention has learned language.

10. Test what an attention plot can explain

Weights 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.

V_changed=V.copy(); V_changed[1]=[100.,100.]
changed_context=weights@V_changed
removed_weights=softmax(np.where(np.array([[True,False,True]]),scores,-np.inf))
print("Original weights:",weights)
print("Same weights, changed values:",changed_context)
print("Remove B and renormalize:",removed_weights@V)
assert not np.allclose(changed_context,context)
evidence={"Q":Q.tolist(),"K":K.tolist(),"V":V.tolist(),"scores":scores.tolist(),
          "weights":weights.tolist(),"context":context.tolist(),
          "causal_weights":causal_weights.tolist(),"seed":SEED,
          "interpretation":"Attention weights are routing weights, not proof of causal importance."}
(ARTIFACTS/"w07-h2-attention.json").write_text(json.dumps(evidence,indent=2))

Checkpoint. State exactly what the heatmap supports: where the chosen computation routes information. Do not claim it proves why a real-world outcome occurred.

11. Independent practice and solutions

1. Set temperature to 0.5 and 2.0 in Step 6. Predict and then compare the sharpness, keeping Q, K, V fixed.

2. Replace all keys by the same vector. Expected: every available key gets equal weight, regardless of query; masked keys still get zero.

3. Make one value very large. Expected: the context may change greatly with no change to attention weights.

4. 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.

5. 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.

Stop when: your evidence JSON contains scores, weights, context, and a causal-mask assertion; explain all matrix axes to a partner.

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.

Sources

- Bahdanau et al. (2014/2015), joint alignment and translation

- Vaswani et al. (2017), Attention Is All You Need

- PyTorch scaled dot-product attention API

- 3Blue1Brown attention lesson