Chapter 3 · About 15 minutes

Attention

Attention lets each token gather what it needs from itself and the tokens before it.

You will learn

Say what attention does in one sentence
Name the query, key and value and what each one is for
Turn three scores into weights with softmax
Read a real attention grid, and say what it can and cannot tell you

In "Lily wanted to play with her", which earlier word does "her" look at most, in one of the model's heads?

Words need context

In layer 3, head 3, space her looks most at Lily 0.66, space her 0.12, space wanted 0.10.

Attention weights in layer 3, head 3. Each row is a token, and each column is a token it can look at. Each row adds up to 1.
	Lily	space wanted	space to	space play	space with	space her
Lily	1.00	hidden	hidden	hidden	hidden	hidden
space wanted	0.87	0.13	hidden	hidden	hidden	hidden
space to	0.66	0.20	0.15	hidden	hidden	hidden
space play	0.53	0.29	0.14	0.04	hidden	hidden
space with	0.79	0.10	0.08	0.03	0.00	hidden
space her	0.66	0.10	0.09	0.02	0.01	0.12

Top three for every token

Lily looks at Lily 1.00.
space wanted looks at Lily 0.87, space wanted 0.13.
space to looks at Lily 0.66, space wanted 0.20, space to 0.15.
space play looks at Lily 0.53, space wanted 0.29, space to 0.14.
space with looks at Lily 0.79, space wanted 0.10, space to 0.08.
space her looks at Lily 0.66, space her 0.12, space wanted 0.10.

Pick another token, layer or head.

Story

So far, each token has its own list of numbers, and nothing else. But a word means more with the words around it. Attention The step where each token gathers information from itself and the tokens before it, taking more from the ones that match what it is looking for. Open the glossary is how a token gathers that context (Vaswani et al., 2017).

Each token looks at itself and the tokens before it, decides how much each one matters, and pulls in a mix of their numbers. In this head, “her” looks mostly at “Lily”.

Numbers

Each row of the grid is one token, and its weights add up to 1. In layer 3, head 3, “·her” gives 0.66 to “Lily” and spreads the rest thinly. Every other row looks most at “Lily” too.

The top right half is empty because no token can look ahead at the tokens after it. This is called the Causal mask The rule that a token may only look at itself and the tokens before it, never at later ones. Open the glossary. A 0.00 means less than 0.005, not exactly 0.

Formula

Token $i$ gets a weighted mix of the value vectors $\mathbf{v}_j$ of the tokens up to it.

\mathbf{o}_i = \sum_{j \le i} a_{ij} \, \mathbf{v}_j, \qquad a_{ij} \ge 0, \quad \sum_{j \le i} a_{ij} = 1

Code

trace.layers[2].heads[2].weights; // layer 3, head 3, 6 × 6, each row adds up to 1

Questions, name tags and backpacks

A made-up example with small numbers.

Lilyquestion (1, 0)name tag (2, 0)backpack (1, 0)
ranquestion (0, 1)name tag (0, 1)backpack (0, 1)
shequestion (2, 0)name tag (1, 1)backpack (0.50, 0.50)

Step 1 of 5. Match questions with name tags

Each token's question is compared with every name tag using a dot product. A big score means a good match.

Rows are the token asking. Columns are the token being looked at.
Asking token	Lily	ran	she
Lily	2	0	1
ran	0	1	1
she	4	0	2

Step through attention with three tokens.

Story

Each token makes three new lists from its numbers. Its Query The list of numbers a token uses to ask what it is looking for, like a question card. Open the glossary is a question card, saying what it is looking for. Its Key The list of numbers that says what a token has, like a name tag. Queries are matched against keys. Open the glossary is a name tag, saying what it has. Its Value The list of numbers a token shares when others look at it, like a backpack. Open the glossary is a backpack, holding what it will share.

A token compares its question with every name tag using a Dot product A way to compare two lists of numbers. Multiply the numbers in matching places and add up the results. Open the glossary, and takes the most from the backpacks whose tags match best. Press Next to see each step.

Numbers

In the made-up example, “she” asks with (2, 0). Lily’s tag (2, 0) scores 4, the tag of “ran” scores 0 and its own tag scores 2. Each list has 2 numbers, so the scores are divided by the square root of 2, about 1.41, giving 2.83, 0 and 1.41. After Softmax A way to turn any list of scores into numbers between 0 and 1 that add up to 1. Bigger scores get bigger shares. Open the glossary, “she” takes 0.77 of Lily’s backpack.

In our model each head’s lists have 32 numbers, so the scores are divided by the square root of 32, about 5.66.

Formula

Stack the queries, keys and values of all tokens as the rows of $Q$ , $K$ and $V$ . Then

\operatorname{Attention}(Q, K, V) = \operatorname{softmax}\!\left(\frac{QK^\top}{\sqrt{d_k}} + M\right) V, \qquad M_{ij} = -\infty \text{ when } j > i

Each of $Q$ , $K$ and $V$ comes from the token vectors times a learned matrix, for example $Q = X W_Q$ . GPT-2 first normalizes the vectors and also adds a bias. In our model $d_k = 32$ .

Code

const head = trace.layers[0].heads[0];
head.q; // queries, 6 × 32 (keys head.k and values head.v are the same shape)
head.scaledMasked; // scores divided by √32, with -Infinity where a token would look ahead

Like a classroom

Each student holds up a question card, everyone wears a name tag, and each student copies a little from the backpacks of classmates whose tags match. Where this breaks. A real student copies from one or two classmates. A token always takes at least a little from every earlier token, and it can only look at classmates who came before it.

Softmax turns scores into weights

Story

Scores can be any size, even below zero. Softmax turns a row of scores into weights between 0 and 1 that add up to 1.

It raises the number e, about 2.718, to the power of each score, then divides each result by their total. Bigger scores get a much bigger share, but every score that is not hidden gets at least a little.

Numbers

Take the scores 1, 0 and 0. e to the power of 1 is 2.718, and e to the power of 0 is 1. The total is 4.718, so the weights are 0.58, 0.21 and 0.21.

A hidden score counts as minus infinity, and e to the power of minus infinity is 0. So hidden tokens get a weight of exactly 0. This is how the causal mask works.

Formula

\operatorname{softmax}(\mathbf{z})_j = \frac{e^{z_j}}{\sum_k e^{z_k}}

Code

softmax([1, 0, 0]); // [0.576, 0.212, 0.212]

Many heads, many layers

In layer 2, head 1, space her looks most at space with 0.96, space play 0.03.

Attention weights in layer 2, head 1. Each row is a token, and each column is a token it can look at. Each row adds up to 1.
	Lily	space wanted	space to	space play	space with	space her
Lily	1.00	hidden	hidden	hidden	hidden	hidden
space wanted	0.31	0.69	hidden	hidden	hidden	hidden
space to	0.12	0.61	0.27	hidden	hidden	hidden
space play	0.04	0.13	0.58	0.25	hidden	hidden
space with	0.01	0.10	0.16	0.27	0.47	hidden
space her	0.00	0.00	0.00	0.03	0.96	0.00

Top three for every token

Lily looks at Lily 1.00.
space wanted looks at space wanted 0.69, Lily 0.31.
space to looks at space wanted 0.61, space to 0.27, Lily 0.12.
space play looks at space to 0.58, space play 0.25, space wanted 0.13.
space with looks at space with 0.47, space play 0.27, space to 0.16.
space her looks at space with 0.96, space play 0.03.

Try every head. Each one has its own pattern.

Story

One way of looking is not enough. Each layer runs several Attention head One of several attention steps that run side by side in a layer, each with its own queries, keys and values. Open the glossary side by side, each with its own questions, name tags and backpacks. They learn different jobs.

In layer 2, head 1 often looks at the token just before. In layer 3, head 3 looks back at “Lily” from almost every word. Their results are joined and passed on to the next layer.

Numbers

Our model has 4 layers with 4 heads each, 16 heads in all. Each head works with lists of 32 numbers, and the 4 heads together make 128 again. GPT-2 small has 12 layers (Radford et al., 2019).

Formula

The heads’ outputs are placed side by side and mixed by one more learned matrix, plus a bias.

\operatorname{MultiHead}(X) = \big[\,\text{head}_1 \;\cdots\; \text{head}_4\,\big] \, W_O

Code

trace.layers[1].heads[0].weights; // layer 2, head 1
trace.layers[0].attnConcat; // the 4 heads side by side, 6 × 128

Common myth

Attention weights show where a head looked, not why the model decided what it did. In “Ben wanted to play with her”, the head from the first grid puts 0.74 on “Ben”, so it is not working out who “her” means. In other kinds of models, researchers found that quite different weights can lead to the same prediction (Jain and Wallace, 2019). Others argued the weights can still be useful clues, if you test them carefully (Wiegreffe and Pinter, 2019). Treat the grid as a hint, not an explanation.

Recap

Attention lets each token pull in information from itself and the tokens before it.
A token's query is matched with every key using a dot product, giving one score each.
Softmax turns the scores into weights that add up to 1, and the token blends the values with them.
Each layer has several heads, and each head learns its own way of looking.

Go deeper

Vaswani et al., 2017. Attention Is All You Need
Radford et al., 2019. Language Models are Unsupervised Multitask Learners
Jain and Wallace, 2019. Attention is not Explanation
Wiegreffe and Pinter, 2019. Attention is not not Explanation