Step 4 — A Causal Self-Attention Head

Released with Lecture 4 · ~3 hours · Starts from: Step 3 (or solutions/step-03.ipynb).

In what follows, there are blocks of code provided for you. It’s important that you can read the code. If you are not proficient with Python, then I propose the following workflow. For each block of code, keywords are provided for the python syntax which is being used, so you can look up python syntax and functions you don’t recognise. To help, beneath many code blocks you will find Python background foldouts: self-contained tutorials, with live Python boxes you can edit and run right on the page, for syntax that is new in this step (syntax already introduced is covered in earlier steps’ foldouts). Then I suggest that you first experiment with changing the code or printing out intermediate variables, until you have figured out what you think it does. Then point AI toward the page and step number, provide your precise explanation of what you believe each line of the code is doing, and ask for corrections. In this way you will quickly become a code reader (if not a code writer). — Kate

Goal

Implement one head of causal self-attention — the mathematical object at the center of the course — as a standalone, thoroughly tested module. No training this week; the deliverables are correctness proofs executed as assertions and pictures of attention matrices on real text. (In Step 5 this head gets multiplied, wrapped, and stacked into the full GPT.)

Notation follows Lecture 4 and Phuong–Hutter.

On paper first

1. Write out, with all dimensions labeled, the map

   $$X \mapsto \operatorname{softmax}\!\Bigl(\operatorname{mask}\bigl(\tfrac{XW_Q (XW_K)^\top}{\sqrt{d_k}}\bigr)\Bigr)\, X W_V$$

   for $X \in \mathbb{R}^{T\times d}$, $W_Q, W_K \in \mathbb{R}^{d\times d_k}$, $W_V \in \mathbb{R}^{d\times d_v}$. Which pairs of tokens does entry $A_{ij}$ of the attention matrix couple, and what does row $i$ of the output equal, as an expectation?

2. Prove the two properties you will turn into unit tests:

   • Rows of $A$ are in the simplex (nonnegative, sum to 1), so output row $i$ is a convex combination of value vectors — attention cannot leave the convex hull of $\{v_j\}$.
   • Causality: with the mask sending $j > i$ scores to $-\infty$, output row $i$ is a function of $x_1,\dots,x_i$ only. State this as: $\partial(\text{out}_i)/\partial x_j = 0$ for $j > i$.

3. Recompute the variance lemma from lecture: for $q, k$ with iid mean-0, variance-1 entries, $\operatorname{Var}(q\cdot k) = d_k$ — hence the $\sqrt{d_k}$.

Tasks

4.1 The head

import torch, torch.nn as nn        # import with alias (as)
import torch.nn.functional as F

class Head(nn.Module):              # class definition with inheritance (subclass of nn.Module)
    def __init__(self, d, d_head, T_max):
        super().__init__()          # super(): call the parent class's constructor
        self.key   = nn.Linear(d, d_head, bias=False)
        self.query = nn.Linear(d, d_head, bias=False)
        self.value = nn.Linear(d, d_head, bias=False)
        # causal mask, precomputed; `register_buffer` = tensor state that isn't a parameter
        self.register_buffer('tril', torch.tril(torch.ones(T_max, T_max)))

    def forward(self, x):                       # x: (B, T, d)
        B, T, d = x.shape                       # tuple unpacking
        q, k, v = self.query(x), self.key(x), self.value(x)     # (B, T, d_head)
        scores = q @ k.transpose(-2, -1) / k.shape[-1]**0.5     # @ = matrix multiply; ** = power; -1 = last index
        scores = scores.masked_fill(self.tril[:T, :T] == 0, float('-inf'))  # boolean mask (==); float('-inf')
        A = F.softmax(scores, dim=-1)                           # (B, T, T)
        self.A = A.detach()                                     # stash for visualization
        return A @ v                                            # (B, T, d_head)

import math                    # use it with a dot
print(math.sqrt(16))
import math as m               # the same module, shorter name
print(m.pi)
from math import sqrt, pi      # or import specific names directly
print(sqrt(2) * pi)

class Animal:
    def __init__(self, name):
        self.name = name
    def speak(self):
        return self.name + " makes a sound"

class Dog(Animal):                 # Dog inherits from Animal
    def __init__(self, name, trick):
        super().__init__(name)     # run the parent's constructor first
        self.trick = trick
    def speak(self):               # override the inherited method
        return self.name + " says woof"

d = Dog("Rex", "sit")
print(d.speak())
print(d.name, d.trick)             # name came from Animal's constructor
print(Animal("Cat").speak())

lin = nn.Linear(32, 16, bias=False)  # a learnable 16x32 matrix
print(sum(p.numel() for p in lin.parameters()))   # 512 trainable numbers
emb = nn.Embedding(65, 32)           # 65 learned vectors of length 32
head = Head(32, 16, T_max=8)
print(len(list(head.parameters())))  # 3: the query, key, value matrices
print(head.tril.requires_grad)       # False — a buffer, not a parameter

import numpy as np
print(2 ** 10)
print(16 ** 0.5)             # fractional power: a square root
A = np.array([[1., 2.],
              [3., 4.]])
B = np.array([[1., 0.],
              [0., 2.]])
print(A @ B)                 # matrix product: rows of A times columns of B
print(A * B)                 # elementwise — a different animal entirely

import numpy as np
x = np.array([3., -1., 2., -7.])
print(x == 2.)                 # elementwise comparison: a boolean array
print(x > 0)
y = x.copy()
y[x > 0] = 0.                  # assign wherever the mask is True
print(y)
z = x.copy()
z[x == 2.] = float('-inf')     # the masking idiom from the head
print(z)
print(np.exp(float('-inf')))   # e^(-inf) = 0: why -inf, not some big number

Batched matrix multiplication (@ over the last two dims of a 3-tensor) is doing all the work; write out what q @ k.transpose(-2,-1) computes entrywise before running anything.

4.2 Tests — your proofs, executed

With random x = torch.randn(4, 8, 32) and a Head(32, 16, T_max=8):

1. Simplex: A nonnegative; A.sum(dim=-1) all ones; upper triangle of every A[b] exactly zero.
2. Causality by perturbation: run the head, change x[:, 5, :] arbitrarily, rerun; assert outputs at positions 0–4 are bitwise unchanged, positions 5–7 changed.
3. Causality by gradient: out[0, 3].sum().backward(); x.grad[0, j] is zero for all $j > 3$ and (generically) nonzero for $j \le 3$. This is test 2’s derivative form — and a free lesson in what .backward() on a non-scalar slice means.
4. Scale check: print scores.std() with and without the $\sqrt{d_k}$, at d_head ∈ {16, 64, 256}. Watch softmax saturate to one-hot rows without it, exactly per the variance lemma.

4.3 Attention on real text — the pictures

Attention is only interesting when $X$ carries content and position, so build the minimal context: token embedding plus learned positional embedding (Lecture 4’s equivariance theorem is why the positional term must be there — without it, this week’s head literally cannot know token order):

tok_emb = nn.Embedding(V, 32)
pos_emb = nn.Embedding(64, 32)
idx = torch.tensor(encode(text[:64]))[None, :]          # indexing with None adds an axis; (1, 64) real Shakespeare
x = tok_emb(idx) + pos_emb(torch.arange(64))[None, :]
head = Head(32, 16, T_max=64)
out = head(x)
plt.imshow(head.A[0], cmap='Blues'); plt.colorbar()     # semicolon: two statements on one line

import numpy as np
v = np.array([1., 2., 3.])
print(v.shape)
print(v[None, :].shape)        # new leading axis: one "batch" of size 1
print(v[:, None].shape)        # new trailing axis: a column
a = 1; b = 2                   # two statements, one line
print(a + b)

Untrained weights already produce structure (why is the first column heavy? think about what row 1’s simplex constraint forces, and how the mask compounds it). Now the real experiment: train nothing, but plot A for 5 different random seeds, and describe the family of patterns — diagonal-hugging, column-collapse, near-uniform. In Step 6 you’ll plot these same pictures after training and watch actual algorithms (previous token heads, induction-like patterns) appear in this space.

4.4 (Bridge) Attention as smoothing

One experiment tying to the kernel view from lecture: replace q, k with positions only — scores[i,j] = -(i-j)**2 / (2*s**2), masked — and plot what the head computes for several bandwidths s. This is Nadaraya–Watson smoothing along the sequence, the $W_Q{=}W_K{=}0$ degenerate case of attention; content-dependence is the only thing attention adds. (5 lines; do plot it.)

Checkpoints

• All four tests of 4.2 pass; you can state which lemma of Lecture 4 each one is. Row sums equal 1 to $\sim 10^{-7}$ (float32); positions before a perturbation are bitwise identical, not merely close.

• The scale check (4.2 part 4) should reproduce roughly this table — raw score standard deviation tracking $\sqrt{d_k}$, and the largest softmax weight in a row saturating to 1 without the scaling:

  $d_k$	raw std	$\sqrt{d_k}$	scaled std	max weight, raw	max weight, scaled
  16	4.12	4.0	1.03	1.000	0.798
  64	7.95	8.0	0.99	1.000	0.684
  256	18.00	16.0	1.13	1.000	0.467

  (Your numbers will differ in the last digits by seed; the pattern is the point.) Note the unscaled row is already fully saturated at $d_k{=}16$ — the softmax is one-hot to three decimals, so its Hessian $\operatorname{diag}(p)-pp^\top$ is numerically zero and no gradient flows through those weights.

• Masking after softmax instead of before leaves row sums ranging over ≈0.13–1.00 rather than exactly 1 — the operator is no longer row-stochastic and Corollary 3.3 fails.

• Heatmaps for 4.3: strictly lower-triangular, rows summing to 1, visibly different across seeds.

• You can answer, without notes: what breaks if the mask is applied after softmax instead of before? (Try it: which test fails?)

If you’re stuck

• A has nan rows: you masked with 0 instead of -inf before softmax, or masked after softmax and renormalized a zero row.
• Tests 2–3 fail at position $i$ itself: off-by-one in the mask — tril includes the diagonal; a token may attend to itself.
• Shape errors: annotate every line with its shape as a comment. This is not a beginner crutch; nanoGPT does it too.

Going further (optional)

• Prove and test permutation equivariance — and mind the trap. With the mask removed, head(x[:, perm, :]) == head(x)[:, perm, :] for a random permutation of the stream rows (Lecture 4’s Theorem 6.1 as an assertion). This stays true even after you add positional embeddings to x, which surprises everyone: the theorem is about attention as a function of $X$, and adding anything to $X$ doesn’t change that.

  To see position actually break the symmetry you must permute the tokens, not the rows:

  a = head(tok(idx[:, perm]))            # permute tokens, then embed
  b = head(tok(idx))[:, perm, :]         # embed, then permute output
  # equal (equivariant); now redo both with  + pos(arange(T))  added:
  a = head(tok(idx[:, perm]) + P)
  b = head(tok(idx) + P)[:, perm, :]     # NOT equal — symmetry broken

  Verify all four, and write one sentence locating exactly which map lost equivariance. (Measured: the first pair agree to $10^{-7}$; the second differ by $O(1)$.)

• Implement the online/streaming softmax: process keys one at a time keeping running max and running sum, and check it matches — this identity (softmax’s shift-invariance again) is the heart of FlashAttention.

• Cost accounting: time the forward pass at $T = 256, 512, 1024, 2048$ (fixed $d$) and fit the exponent. Compare against the $O(T^2 d)$ count from lecture.

Catch-up

Run solutions/step-03.ipynb for the data pipeline and encode; this step needs only Lecture 4 and reuses no model code from Step 3 — a good re-entry point if you’ve fallen behind.