Ch 9. Attention

How does a token decide where to shine? Each token makes three things from its numbers, using weights learned in training: a query (what it’s looking for), a key (what it has to offer), and a value (what it passes on). Where one token’s query matches another token’s key well, the spotlight shines brighter. Each token then takes in a blend of the values, mostly from the tokens it lights up most.

How does a token decide where to shine? Each token’s vector is turned into three new vectors, using weights learned in training:

• a query: what this token is looking for,
• a key: what this token has to offer, to anyone looking,
• a value: what this token passes on if it gets picked.

To score how relevant one token is to another, multiply the first token’s query with the second token’s key, number by number, and add up the results. When the two vectors point in similar directions (think of the nearby words in Ch 5), the score is high. The scores are divided by a fixed number to keep them from getting too large, then softmax (Ch 2) turns them into percentages. Finally each token takes a weighted blend of the values.

The formula

Written out for every token at once, one attention head computes

Each row of $Q$, $K$, and $V$ is one token’s query, key, or value. They come from the token vectors $X$, one row per token, times three learned weight matrices: $Q=X{W}_Q$, $K=X{W}_K$, and $V=X{W}_V$. $Q{K}^{\top }$ is the grid of scores, every query times every key. $d_k$ is the number of values in each key, and $\sqrt{d_k}$ is the fixed number the scores are divided by. $M$ is the mask: 0 where a token may look, and $-\infty$ where it would look ahead, which softmax turns into a weight of exactly 0. Softmax runs along each row, so each token’s weights add up to 1.^[1]

In GPT-2, each token vector has 768 numbers and each layer has 12 heads, so each head works with $d_k=768/12=64$ numbers, and the scores are divided by $\sqrt{64}=8$. The 12 heads’ outputs are put side by side, back to 768 numbers, and multiplied by one more learned matrix, $W_O$.^[1] Unlike the toy, GPT-2 also normalizes the token vectors before they go in (more in the next chapter).^[2]

The grid has a cell for every pair of tokens, so it grows with the square of the text’s length. With the mask, $n$ tokens need

scores in each head: 524,800 for a full window of 1,024 tokens, in each of GPT-2’s 144 heads.

Attention on its own has a blind spot: nothing in a score says where a token sits. “The dog chased the cat.” and “The cat chased the dog.” contain the same words, so attention would blend them the same way, and who chased whom would be lost. So GPT-2 adds a second list of numbers to each token, a position vector, that says where it is: first, second, third, and so on, up to 1,024 places.

Look again at how a score is made: one token’s query times another token’s key. Nothing in that says where either token sits in the text. Take “the dog chased the cat.” and “the cat chased the dog.” At the final ”.”, one attention step would see the same six tokens in both, just in another order, and blend them in exactly the same way. Who chased whom would be lost.

So before the first layer, GPT-2 adds a second vector to each token’s vector: a position vector, one for each place in the text, learned in training like the token vectors. (This is position in the text, not the position in meaning from Ch 5.) Now “dog” second and “dog” fifth start out as different vectors, and attention can tell the sentences apart. GPT-2 has 1,024 position vectors, one for each place in its context window (Ch 8). There is no vector for place 1,025, which is one reason the window stops there.

Newer models such as Llama 3 add no position vectors. Instead they turn each query and key by an angle that grows with its place in the text, so the score between two tokens also depends on how far apart they are.^[2][3][3][4]

In symbols, GPT-2’s first layer starts each token from

where $i$ is the token’s place in the text, $t_i$ is its token ID, $e_{t_i}$ is that token’s row in the 50,257 × 768 table of token vectors, and $p_i$ is row $i$ of the 1,024 × 768 table of position vectors.

Rotary position embedding, the method Llama 3 uses, works on the queries and keys instead. It splits each one into pairs of numbers and turns pair $j$ of the token in place $m$ by the angle $m{\theta }_j$, with

where $d$ is the length of the query and $b$ is a fixed base. When the query in place $m$ meets the key in place $n$, the two turns partly cancel, and the score depends on the places only through the distance between them, $m-n$.^[3] The original paper used $b=\mathrm{10,000}$; Llama 3 raised it to 500,000 to handle longer texts.^[4]

GPT-2 isn’t very good at that last kind of puzzle. The patterns you can read are a few bright spots in a mostly murky picture.

GPT-2 isn’t very good at that last kind of puzzle. Even the largest GPT-2 answered 70.70% of Winograd questions correctly, where guessing gets 50%.^[5][2] The patterns you can read are a few bright spots in a mostly murky picture.

One set of queries, keys, and values is called an attention head. GPT-2 runs 12 heads side by side in each of its 12 layers, each free to look for something different. Switch heads in the viewer: many don’t have a job you can name, and many put most of their attention on the very first word, a resting spot for when there’s nothing useful to find.^[3]

One set of queries, keys, and values is called an attention head. Each layer of GPT-2 runs 12 heads side by side, each with its own weights, so each can look for something different. With 12 layers, that’s 144 heads. Layers and heads are numbered from 0, as in research papers.

Switch to other heads in the viewer and you’ll see that many don’t have a job you can name. You’ll also see lots of heads putting most of their attention on the very first word. In these sentences, from the fourth word on, that happens more than half the time. Researchers think it’s a resting spot: softmax always has to share out 100%, so a head with nothing useful to find parks its attention there.^[6][6]

Here is one head in full, for the sentence about Mary and John. Each row is a word looking back, each column a word it could look at. Step through what the head works out: a score for every pair, then the mask that hides every later word, then softmax, which turns each row into percentages. The buttons page through all 144 heads.

Here is one head in full, for the sentence about Mary and John: every row is a word doing the looking, every column a word it could look at. Step through the formula from the queries and keys above: first the score for every pair, then the mask, then softmax along each row. Page through all 144 heads with the buttons.

The head works out a score for every pair, even for words that come later; the mask then throws those away, 91 of the 196 squares. For head 9.9, the last word’s scores for “Mary” and “John” are close, 4.25 and 3.23, yet after softmax “Mary” gets 36% and “John” 13%. Softmax turns a small lead in score into a big lead in weight.

The scores exist for every pair, later words included: the head computes them all at once, and the mask throws away the 91 of 196 that look ahead before softmax sees them. For head 9.9, the last word’s scores for “Mary” and the first “John” differ by about 1 (4.25 against 3.23), yet after softmax “Mary” gets 36% and that “John” 13%. Softmax turns small gaps in score into large gaps in weight. “When” scores 4.24 and also gets 36%: the resting spot described above.

Since a weight is $e^s$ divided by the row’s total, two scores that differ by $\mathrm{\Delta }$ get weights in the ratio $e^{\mathrm{\Delta }}$, whatever the rest of the row holds: $e^{4.25-3.23}\approx 2.8$, and 36% is about 2.8 times 13%.