Six notes on attention

Transcript: this part, with timestamps

This is the part of the lecture that turns a working block of PyTorch into a mental model. In P4 you watched wei go from a constant lower-triangular average to a data-dependent matrix of query·key affinities, and you saw values get aggregated instead of raw embeddings. That's the mechanism. But a mechanism you can only run is not a mechanism you can modify, and the next thirty minutes of the video (P6) will modify it hard — four heads in parallel, a feed-forward on top, residuals around both. So Karpathy pauses and draws the boundary between the operation and its costume: six short notes, each isolating one thing you could change without attention noticing.

Outline, with timestamps

• 71:38 — Note 1, attention is communication: nodes in a directed graph, each aggregating a weighted sum from whoever points at it.
• 72:06 — The graph we have: eight nodes, node 1 sees itself, node 8 sees everyone — a triangular DAG, one of infinitely many attention could run on.
• 72:46 — Note 2, no notion of space: attention acts over a set, unlike convolution, which is why position had to be added by hand in P3.
• 73:40 — Note 3, no communication across the batch dimension: four independent pools of eight nodes, thirty-two nodes total, zero edges between pools.
• 74:14 — Note 4, encoder vs. decoder blocks: delete one line and every node talks to every node; keep it and the future stays hidden.
• 75:39 — Note 5, self- vs. cross-attention: it's "self" only because Q, K and V are all read off the same x.
• 76:56 — Note 6, "scaled" attention: unit-variance q and k give wei a variance on the order of head_size.
• 77:43 — Why the softmax cares: diffuse logits in, diffuse distribution out.
• 78:13 — The sharpening demo: multiply the logits by eight and softmax collapses toward a one-hot vector.

Note 1 — attention is message passing, and the graph is an input

The framing that unlocks the rest: forget sequences for a moment and picture a directed graph. Every node holds a vector. Every node also advertises two things — a query, "here is what I'm looking for", and a key, "here is what I contain" — and the strength of the edge from node j to node i is the dot product of i's query with j's key. Normalise those strengths over i's in-neighbours, and each node updates itself with the weighted sum of the values its neighbours are willing to broadcast.

Attention is a communication mechanism.Karpathy, 71:38

Nothing in that description mentions text, order, or time. The graph is a parameter of the problem, not of the model: attention is defined for any set of nodes and any pattern of edges you care to allow. What the lecture's code happens to instantiate is one particular, very rigid graph — the autoregressive one. With block_size = 8 in the notebook, node 1 has exactly one in-edge (itself), node 2 has two, node 8 has eight. That's the lower-triangular pattern, and it lives in the model only as a tril buffer that decides which entries of wei get set to -inf before the softmax. Swap that buffer for a different 0/1 matrix and you have a different graph, with the same weights, the same shapes, and the same training loop.

This is worth internalising because it is the fork in the road that produced the whole zoo. Bidirectional encoders (BERT-style) use the all-ones matrix. Sliding-window attention uses a band. Prefix-LM uses a matrix that is dense in the prompt region and triangular after it. None of those are new operations; they are new masks.

Notes 2 and 3 — the two things attention genuinely cannot do

Note 2 is the sharpest one, and it explains a piece of P3 that can otherwise feel arbitrary. A convolution has spatial structure welded into it: a 3×3 filter means "left neighbour, me, right neighbour", and moving a pixel changes the answer. Self-attention has no such thing. The softmax-weighted sum is permutation-equivariant over the time axis — shuffle the tokens and the outputs come back shuffled the same way, unchanged in content. The model is looking at an unordered bag of vectors.

Which is exactly why position_embedding_table exists. In P3 Karpathy added a second embedding table indexed by torch.arange(T) and summed it into the token embeddings, and at the time it looked like a small nicety. It isn't: without it the model literally cannot tell "abc" from "cba", because the only thing carrying order was thrown away the moment we started doing weighted sums over a set.

One honest wrinkle Karpathy skips. The causal mask is itself indexed by position, so a decoder-only stack is not fully position-blind: token i can attend to exactly i+1 predecessors, and that count is a weak positional signal a network can learn to exploit. Later work found that decoder-only Transformers train tolerably with no explicit position encoding at all for this reason. The note is still the right note — attention proper has no notion of space — but the mask leaks a little of it back in.

Note 3 is smaller and is really a shape-reading exercise. The batch dimension is inert. wei has shape (B, T, T), and the matmuls that produce and consume it are batched matmuls: PyTorch runs B independent (T, T) problems in parallel and never mixes them. In the notebook's directed-graph picture that means not one graph of eight nodes but four separate graphs of eight nodes each — thirty-two nodes, twenty-eight of which any given node is forbidden from ever seeing.

Note the asymmetry that makes note 3 non-trivial rather than obvious: tril is shared across the batch and constant across training, while wei is different for every one of the B sequences and changes every step. The mask says which edges may exist; the data says how much each one carries. And because nothing crosses the batch, any layer that does mix across batch — BatchNorm, most obviously — would be a genuinely different kind of object here. That's a quiet setup for the LayerNorm discussion in P6.

Note 4 — encoder and decoder blocks differ by one line

If you want every token to see every other token — you're classifying the sentiment of a whole sentence, say, and there is no "future" to protect — you delete the masked_fill. That's the entire difference. Karpathy's terminology follows the original paper: a block with the triangular mask is a decoder block (it decodes autoregressively, so it must not see ahead), and one without is an encoder block.

Attention supports arbitrary connectivity between nodes.Karpathy, 75:09

Why the mask is load-bearing here and not just a design flourish: recall from P1 that the targets y are the inputs x shifted left by one. Position i's label is token i+1. If position i may attend to position i+1, the answer is sitting in the input. Training loss will fall through the floor, validation loss will follow it — because validation is scored the same leaky way — and generation, which has no future to peek at, will produce noise. This is the single most common way people break the model while "just experimenting", and it fails silently and impressively.

Two things the "one line" framing glosses over, worth knowing before you go read a real encoder-decoder implementation. First, a full encoder-decoder Transformer differs by more than the mask: the decoder carries a third sublayer, cross-attention into the encoder's output, which the lecture's model never grows. Second, real encoders still use masks — just padding masks, zeroing out attention to filler positions in variable-length batches. Karpathy's dataset has fixed-length chunks, so the issue never arises here.

Note 5 — "self" is a statement about where Q, K and V come from

In Head.forward, all three projections are applied to the same argument: self.key(x), self.query(x), self.value(x). One tensor in, three views out, nodes talking to themselves as a group. That, and only that, is what makes it self-attention.

Break the tie and you get cross-attention: queries from one stream, keys and values from another. The rule that keeps it straight — queries come from the stream you are writing into; keys and values come from the stream you are reading from. A translation decoder produces its queries from the partial output it has generated so far and reads keys and values off the encoded source sentence. A vision-language model produces queries from text tokens and reads image patches. The mechanism is byte-for-byte the same; the two nn.Linear calls just take a different argument.

The shape consequence is the one people trip on. In self-attention wei is square, (B, T, T), because the query set and the key set are the same set. In cross-attention it is (B, T_q, T_kv) — a rectangle, one row per node you're updating, one column per node you're reading. A triangular mask is meaningless on a rectangle and is normally absent: there is no ordering relationship between "the fourth word I'm generating" and "the ninth patch of the image".

Note 6 — why * head_size**-0.5, in numbers

The last note is the one with actual arithmetic behind it, and it's the reason the paper says "Scaled Dot-Product Attention" rather than just "Dot-Product Attention". Take q and k to be roughly unit-Gaussian — zero mean, unit variance per channel — which is approximately true at initialisation with sensible init. Their dot product sums head_size independent products of unit-variance terms, so its variance is on the order of head_size: 16 in the notebook, 64 in the final gpt.py, 128 in GPT-3-scale models. Multiplying by head_size**-0.5 divides the variance by head_size and puts it back at 1, independently of how wide you made the head.

Why unit variance rather than any other constant: wei feeds a softmax, and softmax's sharpness is set by the scale of its inputs, not their pattern. Logits near zero come out near uniform; multiply the same logits by eight and softmax converges toward a one-hot vector, putting essentially all mass on the single largest entry. A head whose attention is one-hot at initialisation is a head that copies one neighbour instead of averaging its neighbourhood — it has thrown away the aggregation that makes attention useful before the first gradient step, and the gradients that reach the non-selected entries are vanishing.

Numbers, from sampling q, k ~ N(0,1) with T = 8 and applying the causal mask, then measuring how much probability mass the last row's softmax puts on its single largest entry (uniform over eight would be 0.125):

Read the two right-hand columns as the point of the whole note: with the scaling, the initial attention distribution is the same shape no matter how wide the head is. Without it, widening the head silently sharpens attention — a hyperparameter change quietly becoming an initialisation change, which is exactly the kind of coupling that makes deep networks refuse to train and gives you nothing to debug.

# the argument, reproduced in four lines
q = torch.randn(B, T, head_size)
k = torch.randn(B, T, head_size)
wei = q @ k.transpose(-2, -1)                      # wei.var() ~ head_size
wei = q @ k.transpose(-2, -1) * head_size**-0.5    # wei.var() ~ 1

The scaling is used just to control the variance at initialization.Karpathy, 78:44

Notice the hedge, and take it seriously: at initialization. Nothing constrains q and k to stay unit-Gaussian once training starts — the projections are learned, and a head is perfectly free to grow its weights and sharpen its own attention, which is often what you want a trained head to do. The 1/√d_k is not a running normaliser; it is a one-time choice of units that stops the model from starting in a degenerate place. Modern stacks that train at scales where even this isn't enough add an actual normaliser — QK-Norm, applying LayerNorm to q and k before the dot product — precisely because the initialisation-time argument doesn't cover step 50,000.

The code at the end of this part

These eight minutes change no code — they annotate code you already wrote in P4. Here is the class as it appears in the finished gpt.py, with each of the six notes pointing at its line:

class Head(nn.Module):
    """ one head of self-attention """

    def __init__(self, head_size):
        super().__init__()
        self.key = nn.Linear(n_embd, head_size, bias=False)
        self.query = nn.Linear(n_embd, head_size, bias=False)
        self.value = nn.Linear(n_embd, head_size, bias=False)
        self.register_buffer('tril', torch.tril(torch.ones(block_size, block_size)))

        self.dropout = nn.Dropout(dropout)

    def forward(self, x):
        # input of size (batch, time-step, channels)
        # output of size (batch, time-step, head size)
        B,T,C = x.shape
        k = self.key(x)   # (B,T,hs)
        q = self.query(x) # (B,T,hs)
        # compute attention scores ("affinities")
        wei = q @ k.transpose(-2,-1) * k.shape[-1]**-0.5 # (B, T, hs) @ (B, hs, T) -> (B, T, T)
        wei = wei.masked_fill(self.tril[:T, :T] == 0, float('-inf')) # (B, T, T)
        wei = F.softmax(wei, dim=-1) # (B, T, T)
        wei = self.dropout(wei)
        # perform the weighted aggregation of the values
        v = self.value(x) # (B,T,hs)
        out = wei @ v # (B, T, T) @ (B, T, hs) -> (B, T, hs)
        return out

Line by line, against the notes (gpt.py L64–L90):

• L69–71 — three nn.Linear(n_embd, head_size, bias=False). All three are applied to the same x in forward: that is note 5's definition of "self". Biases are off because a per-channel constant added to every node's key or query shifts all affinities in that row by the same amount, which softmax then cancels; it would only waste parameters.
• L72 — register_buffer('tril', …). This is note 1's graph and note 4's encoder/decoder switch, stored as state that moves with .to(device) and gets saved in the state dict but receives no gradient. It is allocated at full block_size and sliced to [:T, :T] at use, so shorter contexts (the first few steps of generate) work unchanged.
• L83 — the scaled dot product, note 6. k.shape[-1] is head_size, not C; in the video Karpathy first typed C and corrects it in the video's pinned errata at 1:20:05. With one head where head_size == n_embd the two coincide, which is why the bug is invisible until P6 splits the embedding across four heads.
• L84–85 — masked_fill then softmax(dim=-1). -inf is the identity element for this pair: exp(-inf) = 0, so masked entries contribute nothing and the surviving entries in each row still normalise to 1. dim=-1 matters — normalising over the wrong axis produces columns that sum to one and a model that trains to nonsense.
• L86 — dropout on the attention matrix. Not yet present at this point in the lecture; it arrives with the scale-up in P7. Randomly zeroing entries of wei is dropping edges of the graph during training.
• L88–89 — the aggregation. (B,T,T) @ (B,T,hs) -> (B,T,hs): the output is head_size wide, not C wide, which is what makes concatenating four heads back to C work in P6.

Where people get stuck

• "If attention has no notion of space, why does the mask work?" Because the mask is not part of attention — it is a position-indexed object we bolt on. Attention proper is permutation-equivariant over T; the mask and the position embedding are the two places order enters the model. If you shuffle your tokens and shuffle tril identically, you get the same answers back, shuffled. The consequence people miss: the mask leaks a small amount of position information (how many predecessors a token has) even with the position embedding removed.
• Deleting the mask "to try an encoder" and watching the loss collapse to near zero. This looks like a triumph and is a leak: targets are inputs shifted by one, so an unmasked position reads its own label. Validation loss drops too, since validation is scored the same way, so the usual overfitting alarm doesn't fire. The tell is generation — sampling has no future to read, so text quality falls off a cliff while the numbers look superb. Encoder blocks are correct for classifying a whole sequence; they are not correct for next-token prediction.
• Normalising by C instead of head_size. Karpathy makes this exact slip on camera (corrected in the video description at 1:20:05) and the final gpt.py has it right as k.shape[-1]**-0.5. It is invisible in the single-head model where the two are equal and becomes a real 2× error the moment you split 384 channels across 6 heads of 64. Rule: the scale is always the dimension you summed over.
• Which side gives Q and which gives K/V in cross-attention. Queries come from the nodes being updated; keys and values from the nodes being read. Check yourself on shapes: wei must be (B, T_q, T_kv) and the output must have T_q rows, matching the stream you are writing into. If your output length changed, you have them backwards.

Go deeper, verified

• Attention Is All You Need — Vaswani et al. (2017) · §3.2.1 is note 6 in its original form, including the footnote with the same unit-variance argument; §3.2.3 enumerates the three places attention is used, which is note 5's encoder-decoder cross-attention.
• ng-video-lecture · gpt.py, class Head — Karpathy (2023) · the 27 lines all six notes are about; the finished version, so it also carries the dropout that arrives in P7.
• nanoGPT · CausalSelfAttention.forward — Karpathy (2023) · the same operation written for real: heads folded into the batch dimension, the mask registered as bias, and a fast path that hands the whole thing to F.scaled_dot_product_attention(…, is_causal=True). Line 67 is L83 above with math.sqrt spelled out.
• torch.nn.functional.scaled_dot_product_attention — PyTorch docs · note 4 and note 6 as API surface: is_causal is the triangular mask, scale defaults to 1/√E, and passing separate key/value tensors is note 5.
• BERT: Pre-training of Deep Bidirectional Transformers — Devlin et al. (2018) · field map extra. What note 4's encoder branch turns into: no causal mask, so the training objective has to change too (masked-token prediction instead of next-token). The clean demonstration that the mask and the objective are one decision, not two.
• Flamingo: a Visual Language Model for Few-Shot Learning — Alayrac et al. (2022) · field map extra. Note 5 at industrial scale: gated cross-attention layers interleaved into a frozen language model, queries from text, keys and values from vision. Read it for the Q-from-here / KV-from-there pattern rather than the vision specifics.
• Scaling Vision Transformers to 22 Billion Parameters — Dehghani et al. (2023) · field map extra. Where note 6's hedge bites: §2.1 adds QK-Norm (LayerNorm on q and k before the dot product) after attention logits blew up mid-training, which the initialisation-time 1/√d_k argument does not prevent. Now standard in large training recipes.
• Graph Attention Networks — Veličković et al. (2017) · field map extra. Note 1 taken literally: the same attention aggregation on explicit graphs with arbitrary adjacency. Useful for seeing that the sequence case is the special case.

Exercises

1. Break the mask on purposecode — in the Colab, comment out the masked_fill line and retrain the single-head model for the same number of steps. (1) Record train and val loss. (2) Generate 300 characters. (3) Explain in one sentence why the loss is far lower and the text far worse. (4) Now keep the mask but shift y by two instead of one in get_batch and predict two tokens ahead — does the mask still leak? A good answer names the shifted-targets relationship explicitly and states which of loss and sample quality you should trust when they disagree.
2. Measure the variance argument yourselfcode — sample q, k = torch.randn(B, T, hs) for hs in {16, 64, 256}. (1) Print wei.var() with and without * hs**-0.5. (2) Apply the causal mask and softmax, then print the mean of the largest probability in the last row. (3) Plot peak mass against hs for both. (4) Repeat with q and k scaled by 3 to simulate a trained model. You should reproduce the table above — roughly 0.36 peak mass at every head size when scaled, climbing toward 1.0 when not — and see that the scaling fixes the width dependence, not the magnitude of trained weights.
3. Turn the head into a cross-attention headcode — write CrossHead.forward(self, x, context) that takes queries from x and keys/values from context, with no mask. (1) Assert the shape of wei is (B, T_q, T_kv) for T_q != T_kv. (2) Confirm the output has T_q rows. (3) Say in one line why tril cannot be applied here. A good answer notes that Head and CrossHead differ only in the argument passed to two nn.Linear calls.

None of Karpathy's own EX1–EX4 belongs to this part — EX1 (batched heads) lands in P7, and the rest are whole-model exercises. Exercise 3 above is a decent warm-up for EX4 if you plan to implement a cross-attending variant.