Attention

To understand what a word means, we often need to look at other words in the document, including some that came much earlier. The technique AI models use to do that is called Attention. In this chapter we'll explain how it works.

Some Tokens Need Other Tokens

You often can't tell what a token means on its own. Take "bank": is it a place for money, or the side of a river? You have to look at the tokens around it. A token like "it" is even harder: it can stand for almost anything, so you have to look back to find what it points to.

Different tokens need different things in order for us to understand what they mean. "it" needs the thing it stands for. "opened" needs whoever did the opening. "sky" needs the planet it hangs over. "treaty" needs a name from pages ago.

Searching for a Match

So how do we know which tokens we need to "pay attention to" in order to understand the meaning of a particular token?

The "Attention" algorithm does this by giving each token three things:

• A token's query is what it is looking for, like the words you type into a search box.
• A token's key is what it advertises to the other tokens, like the title of a web page.
• A token's value is the information it provides if you pay attention to it, like the contents of the page.

To understand a token, we take that token's query and compare it against every other token's key to find the tokens whose keys seem most similar to that query.

This might seem a bit confusing, but it makes more sense once you play with the playground below:

Matching with the Dot Product

How does a computer decide whether a query matches a key? So far we've written queries and keys in plain English, but comparing English text is tricky. How do we know that "a place with a sky" matches "a planet"?

In Chapter 4 and Chapter 5 we showed that AI models can represent almost anything using vectors. We also showed that we can use the dot product to compare two vectors to see how similar their meanings are.

Indeed this is how Attention compares a query to a key. The query and key are both vectors, and their dot product gives a single match score we can read as a similarity from 0 (nothing in common) to 1 (a perfect match).

From here on we'll write each query and key in plain English and the match as a single score. Underneath they're still vectors and dot products.

Dividing Your Attention

Usually more than one token matches, at least a little. So instead of crowning a single winner, attention blends everything it looked at, giving each token a say in proportion to how well it matched. To do this, we turn the match scores into percentages that add up to 100%.

Most AI models do this with a function called softmax. Softmax always makes the percentages add up to 100%, and gives a bigger share to bigger scores.

The scores we gave in the previous section range from 0 to 1, but that range doesn't work well with softmax. For softmax to treat one match as clearly better than the rest, the scores need to be much bigger, like 9 instead of 1. So from here on we'll use larger match scores. A real model gets them by giving its key and query vectors larger magnitudes, which stretches out their dot products.

What if nothing matches?

We need one more trick to make softmax work nicely. Softmax requires its percentages to add up to 100%, but what do we do if nothing matches and we don't want to pay attention to any previous tokens? The fix is a sink: an extra token with a small fixed score. When the real tokens all score low, the sink out-competes them and quietly absorbs the attention. When something genuinely matches, it easily beats the sink.

Real models don't always add a dedicated sink token. Often they just learn to dump spare attention on a comma, the first token, or the token itself.

What Did You Find?

The reason we pay attention to another token is to extract information from it. That information is the token's value, which is yet another vector.

Attention combines the values of all the tokens, weighting each by its softmax percentage.

It's kind of crazy that this works. Why should you be able to blend these value vectors together and get something that makes any sense? It works because of something we saw in Chapter 5: the vectors AI models use are built so that if you average two of them, the result combines their meanings.

Where the Vectors Come From

Where do a token's query, key, and value actually come from? You don't need to write them by hand. The model uses a neural network to compute them from the token's embedding, and we use optimization to train the neural network to produce the queries, keys, and values that do the best job at predicting the next token in the training data.

Multi-Headed Attention

In the examples we've given so far, each token only has a single query. But in practice, there are often lots of different kinds of information we might want in order to understand the meaning of a token. Where is this happening? Who did this action? Which of several meanings of this word might be relevant?

So in practice we usually do multi-headed attention. We have multiple attention heads, each of which computes a different query, key, and value for each token. We also usually have multiple layers of attention, where each layer adds more information on top of what was gathered by the previous layers.

You'll see this in more depth when we discuss Transformers in Chapter 9.

What We've Built

Attention is:

1. Query, key, value. Each token looks for something (query), advertises what it has (key), and offers an answer (value), like a search.
2. Match scores. The dot product of a query and a key scores how well two tokens fit.
3. Softmax. Turn those scores into percentages that add up to 100%.
4. The sink. A safe place for attention to go when nothing in the sentence fits.
5. Blended values. Combine the chosen values using those percentages.
6. Many heads. Several heads run side by side, each looking for something different.

One thing is still missing: attention has no idea where the tokens are, only what they are. "The dog chased the cat" and "the cat chased the dog" give exactly the same scores. The next chapter fixes that with a trick involving rotation.