Self-attention and multi-head attention

In self-attention, queries, keys and values all come from the same sequence. Each token's embedding ${\mathbf{x}}_i$ is projected three ways by learned matrices:

Each token asks a question (its query), advertises what it contains (its key), and offers content (its value). In "The animal didn't cross the street because it was too tired", a well-trained model's query for "it" matches the key of "animal" strongly, so the output for "it" mixes in information about the animal. Unlike an RNN, any two positions interact directly in one step, however far apart, and all positions are computed in parallel, which suits GPUs perfectly.

The causal mask

A language model is trained to predict each next token. If position 3 could attend to position 4, it could simply read the answer. A causal mask blocks this: set the scores for every future position to $-\infty$ before the softmax, so their weights become exactly 0. Each token sees only itself and earlier tokens, and one forward pass trains the prediction at every position at once.

Multiple heads

One attention pattern per layer is limiting: a token might need to look at its subject, the previous word, and a matching bracket all at once. Multi-head attention runs $h$ attention operations in parallel, each with its own smaller projections of size $d_{\text{model}}/h$, then concatenates the results and mixes them with one more matrix $W_O$. With $d_{\text{model}}=512$ and 8 heads, each head works in 64 dimensions, so the total cost is about the same as one big head.

Researchers studying trained transformers have found heads with recognisable jobs, such as attending to the previous token, or "induction heads" that find an earlier occurrence of the current token and copy what followed it (Olsson et al., 2022), a mechanism behind in-context learning.

The cost

The score table is $n\times n$ for a sequence of $n$ tokens, so compute and memory grow quadratically with length: doubling the context quadruples the attention cost. Much engineering (FlashAttention, sparse and linear attention variants, sliding windows) goes into making long contexts affordable.

Self-attention projects each token into a query, key and value, so every token can draw on every other in one parallel step. A causal mask hides the future; multiple heads look for different relationships.