Attention

A word needs its neighbours

What does it mean in The cat sat on the mat because it was tired? The cat, of course. Change the last word to soft and the answer flips to the mat. The meaning of a word depends on the other words around it, sometimes far away.

Attention is how a transformer lets every token look at the other tokens and pull in what it needs. On the left are the tokens that ask, on the right the tokens they look at. A thicker line means more attention.

The model on this page is a toy with hand-made weights, small enough to read every number. Real models learn their weights in training, but they compute attention with exactly these formulas.

Try it. Switch the last word between tired and soft and watch where the last token looks. Then click it.

Queries and keys

Every token gets three vectors, made by multiplying its current vector by three learned matrices. The query says what the token is looking for. The key says what the token has to offer. The value is what it will hand over if someone looks at it.

To decide where to look, a token takes the dot product of its query with the key of every token. A big product means the question and the offer match. The products are divided by the square root of the vector size, here the root of 4, so the numbers do not grow too big in large models.

Try it. Pick tired as the asking token. Its query wants something alive, and only cat has a key that says so.

Softmax turns scores into weights

The scores of one token go through softmax, the same function that turns a language model's scores into probabilities. The result is a set of attention weights that are all positive and add up to 1.

Softmax exaggerates differences. A score that is just 2 points higher gets more than 7 times the weight, so one or two tokens usually take most of the attention while the rest get crumbs.

Values are mixed into the token

Now every token collects the values of the others, each multiplied by its weight, and adds them up. The sum is the new vector of the token, now with context mixed in.

Look at it before and after. Going in, it is a bare pronoun. Coming out, it carries a share of cat and mat. For tired the biggest share comes from the cat, so the next layers know who is tired.

Try it. Compare the output of tired with the output of soft.

Every token at once

All tokens do this at the same time, and the weights form a square table with one row per asking token. This is why transformers train so fast on graphics cards. The whole table is one big matrix multiplication.

Models that write text, such as GPT and Llama, use a causal mask. A token may look only at itself and the tokens before it, never at the future, because during writing the future does not exist yet. Models that only read, such as BERT, look both ways.

Try it. Turn the mask off. The first tokens can now see the whole sentence, and sat finds its subject cat together with the mat that comes later.

Many heads look for different things

One attention pattern is not enough, so a layer runs many of them side by side. Each is called a head and has its own query, key and value matrices. Their outputs are glued together and mixed again.

Researchers who open real models find heads with clear jobs. Some look at the previous token. Some link a pronoun to its noun. Many park most of their attention on the first token when they have nothing useful to do, which is called an attention sink. Our three heads copy these three patterns.

Formula

softmax(Q · K^T / √d) · V

Introduced

Attention Is All You Need, 2017

Heads

12 per layer in GPT-2 small, 32 in Llama 3.1 8B, where every 4 query heads share one key and value head

Cost

Grows with the square of the context length