05 Notebook: Self-Attention Properties

April 2026 (239 Words, 2 Minutes)

Goal

Take the attention function you wrote in notebook 04 and put it under a microscope. Four properties to investigate experimentally:

• Permutation equivariance, as a numerical assertion on your own code
• Causal masking and the lower-triangular attention pattern
• Causal prefix invariance — the property that makes KV-caching correct
• Positional embeddings — what attention literally cannot do without them, plus a side-by-side look at sinusoidal / learned-absolute / RoPE and the extrapolation cliff that motivates Lecture 6

Open the Notebook

Open in JupyterLite — runs directly in your browser, no installation required.

Pure NumPy and matplotlib only. No PyTorch, no transformers. You will reuse your attention function from notebook 04.

What You Will Do

1. Implement permutation_matrix — build a (T, T) permutation matrix from a list of indices
2. Implement check_equivariance_max_defect — verify ||Attn(PX) − P · Attn(X)||_F < 1e-10 across many random P
3. Adding a causal mask — three lines to turn self-attention into causal attention; side-by-side heatmaps
4. Implement prefix_invariance_max_defect — verify Y_full[:T-1] == Y_prefix bitwise
5. Causal violates strong permutation equivariance — measure how badly when P swaps a future token into the past
6. The same-multiset eye-opener — “the cat ate the mouse” vs “the mouse ate the cat” produce identical outputs without position
7. Implement my_sinusoidal_position — write the Vaswani formula yourself
8. Three positional schemes — sinusoidal, learned-absolute, RoPE side by side on the same input
9. Implement rope_angle_grid — isolate the (T, d/2) angle table at the heart of every RoPE implementation
10. Positional dilution at large T — cos(P[0], P[t]) for t from 1 to 10 000
11. Extrapolation failure — what each scheme does when asked for a position past training length

Deliverables (AI Diary entry — notebook §11)

• Why is the equivariance defect at floating-point noise level (≈ 1e-15)? What makes the property exact, not approximate?
• How does causal prefix invariance enable KV-caching? What gets cached, what gets recomputed, and what would break if causal attention were replaced with self-attention?
• Why does RoPE rotate Q and K but not V?
• The architectural change from GPT-2’s learned-absolute (n_max = 1024) to LLaMA’s RoPE (no hard cap). Connect to what you saw in section 11.

Reference

• Lecture 5 — Inside the Transformer
• Chapter 8 — Causal Attention
• Chapter 10 — Positional Encoding
• §10.2 — Permutation equivariance lemma