Attention: looking back over the whole story

The problem: remembering the story

Read “Lily found a shiny shell on the beach. She put it in her pocket and ran home to show her mom. Her mom smiled and asked where she had found it.” and you know who comes next: Lily. Language keeps pointing back, to people and things introduced earlier and to what was said a while ago. A translation needs the whole sentence, an answer about a document needs the document, and a chat assistant has to keep track of the conversation. Attention was invented for this: first, in 2014, to let networks that translate look back at the whole sentence they are translating, and then, in 2017, as the core of the transformer, the design behind today's large language models.

Why is it hard?

Chapter 6's network read a window of the last 8 tokens, and making the window longer does not fix it. Every slot of the window has its own weights, 32 × 128 = 4,096 of them in the hidden layer, so a window of 256 tokens would need more than a million weights in that layer alone, and what one slot learns about names would not carry over to the others. Counting does even worse, because long stretches of text almost never repeat. Yet the word that comes next was often said long before. Here is a story the networks never saw. A word is marked when it already appeared earlier in the story, but not among the 8 tokens before it, out of chapter 6's reach; orange marks the story's name. (The 100 most common tokens, such as “the”, “she” or “.”, repeat all the time and are left unmarked, except a story's name.)

What attention does differently

• It looks things up by content, not by place. Every token offers a short list of numbers, its key, and the position that is reading asks with a list of its own, its query. Tokens whose keys match the query get the most weight, wherever they are.
• The same weights work at every distance. The weights that make queries and keys are shared by all positions, so reading 256 tokens back takes no more of them than reading 8. Only one small score per distance is added, so that order still counts.
• Several lookups at once. The network makes 8 lookups side by side, its heads, and each learns to ask its own question.

Measuring…

How attention works, step by step

Before the real network, a toy small enough to see every step. Its numbers are made up, and kept small: queries and keys have 2 numbers instead of 8, so they can be drawn as arrows, and values 3 instead of 16. The arithmetic is the real thing, computed from the arrows as you watch, except for two additions the real network makes to each score, which Tool 2 covers.

That is all attention is: compare (a query against the key of every token up to itself, by dot product), weigh (softmax), and blend (a weighted average of the values). The real network does exactly this, with lists of 8 and 16 numbers and 8 heads at once, and with two additions to each score: the dot product is divided by √8, so that scores do not grow with how many numbers the lists hold, and a learned score for the distance back is added (Tool 2). In the toy, the queries, keys and values were picked by hand. In the real network they are made from each token's row by small tables of weights, and those tables are learned by gradient descent, like every other weight.

What this is

Now the real network. This film follows one position of a story as it looks back: one head's query meets every earlier key, the scores become weights, and the weights pick what to bring back.

The question

Here is a short story whose first sentence names its hero. Its next word is a natural place for the hero's name to come back, though “she” would fit too. Change the name and see which network follows.

Chapter 6's network sees only the last 8 tokens, outlined above, and the name is not among them: it can only guess a likely name, or something else. The attention network reads the whole story. Try a rare name too, such as “fred”, or an invented one; Secrets 2 and 4 explain what happens. The tools below show how the reading works.

Tool 1: a lookup by content

Every token gets three short lists of numbers, each made from its row of the embedding table by a small table of weights: a key, what it offers to be found by; a value, what it hands over when it is read; and a query, what it is looking for. To read, a position compares its query with the key of every token up to itself. The comparison is the dot product: multiply the two lists number by number and add up. Softmax turns the scores into weights that add up to 1, and the position takes the weighted average of the values. Queries and keys have 8 numbers each, values 16. Click any token to make it the one reading, and pick a head.

Not chapter 6's closeness. In chapter 6, two words were close when their rows pointed the same way, measured by the cosine of the angle between them. Attention's score is not that. The query and the key are made from the rows by two different tables, and each head has its own pair, so each head learns its own sense of which keys fit which queries. On a wide screen, the table's last column gives that cosine for this network's own rows, only to compare. Over all 2.6 million pairs of a token and a token it can read in the 100 test stories, the dot product does not follow the cosine of the two rows in any head (correlations between −0.12 and 0.09).

Between a query and a key, the angle counts, but not alone. The dot product is the query's length times the key's length times the cosine of the angle between them (each arrow's length, not how many numbers it holds). Dividing by √8 is the same for every pair, so it is not chapter 6's cosine, which divides the dot product of two rows by both their lengths. The angle still orders a position's keys much as the dot product does, but the keys' lengths often decide which comes first: a long key that points further off can beat a shorter one that points closer to the query's way. In 48% to 65% of the positions of the test stories, depending on the head, the key closest in direction to the query is not the one with the largest dot product. (In symbols, with x for a token's row: q_i · k_j = x_i^T W_q^T W_k x_j. The 32 × 32 table W_q^T W_k is made of two learned tables, differs in each head, has rank at most 8 and is not symmetric, so it measures no angle between rows at all.)

The scores are a soft version of looking something up. A dictionary finds the one entry whose key equals what you ask for; attention gives every entry some weight, more the better its key matches, and returns a blend. Because every step is a sum, a product or softmax, gradients flow through all of it, and the small tables that make queries, keys and values are learned like any other weights.

Tool 2: where, as well as what

A dot product of a query with a key says nothing about where the key sits: move the earlier tokens around and the weights move with them. But order matters (“the dog chased the cat” is not “the cat chased the dog”), so each head also learns a score for every distance back, from 0, the token itself, to 255, which is added to the dot product before softmax. A head with a high score at distance 1 reads the token just before, whatever it is. A head with flat scores reads by content alone. Here are the 8 heads' distance scores, as learned.

GPT-2 marks order differently: it adds a learned row for each position, 0 to 1,023, to the token's row. Scores that depend on how far apart two tokens are, as here, are closest to those of T5 (a learned score per distance, shared by groups of distances) and ALiBi (a fixed penalty growing with distance). Rotary position embeddings, used in many recent large models, also make the score depend on the distance, by turning queries and keys through angles that grow with position.

Tool 3: the network around it

The rest is chapter 6's network. The hidden layer reads the token's own row of 32 numbers and the 8 heads' results of 16 numbers each, 160 numbers in all; 128 tanh neurons read them, and an output layer scores the 2,000 tokens. Type a phrase and follow its last token through.

So the numbers the hidden layer reads for a token now depend on the story before it: the same token in two stories brings its own row, but the heads bring back different results. The token's own row is still the same in every story, but what the hidden layer reads alongside it now comes from anywhere in the window, chosen by content, rather than from 8 fixed slots. That is the start of the fix chapter 6 promised for “one word, one row”. In a single layer, these rebuilt numbers serve only to guess the next token; chapter 8 stacks layers, so that later ones can read them.

In all, 352,848 numbers to learn, about as many as chapter 6's 354,896: the attention network spends no more weights, but reads up to 256 tokens back instead of 8. Training is chapter 6's too: the same loss, −log of the probability given to the token that came next, and the same Adam. What changes is what a step reads. No position ever reads a later token, so one forward pass over a whole story makes a guess at every position at once, and the gradients of all of them are added up. Each training step reads 8 whole stories, about 2,000 positions, instead of chapter 6's 64 positions picked at random.

Attention and the transformer

Attention is one part of a larger design, the transformer, which is what “T” stands for in GPT. This page's network is not yet a transformer: it is one layer of attention with chapter 6's layers around it. This animation goes from chapter 6's network to this one, then to a transformer block, a stack of blocks, and a GPT.

The first transformer, in 2017, was built for translation. It had two stacks: an encoder of 6 blocks that read the sentence to translate, every token seeing every other, and a decoder of 6 blocks that wrote the translation one token at a time. In each decoder block, one attention read only the tokens already written, as here, and a second attention read the encoder's rows. GPT models keep only the decoder kind, without the encoder or the decoder's reading of it: a stack of blocks with causal attention that guesses the next token. How big they get:

The first transformer's sizes are those of its base model. Later models changed many details, such as how positions are marked, but keep this shape: embeddings, a stack of attention-and-MLP blocks with residual connections and layer norms, and an output layer.

The words, and how they fit

Attention

A way for each position to read other positions by content: compare its query with their keys, turn the scores into weights with softmax, and take the weighted average of their values.

Query, key and value

Three short lists made from each token's row by three small tables of a head: what the token is looking for, what it offers to be found by, and what it hands over when it is read. A query and a key are compared by their dot product, a match the head has learned, not the closeness of the two words' rows.

Self-attention

Attention whose queries, keys and values all come from the same text, as on this page. When the keys and values come from another text, as when a translation reads the sentence it translates, it is called cross-attention.

Causal attention

Attention in which a position reads only itself and earlier positions, never a later one (also called masked attention). It is what lets one pass over a story train every position at once at guessing the next token, without any position seeing its own answer.

Head

One attention lookup, with its own small tables for queries, keys and values. Multi-head attention runs several heads side by side and joins their results; this page has 8 heads, each bringing back 16 numbers. In a transformer, one more table then mixes the joined results back to the width of a row, so they can be added to it; this page has no such table and lays them beside the row instead.

MLP

A small network of the kind in chapters 1 to 4 (a multilayer perceptron, also called a feed-forward layer), applied to each position on its own. Attention moves information between positions; the MLP works on what each position holds.

Residual connection

Adding a part's input back to its output, so a token's row passes through a block and only gains what the block adds. This page lays the token's own row beside the heads' results instead.

Layer norm

Rescaling each row to an average of 0 and a spread of 1, then applying a learned scale and shift. It keeps the numbers in a steady range as blocks are stacked.

Transformer block

Multi-head attention and an MLP, each with a layer norm and a residual connection. Rows in, rows out, so blocks can be stacked.

Transformer

A network built from a stack of transformer blocks, on top of token rows and a way of marking positions.

GPT

A generative pre-trained transformer: a stack of blocks with causal attention, trained to guess the next token on a large amount of text. Chapter 8 builds a tiny one.

Secret 1: the heads found jobs

Nobody told the heads what to look for. All 8 started alike, with small random tables and every distance score at 0, so at first each one spread its weight almost evenly over the story. Here is where each head's weight goes in the 100 test stories, by distance back, and which tokens it reads most, at several moments of training.

Measuring…

Secret 2: it copies the names it knows

Does the network bring back the name from the story, or just guess a common one? Swap the name everywhere in a story for another, and watch the places where the name comes back.

Measuring…

Secret 3: switch a head off

Secret 1 shows where each head looks. To see what each head is for, switch it off: its 16 numbers become 0, and everything else stays as trained. Pick a head to see how the network's guesses on the 100 test stories change without it.

Measuring…

Secret 4: what one layer cannot do

Copying a whole name is one lookup: find a name earlier in the story and bring it back. Finishing a name made of pieces is two. Give a story's hero an invented name, such as “zarko”, cut into the pieces “z”, “##ark” and “##o”. At a later mention, after “z”, the right next piece is the one that came right after “z” before: the network must find the earlier “z” and then read the token after it.

Measuring…

The algorithm: attention, and the network

Five parts, from the weights to a training step. This is the JavaScript this page actually runs, or the same program in Python. The tokenizer is chapter 6's.

The playground

Write the start of a story, click any token to see where each head looks from there, and let the network write on. Or train your own network from scratch and watch its heads find their jobs.