Step 4 — A Causal Self-Attention Head

Released with Lecture 4 · Starts from: Step 3 (or solutions/step-03.ipynb) · Solution: notebook · read online

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 calculation from lecture (Remark 2.6): for $q, k$ with iid mean-0, variance-1 entries, $\operatorname{Var}(q\cdot k) = d_k$ — hence the $\sqrt{d_k}$.

From proofs to executable assertions

This step isolates one causal head so that its structural properties can be tested before training a full transformer. This is a useful scientific order: first establish what the implementation must satisfy for every parameter value, then examine the patterns produced by particular parameters.

A minibatch of size $B$ adds a leading coordinate to every shape:

Matrix multiplication acts on the final two coordinates and preserves the batch coordinate. Before running the implementation, writing these shapes beside every line catches most transposition mistakes.

Each main test below is a theorem of Lecture 4 specialized to finite arrays.

Some equalities are structural and some are numerical. A stored causal mask can make the upper triangle exactly zero. Floating-point row sums are usually only close to one. A perturbation test should compare outputs that execute the same deterministic arithmetic on the same earlier data; gradient tests should allow for the numerical tolerance appropriate to the chosen data type.

A good unit test is a small proof with numbers substituted. Proposition 4.3 says that a future variable is absent from an earlier output formula. The perturbation test checks the resulting function equality; the gradient test checks its differential consequence. Agreement between the two makes an indexing accident much less likely.

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.

A note on the line self.A = A.detach(): storing an attention matrix for plotting should not keep an entire backpropagation graph alive, so the head stashes a gradient-detached copy for visualization. Detaching the diagnostic copy does not alter the returned attention output; detaching the matrix used in the returned output would incorrectly cut the training gradient.

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

This experiment supplies a controlled comparison for the pictures of 4.3. Its bandwidth is known, so a narrow diagonal or a broad running average has a mathematical explanation (Lecture 4, Example 8.4, and the heatmaps in its §8.2). Trained heatmaps can later be compared with these baselines.

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

What this head changes from Step 3

The Step 3 neural $n$-gram compresses a fixed concatenated context with one MLP. Its first-layer parameter count grows with the chosen context length. The Step 4 head uses shared projections at every position, and its parameter count is independent of the current $T$. Its attention matrix then selects a different convex combination for every input and query position.

This does not make one head a language model. It has no multi-head output projection, residual connection, normalization, position-wise MLP, or unembedding to vocabulary logits. Lecture 5 assembles those pieces. Step 4’s goal is narrower and foundational: make the data-dependent communication operator correct before multiplying and stacking it.

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.