Tutorial · 8

Attention

Attention is a weighted average where the weights are computed from the data. That is the whole of it, and it is three lines of arithmetic you can read. This tutorial writes those lines out, trains them on a task built for them, reaches for nn.MultiheadAttention once the hand-written version is understood, and finishes on the thing attention genuinely cannot do.

1 · Three lines, no training

Six tokens, each a vector. Compare every token with every token (Q·Kᵀ), turn the comparisons into weights that sum to one (softmax), and average the values with them. The two copies of "the" have the same vector, so they end up attending to each other — visible in the map before anything has been trained.

Row 0 and row 4 are the same row: identical vectors compare identically against everything. That is attention having no notion of where a token sits, which block 4 is about.

2 · A task that needs it

Eight positions. Each carries a random value, and exactly one carries a flag. The answer is the value at the flagged position — a product of two inputs, so no linear model can reach it, and a network that averages everything cannot either. One query vector, learned, is enough: it has to learn to look at the flag.

3 · The same thing, packaged

nn.MultiheadAttention is that arithmetic with the projections built in and several heads running side by side. It also takes a mask, and the useful mask is the causal one — it fills the upper triangle with -Infinity so that softmax gives those positions zero weight, which is what makes a language model unable to read ahead.

Note the shape it is handed: [length, batch, embed], because that is torch's default and this class follows it. batchFirst is the option, not the rule — a layer that reads [batch, length, embed] unconditionally is the one setting torch does not have.

A mask is easy to believe and easy to get wrong, so this measures it: change the last token and see whether the first position's output moves.

Exactly zero, not nearly zero. The masked positions contribute nothing at all, because softmax of -Infinity is 0 and not a small number — worth checking rather than assuming, since a mask that is merely very negative gives a leak that no loss curve will show you.

4 · Attention does not know what order anything is in

Now the answer is "the value at position 2" — no flag to find. Attention compares content with content, and the tokens are indistinguishable, so it cannot get there. The fix is to write the position into the token, and how you write it decides whether the fix works: a ramp t/7 only lets a query prefer "later" or "earlier", while a sine and cosine put the positions on a circle where a query can point at one of them.