Self-Attention & Multi-Head

Every token looks at every other token and decides how much to draw from each — computed as scaled dot-product attention, run in parallel heads.

9 min read Watch it move Build it

Self-attention lets every token in a sequence look at every other token *in the same sequence* and decide how much to focus on each. A word asks a question — its query — and every word offers a label — its key. Matching queries against keys gives focus weights, and those weights blend the words' actual content — their values. It's the operation that lets *"it"* in a sentence figure out which earlier noun it refers to.

Three vectors per token: query, key, value

Each token's embedding x is multiplied by three learned weight matrices to produce three vectors. The query is what this token is looking for, the key is what it advertises to be matched against, and the value is the content it will pass on once focus is decided. Same input, three different projections.

# X: (n tokens, d_model). W_Q, W_K, W_V are learned.
Q = X @ W_Q   # (n, d_k)  what each token is looking for
K = X @ W_K   # (n, d_k)  what each token offers
V = X @ W_V   # (n, d_v)  the content each token passes on

Scaled dot-product attention

1Score every query against every key: scores = Q @ K.T — an n×n grid where entry (i, j) is how much token i cares about token j.
2Divide by sqrt(d_k). Without this scaling, large dot products push softmax into tiny gradients; the sqrt(d_k) keeps the numbers stable.
3Apply softmax along each row, so every token's focus weights are positive and sum to 1.
4Multiply the weights by V: each token's output is a weighted blend of all the values.

scores  = (Q @ K.T) / sqrt(d_k)   # (n, n) match grid
weights = softmax(scores, axis=-1) # each row sums to 1
out     = weights @ V              # (n, d_v) blended values

Worked example

For the sentence the cat sat, the query for sat scores highest against the key for cat (its subject). After softmax its row might read the: 0.1, cat: 0.7, sat: 0.2, so sat's output vector is mostly cat's value — the model has wired the verb to its subject.

Multi-head: several attentions at once

One attention computation can only track one kind of relationship. Multi-head attention splits the model dimension into h smaller heads, runs the whole query/key/value process independently in each, then concatenates the heads and projects them back with W_O. In the original Transformer, d_model = 512 is split into 8 heads of size 64. One head might follow subject-verb links while another tracks adjectives — and they run fully in parallel.

Cost grows with the square of length

The score grid is n×n, so attention is O(n²·d) in both time and memory. Doubling the sequence length quadruples the cost — which is exactly why long-context models need specialized attention variants.

OperationTimeSpace

Q/K/V projection · three matrix multipliesO(n·d²)O(n·d)

Scaled dot-product · the n×n score gridO(n²·d)O(n²)

Multi-head · h heads, parallel over tokensO(n²·d)O(n²·h)

Check yourself

Why are the attention scores divided by the square root of the key dimension?