Attention: words looking at words

Read this: "The trophy didn't fit in the suitcase because it was too big." What does "it" mean? The trophy. Change "big" to "small" and "it" becomes the suitcase. To predict well, a model has to work out which earlier words matter for the word it's on. The mechanism that does this is called attention. It's the core of the transformer, the design almost every modern language model is built on.

How attention works, roughly

At each layer, every token turns its list of numbers into three new ones.

• A query, which describes what it's looking for. For "it", something like "I'm a pronoun. Which thing am I pointing to?"
• A key, which advertises what it offers. "Trophy" might advertise "I'm a physical object, and I'm the subject".
• A value, the information it hands over to any token that attends to it.

Each token compares its query with the key of every earlier token. Strong matches get a big share of its attention and weak matches get almost none. Then it blends together the values it attended to, weighted by those shares, and adds the blend to its own numbers. "It" leaves the layer carrying information about the trophy.

Those plain-English descriptions are a cartoon. Real queries and keys are lists of numbers that nobody wrote. They're learned in training like everything else. The matching and blending, though, is exactly how it works.

When a model is writing text, each token can only look backward, at itself and the tokens before it. It can't peek at words that haven't been written yet.

Choose a query token and press Play to trace its links. Only that token and earlier tokens can contribute to its update.

Many heads, many layers

One attention pattern can't track everything, so each layer runs many of them side by side, called heads. Researchers have found heads that seem to have recognizable jobs. One looks at the previous token. Another finds an earlier spot where the current text appeared before and copies what came next. Plenty of heads have no job anyone has figured out.

Between attention steps sits a second kind of block, usually called the feed-forward layer. It works on each token separately, and much of what the model knows seems to be stored in these blocks. A large model stacks dozens of attention and feed-forward pairs. A token's numbers get refined at every one. After the last layer, the numbers at the final position are turned into the scores for the next token, the ones you played with in the first lesson.

Why this design won

Before transformers, the leading language models read text one word at a time, carrying a running summary forward. Connections between distant words faded, and training had to step through each text in order, which made it slow to spread across many chips. Attention lets every token look directly at every earlier one, and training can handle all the positions in a text at once. That fit the graphics chips used for training very well, and it made training on enormous amounts of text practical.

A team at Google introduced the transformer in 2017, in a paper titled "Attention Is All You Need". Models have grown enormously since then. The basic recipe has barely changed.

There is a cost. Every token compares itself with every earlier one, so doubling the length of a text roughly quadruples the attention work. That's one reason there's a limit on how much a model can read at once.