Self-attention: how every token looks at every other token

Up to this chapter we've turned text into tokens (chapter 2) and tokens into vectors (chapter 3). The interesting question is now: how does a token know about the rest of the sentence? The answer modern LLMs use is self-attention.

The intuition

Intuitively, each token looks over all the earlier tokens and pulls in the information it needs — like re-reading the sentence and focusing on the few words that matter for what comes next.

Imagine you're the model and someone hands you the prompt "The cat sat on the ___" and asks for the next word. To guess well you have to look back at the earlier tokens: cat tells you the subject is an animal, sat tells you it's resting on something, on the tells you a noun is coming next. Self-attention is a learned, tunable-by-training version of that "look back": for every token in the sequence the model decides how much it should pay attention to every other token. The panel on the right runs exactly this prompt and shows you the attention pattern.

Same word, same starting vector

Here's the cleanest way to see why that look-back is necessary. The embedding lookup from chapter 3 is context-free: it's a table indexed by token id, so the word bank fetches the same row of embed_tokens.weight whether the sentence is about money, an aircraft turn, or a riverside picnic. At the moment the last token enters the layer stack, the model literally cannot tell those meanings apart — the three starting vectors are byte-identical.

Pulling those meanings apart is the whole job of the layers above, and you can watch the result at the other end of the pipeline: feed this model three prompts that all end in the very same token and its next-token bets diverge completely. One caveat to keep you honest — in this model the disambiguation work is shared by the 6 full-attention layers and the 18 GatedDeltaNet linear-attention layers (chapter 12), so "attention did it" is shorthand for "the context-mixing stack did it."

Q, K, V — three projections of the same vector

Each token comes in as a vector x (its embedding, after earlier layers have refined it). The model linearly projects x three different ways (a "projection" here just means multiplying the vector by a learned matrix to get a new vector):

• Q = x · Wq — the query: what is this token looking for?
• K = x · Wk — the key: what does this token offer if you look at it?
• V = x · Wv — the value: what should this token contribute when it gets looked at?

Each is a vector of length d_head (the per-head dimension — the same number the original "Attention Is All You Need" paper calls d_k and the model's config calls head_dim). Crucially these aren't three different tokens — they're three different views of the same token, learned separately so attention can do something more interesting than just "compare embeddings." In this model: x is 1,024 numbers → 8 query vectors of 256 each, plus 2 key/value pairs of 256 (the 8-vs-2 trick is GQA, next chapter).

One more thing a beginner has to get right: the same three matrices Wq, Wk, Wv are applied to every token at every position — they're learned once and shared across the whole sequence, so the projection itself carries no notion of where a token sits; positional information is injected separately (by RoPE, chapter 6).

A naming note while it's fresh: because Q, K and V all come from the one running sequence, this whole chapter is self-attention. Its sibling, cross-attention, draws queries from one sequence and keys/values from another — the original use was machine translation, where a decoder query reads an encoder's keys. Cross-attention doesn't appear anywhere in this model family, though — not even in the full multimodal checkpoint. Vision fuses in via early fusion instead: projected image-patch tokens get spliced directly into the same self-attention sequence as the text tokens, which is off in this demo — none of the 6 full-attention layers here do cross-attention.

The formula

The core operation, for a single head, is one line of math:

Let's read it left to right. Q · Kᵀ (the ᵀ flips rows and columns — notation for "dot every query with every key") is a [seq_len, seq_len] matrix: cell (i, j) is the dot product of token i's query with token j's key — a raw score for "how well does i's question match j's offer?"

We divide by √d_head for a numerical reason: as d_head grows, dot products grow on average too, and without scaling the softmax would saturate (one cell at 1.0, the rest near 0.0) and gradients would vanish (a gradient is the training signal that tunes the weights — covered in the Training chapter; don't worry if it's fuzzy here). Dividing by the standard deviation of a random dot product (≈ √d_head) keeps the scores in a regime where softmax is informative.

The widget below makes that one sentence tangible. Drag the head dimension and watch the same set of key alignments: without the scale, the distribution collapses toward a single key as d_head grows; with the ÷√d_head scale it stays exactly the same at every width.

softmax turns each row of the score matrix into a probability distribution: every row sums to 1. That's the per-token attention pattern. Multiplying that [seq_len, seq_len] distribution by V mixes the values together, weighted by the pattern — and that mixture is what flows out of the attention layer as the new representation of token i.

One thing that's easy to picture wrong: that mixed value vector is not a wholesale replacement for the token. It's a small correction that gets added back onto the vector the token carried in — attention writes a delta onto the running representation, it doesn't overwrite it. That running representation has a name, the residual stream, and the add-don't-replace rule is the seam every block in this model is wired around; the full transformer block chapter shows the wiring.

That [seq_len, seq_len] matrix is also where attention gets expensive. Three optional deep-dives go one level further — into the memory cost, how to put a number on it, and the trick that sidesteps it. All are skippable on a first read.

Causal masking — and why it's the load-bearing wall of decoder LLMs

Self-attention as defined so far is symmetric: every token can see every other token. For an encoder (think BERT — fill-in-the-blank tasks) that's fine. For a generative LLM it's a catastrophe. The whole point of training is to predict the next token from what came before. If the model is allowed to look at what comes after while training, it can short-circuit the entire problem by copying.

So before the softmax we set the upper-triangle of QKᵀ to -∞ (equivalently, add a mask matrix that's 0 on and below the diagonal and -∞ above, then softmax each row). After softmax those cells become exactly 0 — token i can only attend to keys 0 … i. The heatmap on the right is lower-triangular for exactly this reason. The bottom row — the last token of the prompt — is the only row that sees the whole prompt, which is why that row produces the next-token prediction.

The toggle widget below makes this concrete: flip the mask off, and watch the model's attention for the very first token reach forward into the future of the sequence.

A subtle but important consequence: the mask is what makes parallel training possible. With the mask in place, you can shove the entire sequence into the model in one forward pass and compute the loss at every position simultaneously — each position's prediction is independent of the others because none of them can see ahead. Without the mask, you'd have to feed tokens one by one. The triangular sparsity pattern is the single architectural decision that makes decoder-only training tractable at scale.

What each cell in the heatmap means

The heatmap on the right is exactly the softmax(QKᵀ/√d) matrix for one layer and one head:

• rows are queries (the token doing the looking)
• columns are keys (the token being looked at)
• a bright cell at (i, j) means "when computing the next representation of token i, the model pulls a lot from token j's value"
• every row sums to 1; the upper triangle is 0 because of causal masking

Look at the long edges the fan-out graph draws under the heatmap: a single full-attention step can route a far-back token's content straight to the current position, no matter how many tokens sit in between. Distance is free here — only relevance costs, because the score is a dot product, not a function of how far apart two tokens are. (This free-reach claim is about the 6 full-attention layers; the 18 GatedDeltaNet layers reach back through a running state instead, and the KV-cache chapter covers why the model still needs both.)

Why multiple heads, multiple layers

Different heads end up specializing — one might track "the previous token," another "the start of the sentence," another "syntactically related noun." Stacking attention layers lets later layers compose these patterns: chapter 5 looks at heads, chapter 9 at the full stack.

One subtlety about this model: it is hybrid. Only every 4th layer — 6 of its 24 — uses the full softmax self-attention you have been reading about. The other 18 use a cheaper linear-attention variant (GatedDeltaNet) that keeps a fixed-size running state instead of attending over all past tokens (covered in the KV-cache chapter). The live heatmap on the right reflects that: scrub the layer slider onto one of the 6 full-attention layers to see real scores; land on a linear layer and the panel just notes that it does not expose a softmax attention matrix.

A note on the head slider: this model has 8 query heads but only 2 key/value groups that those heads share — grouped-query attention (GQA), which the next chapter unpacks. For now, just read each head as its own attention pattern.

Working a tiny example by hand

The matrix algebra is easier to feel with concrete numbers. The widget below walks through two tokens — "The" and "cat" — through every step of the attention formula. All numbers are computed from fixed Q, K, V values; step through and watch the score matrix, scale, mask, softmax, and final output appear.

x[The] = [1.00, 0.20]
x[cat] = [0.30, 0.90]

Q[The] = [1.00, 0.20]   Q[cat] = [0.30, 0.90]
K[The] = [1.00, 0.20]   K[cat] = [0.30, 0.90]
V[The] = [0.54, -0.16]   V[cat] = [0.33, 0.54]

out[The] = [0.54, -0.16]
out[cat] = [0.42, 0.24]

The causal mask, cell by cell

A second view of the same mask: a 5×5 grid where you can click any cell to see what "i sees 0..i" means concretely for this prompt.

Hit Run on the right. The model's predicted next token appears at the end of the prompt with a rainbow shimmer, and the heatmap shows the real post-softmax attention scores from the forward pass that produced it. Edit the prompt or change layer / head to watch the pattern (and the prediction) change.