Step 4 — A Causal Self-Attention Head¶

Solution notebook. No training this week. The deliverables are Lecture 4's propositions executed as assertions, plus pictures of attention matrices on real text. Runs in seconds.

Every assertion below is annotated with the result it verifies.

4.1 The head¶

$$\operatorname{Attn}(X)=\operatorname{softmax}\!\Bigl(\operatorname{mask}\bigl(\tfrac{XW_Q(XW_K)^\top}{\sqrt{d_k}}\bigr)\Bigr)XW_V$$

register_buffer holds the causal mask as non-parameter state, so it moves with .to(device) but is not trained. We stash A for plotting.

In [ ]:
import torch, torch.nn as nn, torch.nn.functional as F
import matplotlib.pyplot as plt
torch.manual_seed(0)

class Head(nn.Module):
    def __init__(self, d, d_head, T_max, causal=True):
        super().__init__()
        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)
        self.causal = causal
        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
        q, k, v = self.query(x), self.key(x), self.value(x)
        scores = q @ k.transpose(-2, -1) / k.shape[-1]**0.5    # (B,T,T)
        if self.causal:
            scores = scores.masked_fill(self.tril[:T, :T] == 0, float('-inf'))
        A = F.softmax(scores, dim=-1)
        self.A = A.detach()
        return A @ v                                            # (B,T,d_head)

d, dh, T, B = 32, 16, 8, 4
head = Head(d, dh, T)
x = torch.randn(B, T, d)
out = head(x)
print('out', out.shape, '| A', head.A.shape)

4.2 The proofs, as assertions¶

In [ ]:
A = head.A

# Prop 3.1 — rows of A lie in the simplex
assert (A >= 0).all()
assert torch.allclose(A.sum(-1), torch.ones(B, T), atol=1e-6)
print(f'Prop 3.1  simplex rows        OK  (max |rowsum-1| = '
      f'{(A.sum(-1)-1).abs().max():.1e})')

# Cor 3.2 — output lies in the convex hull of the value vectors
v = head.value(x)
lo, hi = v.min(dim=1).values, v.max(dim=1).values
assert ((out >= lo[:, None, :] - 1e-6) & (out <= hi[:, None, :] + 1e-6)).all()
print('Cor 3.2   output in conv hull  OK')

# Cor 3.3 — every allowed weight is strictly positive
allowed = torch.tril(torch.ones(T, T)).bool()
assert (A[0][allowed] > 0).all()
print(f'Cor 3.3   allowed weights > 0  OK  (min = {A[0][allowed].min():.2e})')

# Prop 5.1 — the mask makes A lower-triangular
assert (A.triu(diagonal=1) == 0).all()
print('Prop 5.1  A lower-triangular   OK')

Causality, two ways¶

By perturbation (positions before the change must be bitwise identical, not merely close) and by gradient — the derivative form of the same statement.

In [ ]:
x2 = x.clone(); x2[:, 5, :] = torch.randn(B, d)
out2 = head(x2)
assert torch.equal(out[:, :5, :], out2[:, :5, :])
assert not torch.allclose(out[:, 5:, :], out2[:, 5:, :])
print('Prop 5.1  perturbation test    OK  (positions 0-4 bitwise identical)')

xg = x.clone().requires_grad_(True)
head(xg)[0, 3].sum().backward()
assert (xg.grad[0, 4:] == 0).all()
assert (xg.grad[0, :4].abs().sum(-1) > 0).all()
print('Prop 5.1  gradient test        OK  (d out_3 / d x_j = 0 for j > 3)')

Lemma 4.5 — why the $\sqrt{d_k}$¶

Raw scores have standard deviation $\sqrt{d_k}$, so without the scaling the softmax saturates to one-hot as $d_k$ grows — its Hessian $\operatorname{diag}(p)-pp^\top$ goes numerically to zero and no gradient flows through the attention weights.

In [ ]:
print(f"{'d_k':>5} {'raw std':>9} {'sqrt(d_k)':>10} {'scaled std':>11}"
      f" {'max wt raw':>11} {'max wt scaled':>14}")
for dk in (16, 64, 256):
    q, k = torch.randn(1, T, dk), torch.randn(1, T, dk)
    raw = q @ k.transpose(-2, -1)
    sc = raw / dk**0.5
    print(f'{dk:>5} {raw.std():>9.2f} {dk**0.5:>10.1f} {sc.std():>11.2f}'
          f' {F.softmax(raw,-1).max():>11.3f} {F.softmax(sc,-1).max():>14.3f}')

Reference: raw std 4.12 / 7.95 / 18.00 against $\sqrt{d_k}$ = 4 / 8 / 16, and the unscaled max weight is 1.000 already at $d_k=16$ — fully saturated to three decimals.

Exercise 4 — masking before vs after softmax¶

Zeroing $A$ after the softmax (instead of $-\infty$ before) destroys row-stochasticity, and with it Corollary 3.3's guarantee.

In [ ]:
s = torch.randn(1, T, T)
right = F.softmax(s.masked_fill(torch.tril(torch.ones(T,T))==0, float('-inf')), -1)
wrong = F.softmax(s, -1) * torch.tril(torch.ones(T, T))
print(f'-inf before softmax : row sums {right.sum(-1)[0].min():.3f} '
      f'.. {right.sum(-1)[0].max():.3f}')
print(f'zeroing after       : row sums {wrong.sum(-1)[0].min():.3f} '
      f'.. {wrong.sum(-1)[0].max():.3f}   <- not stochastic')

4.3 Equivariance — and the trap¶

Theorem 6.1 is about $\operatorname{Attn}$ as a function of $X$, and it stays true however $X$ was built — including after adding positional embeddings. Permuting the rows of $X$ therefore will not show you the symmetry breaking, and most people set this experiment up that way the first time.

Equivariance is lost one step earlier, in the embedding map $\iota(x)_t = E_{x_t} + p_t$, because the positional term does not travel with the permuted token. You must permute tokens.

In [ ]:
hn = Head(d, dh, T, causal=False)      # no mask, so S_T acts
perm = torch.randperm(T)

# (a) the layer is equivariant in X ...
assert torch.allclose(hn(x[:, perm, :]), hn(x)[:, perm, :], atol=1e-5)
print('Thm 6.1   equivariant in X                     OK')

# (b) ... and STAYS equivariant when X contains positional vectors
xp = x + nn.Embedding(T, d)(torch.arange(T))[None]
assert torch.allclose(hn(xp[:, perm, :]), hn(xp)[:, perm, :], atol=1e-5)
print('          still equivariant with pos-emb in X   OK  <- the trap')

# (c) the composite TOKEN -> output: equivariant without positions ...
tok, pos = nn.Embedding(V := 65, d), nn.Embedding(T, d)
idx = torch.randint(0, V, (1, T))
P = pos(torch.arange(T))[None]
a1, b1 = hn(tok(idx[:, perm])), hn(tok(idx))[:, perm, :]
assert torch.allclose(a1, b1, atol=1e-5)
print(f'Thm 6.1   token-perm, NO pos-emb               OK  '
      f'(diff {(a1-b1).abs().max():.1e})')

# (d) ... and BROKEN with them. This is Corollary 6.2.
a2, b2 = hn(tok(idx[:, perm]) + P), hn(tok(idx) + P)[:, perm, :]
assert not torch.allclose(a2, b2, atol=1e-3)
print(f'Cor 6.2   token-perm, WITH pos-emb        BROKEN  '
      f'(diff {(a2-b2).abs().max():.1e})')

4.4 Attention patterns on real text¶

Untrained weights, five seeds. The structure you see is a prior: row 1 has only one position available so $A_{11}=1$ by the simplex constraint, row 2 splits between two, and so on — the mask alone forces early positions to carry weight. In Step 6 you will plot these same pictures after training and find actual algorithms in them.

Setup: data and tokenizer (Step 1's solution, reproduced)¶

Every notebook in this series is self-contained: it re-creates what it needs from earlier steps in one compact cell, so you can run it top to bottom without opening the others. On Colab, uncomment the download.

In [ ]:
# !wget -q https://raw.githubusercontent.com/karpathy/char-rnn/master/data/tinyshakespeare/input.txt
import torch

with open('input.txt') as f:
    text = f.read()
vocab = sorted(set(text))
V = len(vocab)
stoi = {ch: i for i, ch in enumerate(vocab)}
itos = {i: ch for i, ch in enumerate(vocab)}
encode = lambda s: [stoi[c] for c in s]
decode = lambda ids: ''.join(itos[i] for i in ids)
ids = torch.tensor(encode(text), dtype=torch.long)

assert len(text) == 1_115_394 and V == 65
print(f'{len(text):,} characters, vocab {V}')
In [ ]:
Tc = 48
tok_emb = nn.Embedding(V, 32)
pos_emb = nn.Embedding(Tc, 32)
idx = ids[:Tc][None, :]
chars = [decode([i]) for i in idx[0].tolist()]
labels = [{'\n': '\\n', ' ': '␣'}.get(c, c) for c in chars]

fig, axes = plt.subplots(1, 5, figsize=(16, 3.4))
for s, ax in enumerate(axes):
    torch.manual_seed(s)
    h = Head(32, 16, Tc)
    xin = tok_emb(idx) + pos_emb(torch.arange(Tc))[None]
    h(xin)
    ax.imshow(h.A[0], cmap='Blues')
    ax.set_title(f'seed {s}', fontsize=10); ax.set_xticks([]); ax.set_yticks([])
plt.suptitle('untrained causal attention, five seeds '
             '(strictly lower-triangular, rows sum to 1)')
plt.tight_layout()

4.5 Attention as kernel smoothing¶

Set $W_Q = W_K = 0$ and attention becomes a running mean; hardwire scores by position, $S_{ij} = -(i-j)^2/2s^2$, and it becomes a Gaussian smoother of bandwidth $s$ — i.e. a convolution. This is Nadaraya–Watson along the sequence. Content-dependence is the only thing real attention adds on top, which is exactly why it is powerful: the 'bandwidth' and 'location' can vary per token and per input.

In [ ]:
i = torch.arange(Tc).float()
mask = torch.tril(torch.ones(Tc, Tc))
fig, axes = plt.subplots(1, 4, figsize=(14, 3.4))
for ax, s in zip(axes, [0.5, 2.0, 8.0, float('inf')]):
    if s == float('inf'):
        S = torch.zeros(Tc, Tc); title = 'W_Q=W_K=0 (running mean)'
    else:
        S = -(i[:, None] - i[None, :])**2 / (2 * s**2); title = f'bandwidth s={s}'
    Ak = F.softmax(S.masked_fill(mask == 0, float('-inf')), dim=-1)
    ax.imshow(Ak, cmap='Blues'); ax.set_title(title, fontsize=10)
    ax.set_xticks([]); ax.set_yticks([])
plt.suptitle('positional-kernel attention: the degenerate, content-blind case')
plt.tight_layout()

Exercise 7 — rank of the score matrix¶

$S = QK^\top/\sqrt{d_k}$ factors through $\mathbb R^{d_k}$, so $\operatorname{rank} S \le d_k$ regardless of $T$. A single head can only produce a low-rank family of patterns — which is the reason Lecture 5 runs several in parallel.

In [ ]:
big = Head(64, 8, 32, causal=False)
xb = torch.randn(1, 32, 64)
big(xb)
S = (big.query(xb) @ big.key(xb).transpose(-2, -1))[0]
print(f'score matrix is {tuple(S.shape)} with d_k = 8; '
      f'rank = {torch.linalg.matrix_rank(S).item()}')
assert torch.linalg.matrix_rank(S).item() <= 8

→ Continue with Step 5: multiply this head, wrap it in residuals and LayerNorm, and count every parameter.