Self Attention

Picture a dinner party. Twenty guests around a long table, everyone talking at once. Now picture the same party the way a RNNs would have to run it: one guest speaks, the next guest listens to a compressed summary of everything that's been said so far, then whispers their own summary to the guest on their right. By the time the rumor has crossed the room, the first joke has been rewritten, forgotten, and re-forgotten through nineteen hops of a very tired game of telephone. Long-range conversation is possible in principle and miserable in practice.

Self-attention runs the party differently. Every guest turns to every other guest at the same instant and silently asks the same question: how relevant are you to me right now? No passing notes down the chain. No summarizing. Guest 100 hears guest 1 with the same clarity as guest 99 — the shortest path between any two people in the room is always one look across the table. That single architectural decision — “everyone reads everyone, in parallel” — is the engine underneath BERT, GPT, Claude, every modern LLM you've touched.

This lesson builds the room from the ground up. We'll meet the three things each guest carries, derive the Q/K/V projection trick, unpack the scaled dot-product formula a factor at a time, see why the √d_k in the denominator is mathematically non-negotiable, watch a live attention matrix light up over a real sentence, and implement causal attention from scratch in three flavors.

Three learned matrices. One matmul, one scale, one softmax, one more matmul. That is all of attention. Every paper you'll read in this section is a variation, optimization, or re-packaging of the diagram above — the same dinner party, re-catered.

Now the reveal. Each guest at the table arrives carrying three things — not two, not five, three — and once you see them separately the whole mechanism clicks. Given a sequence of N token embeddings (our word embeddings for the guests) X ∈ ℝ^(N × d), attention projects each guest three ways:

Same guest, three different pitches. A guest's query is “what am I interested in right now?” — the topic they're trying to follow. Their key is “here's what I have to offer” — the label they're broadcasting across the room, the nametag that tells everyone else what they're about. Their value is “here's what I'll actually say if you call on me” — the contribution they make once they're chosen. The weights W_Q, W_K, W_V are the only learnable parameters in a single-head attention block, and they're what distinguish a great party host from a random one — they've learned which questions to ask, which nametags to print, and which things to say when called on.

Why three roles and not one? Because the question a guest is asking (“who here is talking about music?”) is genuinely different from the nametag they're wearing (“I play jazz”) which is different from the story they'll tell when called on (a fifteen-minute anecdote about a bass player in 1974). Collapse any two of those roles into one and the mechanism stops working — the room loses the ability to look for something different from what it's offering.

Four guests at the table, fully projected. Highlight any query row and the widget draws a dot product with every key — that's one guest turning to the room and sizing up every nametag. The values don't enter the picture yet; they're waiting patiently on the right, ready to be weighted and summed once softmax has decided who gets the floor. Notice the shapes: Q and K are the same width (d_k, so their dot product is scalar), V can be a different width (d_v) — in practice they're usually equal, but nothing in the math requires it.

One line. Four operations. Let's unpack it piece by piece, because every factor in there is doing work.

Step 1 — raw scores. Q · Kᵀ is an N × N matrix. Entry (i, j) is guest i's query dotted with guest j's key — how interested this guest is in what that guest is offering. Every pair in the room scored at once, in a single matmul. This is the payoff of the matrix formulation: the GPU does the whole thing with no Python loops in sight.

Step 2 — the √d_k scaling. This looks arbitrary. It isn't. Suppose q and k are vectors of dimension d_k with unit-variance entries. Their dot product is a sum of d_k zero-mean unit-variance products, so by the central limit theorem the dot product has variance d_k and standard deviation √d_k. For d_k = 64, scores routinely land at ±8. Feed that into softmax and one entry dominates so hard that the gradient on every other entry goes to ~0 — the softmax has saturated. A single guest is shouting so loudly that nobody else in the room can be heard. Dividing by √d_k pulls the variance back to 1, keeping the distribution soft and gradients alive.

Step 3 — softmax. Applied along the keys axis (rows of the score matrix). Row i becomes a probability distribution over keys: exactly how guest i splits their attention budget across the room. They have one unit of listening to spend, and softmax tells them how to spend it. The “wrong axis” mistake — softmaxing along queries instead of keys — is one of the most common bugs in from-scratch attention. Read Step 3 twice and tape it to your monitor.

Step 4 — weighted sum of values. Multiply the N × N attention weights by the N × d_v value matrix. Output row i is a convex combination of every guest's value vector, weighted by how much guest i was listening to each of them. Each output is a context-aware impression of that position — a single guest's final takeaway from the room, assembled out of what everyone else said, pre-weighted by how much they mattered.

That's an actual attention matrix from a small model run on a real sentence. Scrub the query selector and watch one row at a time light up — that row is a single guest showing you who they ended up listening to. The bright diagonal is a hint that guests attend to themselves a lot (sensible — a token's own embedding is usually the best single source of information about it, in the same way that knowing what you came to the party to talk about is a decent prior on what you'll say next). The off-diagonal heat is where the interesting learning lives: a pronoun listening to its antecedent, a verb listening to its subject, a closing bracket listening to its opener. Every pattern you see here is emergent — no one programmed “link pronoun to antecedent,” the network learned to eavesdrop that way because doing so lowered the loss.

Three implementations of the same block. Pure Python on a tiny 4-guest party so you can see every index. NumPy with einsum to do the whole thing in one expression. PyTorch's F.scaled_dot_product_attention — the call you actually use in production, which dispatches to FlashAttention kernels under the hood.

Now with NumPy. The whole block collapses to five lines. einsum makes the axis contractions explicit — "nd,md->nm" is literally “for each (n, m) pair, sum over the d axis,” which is the definition of Q · Kᵀ. Every guest scoring every other guest, one index at a time, but all at once.

And now PyTorch. In real code you never handwrite attention — you call F.scaled_dot_product_attention, which dispatches to FlashAttention or memory-efficient kernels depending on your hardware. Or for a full multi-head block, nn.MultiheadAttention. What used to be 20 lines of NumPy is a single call with a causal flag.

One more piece — causal masking. A language model generating text can't be allowed to peek at future tokens while predicting the next one. At our dinner table: guest t is only allowed to listen to guests who've already spoken. Letting them eavesdrop on guests who haven't arrived yet would be training-time cheating of the most embarrassing kind. The fix is a mask applied to the score matrix before softmax: set every entry above the diagonal to -∞, so after softmax those weights are exactly 0 and never leak into the output.

This is what is_causal=True does inside PyTorch's SDPA — it fills the upper triangle with -∞ before the softmax runs. In a decoder-only model like GPT, every attention layer is causal. In an encoder like BERT, none of them are — everyone at the table speaks at once and hears everyone else. In an encoder-decoder like T5, the encoder is bidirectional and the decoder is causal with an extra cross-attention step.

What to carry forward. Attention is a dinner party rendered in linear algebra: every guest turns to every other guest with a query, a key, and a value — three learned projections, one scaled dot product, one softmax, one weighted sum of values. The √d_k keeps the softmax from saturating at reasonable d_k. Softmax goes along keys. Masks go in before softmax. The whole room is permutation-equivariant, which is why we'll staple positional encodings to the input. And the O(N²) memory cost — every guest scoring every other guest — is the single biggest bottleneck in modern transformer scaling, which is why your next six papers will all be about sidestepping it.

Next up — Multi-Headed Self Attention. One conversation at one table is fine — but the room would be richer if every guest could have several conversations at once, each specializing in a different angle. A pronoun probably wants to find its antecedent and its syntactic role and its verb, simultaneously. The answer is multi-headed self attention: run several attention heads in parallel with different W_Q / W_K / W_V, concatenate their outputs, and project down. We'll derive why n_heads matters, what each head tends to specialize in, and how a 12-head attention block is still a single matmul if you squint at it right.