How Attention Works

Read this sentence: the animal didn't cross the street because it was too tired. You knew at once that it is the animal. Now change one word: the animal didn't cross the street because it was too wide. Now it is the street. Nothing about the word it changed, and the word that settles the question sits four words further on. Puzzles like this oneSentences built so that one word flips what a pronoun refers to are called Winograd schemas, after the computer scientist Terry Winograd, who used one in 1972. A set of them was proposed as a test of machine understanding in 2012 . were long a test of whether a program understands language at all.

Modern language models handle them through a mechanism called attention. This page builds attention up from nothing, on models small enough that every number can be shown. The first one is a single attention head with six numbers per word, which I set by hand so it would do this one job. Switch the last word between tired and wide, and watch where it looks:

With tired, it puts 75% of its attention on animal; with wide, 75% goes to street. The rest is spread thinly over the other words. Each word in a real model does this kind of looking around, in dozens of heads at once and in every layer, and what each head finds gets mixed into the word's numbers before the next layer reads them. By the top of the model, the numbers for it carry a good deal of animal or street, and that is how the model "knows" which one was meant.

The rest of this page takes the figure apart. We'll need four ideas: that words are points in a space, that a dot product measures how much two points agree, that a query and a key decide who looks at whom, and that a value decides what gets passed along. Then we'll put many heads together, teach one to look at the right place from examples, and see why the whole scheme gets expensive as texts get long.

Words as points

A computer can't do arithmetic on the word cat. So the first thing a language model does is look each word up in a table and replace it with a list of numbers, a vector. In a large model each list is a few thousand numbers long, and the table is learned along with everything else. The useful property is that words used in similar ways end up with similar lists.

We can see where that property comes from without training anything. Below, the page has counted, in a little corpus of 27 sentences about animals and streetsThe whole corpus is in the page's source, in the CORPUS string. It is tiny on purpose: real models learn from trillions of words., which words appear just before and just after each of 16 words. The counts are weighted by how surprising they areThe weighting is pointwise mutual information: how much more often two words appear together than they would by chance . Negative values are set to zero., giving each word a list of 148 numbers, and then squeezed down to the two directions along which the words differ most, so every word becomes a point on a page. Drag the probe around the map:

Three groups appear without anyone telling the program what the words mean. The animals sit together at the bottom left, the places on the right, and the describing words at the top. They got there only because of the company they keep: cat and dog both follow the and come before slept or ate, while tired and wide both follow was and come before and.

Two directions can't hold much, though. On this map tired and wide are neighbours, because in the corpus they are used in the same grammatical slot, even though one describes animals and the other describes places. Squeezing 148 numbers into two threw away the difference. In the full list of 148, it's still there. Pick a word and compare its full list with every other word's:

In the full lists, tired is closest to sleepy, and wide to busy. Cat is closest to horse and dog; street to river, road, and park. A few neighbours are accidents of the tiny corpus: bird comes second for tired, because in 27 sentences the bird happens to turn up in the same slots as describing words, and more text would wash that out. More numbers keep more distinctions. That is why the vectors in large models are thousands of numbers long, and why they are learned rather than counted: training finds whatever directions help predict text, including ones no linguist would think to name.Counting and squeezing, as here, is close to how word vectors were made in the 1990s and 2000s. Learned vectors became standard in the 2010s.

Agreement

We'll need one tool over and over: a single number that says how much two vectors agree. It's called the dot product, and it's as simple as arithmetic gets: multiply the two lists entry by entry and add up the results. In two dimensions you can watch what that number means. Drag the tips of the orange and blue arrows:

The arithmetic is the first line of the readout. The second line says the same thing geometrically: the dot product is the length of one arrow, times the length of the other, times the cosine of the angle between them. That's the length of the orange arrow's shadow on the blue one, times the blue arrow's length. Turn the arrows to right angles and the shadow vanishes: the dot product is zero, the two have nothing in common. Point them apart and it goes negative.

The same recipe works for lists of any length, where there is no picture to draw:

Two lists of 4,000 numbers that agree, in the sense that their big entries tend to be in the same places with the same signs, have a large dot product. Two that have nothing to do with each other get a small one, positive and negative products cancelling out. The similarities in were dot products too, of lists first scaled to length 1, which turns the dot product into the cosine of the angle.

Queries and keys

Back to the sentence. For it to find animal, two things have to happen. It has to say what kind of word it's looking for, and every other word has to say what kind of word it is. Attention gives each word two new vectors for these jobs: a query, shown in orange, and a key, in blue. The score between two words is the dot product of one word's query with the other word's key: large when what one looks for matches what the other offers.

In our hand-built head, queries and keys have just two numbers each, so they can be drawn. The key of animal points right; the key of street points up; the other words have short keys near the middle, because they have little to offer this head. Drag the query around and watch the scores it gives each word:

Point the query right and animal scores highest; point it up and street does. A longer query makes every score bigger, the good matches much bigger. Point it away from both and the scores go negative: a negative score is a word the query actively doesn't want.

Where do queries and keys come from? Each is the word's vector multiplied by a matrix, one matrix for queries and another for keys. The matrices are the head's learned parameters, its whole personality. Our hand-built head describes each word with six numbers: how much it is an animal, how much a place, how much it describes animals, how much it describes places, and two leftover numbers for everything else. The query matrix maps describes animals to point right and describes places to point up. One more ingredient: by the time this head runs, an earlier layer has already copied some of the describing word at the end of the sentence into it.That earlier step is itself attention, in an earlier layer. Copying information about one word into another's vector is exactly what attention does, as we'll see in the section on values. Change how much was copied, and how strong the query matrix is:

Each number of the query is a dot product too: the description of it against one column of the matrix. With nothing copied in, the description of it has no describing-word part, so its query is zero, and it scores every word the same. Copy tired in and the third number lights up, the matrix turns it into "point right", and the query lands on animal. Make the matrix bigger and the query gets longer, and the attention more decisive: at 8, animal gets 96%. Set it to 1 and animal gets only 17%, against about 9% for every other word.

A trained model works the same way, with numbers no one chose. Its word vectors have thousands of entries, its query and key matrices map them to vectors of, typically, 64 to 128 numbers, and the "directions" that queries and keys use are whatever helped predict text during training. Some turn out to mean something a person can name; many don't.

From scores to weights

Scores can be any size and any sign, but what we want is a recipe: how much of each word to take. That means numbers that are never negative and add up to one. The standard way to get them is called softmax, and it works in two steps. First, raise e \approx 2.718 to the power of each score, which makes every number positive and stretches the gaps between them. Then divide each result by their total. Drag the scores up and down and follow them through both steps:

The exponential is what makes softmax decisive. A word whose score is 1 higher than another's gets e \approx 2.7 times the weight; 3 higher, about 20 times. Pull the animal bar down to the level of the others and its weight collapses to an equal share; push any word's score up by a few points and it takes nearly everything. The weights always add up to one, so attention is a budget: more for one word means less for the rest.

Dividing every score by a temperature before the exponential controls how decisive it is. Turn it down and the best score takes almost all the weight; turn it up and the weights flatten toward an even split. The same knob appears when a language model picks its next word, where it sets how adventurous the choice is.

Here is that effect, measured. For each width d, the page draws 400 random queries and compares each against 8 random keys, with every entry a random number from a bell curve. Change the width:

At width 2, the typical score is about 1 away from zero, and the largest of 8 weights averages 0.41. At width 512 the typical score is about 20 away from zero, and without the division the largest weight averages 0.95: attention has made up its mind before learning anything. Divided by \sqrt{d}, the scores have the same spread at every width, and the largest weight stays near 0.36. This is why the attention formula carries a \sqrt{d}, and why our two-number head divided by \sqrt 2 all along.The original paper gives exactly this reason: for large widths, unscaled dot products grow large and push the softmax into regions where its gradients are tiny, which makes learning slow .

What gets passed along

Weights say how much to take from each word. The last question is what to take. Each word gets a third vector, its value, shown in green, made by a third matrix. The output of the head for a word is the weighted sum of every word's value, using that word's attention weights.

In our head, values have two numbers: how much the word is an animal, and how much it is a place. Build the output for it one word at a time, largest weight first:

The first term does almost all the work: with tired, animal has weight 0.75 and value (1, 0), so after one word the output is already (0.75, 0), and all eleven together give (0.76, 0.02). The other ten words add small corrections that mostly cancel. Switch to wide and the output swings to point at place. This output is added to the vector of it, and from then on, everything in the model that reads the word it reads a little "animal" or a little "place" along with it.

Notice that the output is always an average of the values, so it can only be something the values already say: the red arrow can never leave the region the green arrows span. Attention moves information between words. It doesn't invent any.The parts of a transformer that compute new things from what a word has gathered are the feed-forward layers, which work on each word separately. We'll meet them briefly at the end.

Every word at once

So far we've followed one word. In a real model every word does the same thing at the same time: every word has a query, scores every word's key, and gathers its own weighted sum of values. Put all the weights in a grid, one row per word doing the looking, one column per word being looked at, and you get the attention matrix. Point at a row, or step through them:

Two rows stand out. The row for it is the one we've been studying. The row for tired is sharp too, and points at the same word: tired describes animals, so its query looks for an animal, and it finds one. Every other row is nearly flat. Those words have nothing to ask this head, so they spread their attention evenly, and the even average of everyone's values they get back is close to zero. A head that has nothing to say for a word mostly stays out of its way.Trained models often park this leftover attention on the very first word of the text instead of spreading it, which serves the same purpose: the first word's value can be learned to be nearly empty.

All of this is a few matrix multiplications. Stack every word's vector as a row of a matrix X; then \orange{Q} = X\orange{W_Q}, \blue{K} = X\blue{W_K}, and \green{V} = X\green{W_V} hold every query, key, and value at once, and:

Read from the inside out: \orange{Q}\blue{K}^{\top} is every query against every key, the grid of scores; dividing by \sqrt d keeps them in range; softmax, applied to each row, turns them into the matrix of weights above; and multiplying by \green{V} takes each row's weighted sum of values. This one line is the heart of every transformer, the kind of model behind modern chatbots .

No peeking ahead

There's a catch for models that write. A language model produces text one word at a time, each time predicting the next word from the words so far. When it's working on it, the word tired doesn't exist yet. So in such a model each word may only attend to itself and the words before it. That rule is called a causal mask: every score for a later word is set to minus infinity before the softmax, which makes its weight exactly zero.

Watch the matrix as the sentence is written word by word. With the mask on, each new word adds one row at the bottom and one column on the right, and nothing above it changes. Turn it off and see what happens to the rows already written:

Without the mask, adding a word changes the rows of words already written, because every one of them could see the new word. A model that has to rewrite everything it computed whenever a word arrives does a great deal of repeated work. With the mask, each word's row is fixed the moment the word is written, so a model can keep every past word's key and value in memory, a store usually called the key–value cache, and for each new word compute only one new row.

The mask also changes where the work happens. Step to word 8, it, with the mask on: it can't see tired, nothing has been copied into it, and its row is flat. The question of what it refers to can only be settled later. Step to word 11: tired itself looks back and finds animal. In a model that writes left to right, information flows only forward, and later words do the gathering.

Many heads

One head computes one kind of relationship. A real model runs many side by side, each with its own three matrices, and adds what they find.Strictly, the outputs of all the heads in a layer are placed side by side and multiplied by one more matrix, which is the same as letting each head add its own projected output. Trained heads often turn out to do recognizable jobs. Here are three that show up again and again in real models, built by hand on a sentence with some repeated words. Pick a head and a word:

The previous-word head is the simplest: every word looks one step back. It knows nothing about meaning, only position, and it's useful because it lets each word carry a note of what came just before it. The same-word head finds earlier copies of the current word: the second cat looks at the first. When a word has no earlier copy, it rests on the first word, which contributes almost nothing.The chips show spaces as dots. Like real tokenizers, this sentence treats the at the start and the with a space before it as different tokens, so the first the has no exact copy later on.

The third, the induction head, is the clever one. At the second cat it doesn't look at the first cat; it looks at the word right after the first cat, sat. It can do this because its keys don't come from each word itself but from the word before it, which a previous-word head in an earlier layer has already copied in. Its query asks "where have I been before?" and the answer it finds is "just before sat". If a model is guessing the next word, that's a strong hint: last time the cat appeared, sat came next. Two simple heads in two layers, working together, let a model continue patterns it has seen only once in the text so far. Researchers found exactly these induction heads in trained transformers and argued that they account for much of a model's ability to learn from examples given in its prompt .The pairing of a previous-word head with an induction head was first described in a study of small attention-only models .

Where words are

Look back at how scores are computed: a query against a key, nothing else. Nothing in that recipe knows the order of the words. If a head reads the whole sentence at once, without a mask, then shuffling the words just shuffles the rows and columns of its matrix; each word gathers exactly the same mixture as before. Try it on a sentence where order is the whole point. This head, reading both ways, looks for the nouns around a verb:

Without positions, bites splits its attention evenly between dog and man in both sentences, and gets back the same mixture. As far as this head is concerned, the two sentences are identical. Turn positions on and give the head a reason to care about them, a preference for the word right before the verb, and now dog bites man and man bites dog come out different.

So transformers add information about position to each word's vector before the first layer. The original design used a fixed pattern of waves: for position p, a list of sines and cosines of p at many different frequencies . Each row below is one position's list of 16 numbers; low-numbered columns wave quickly from row to row, high-numbered ones slowly. Pick a position and compare its list with every other position's:

A position is most like itself (a dot product of 8: each sine–cosine pair contributes \sin^2 + \cos^2 = 1), nearly as like its neighbours (about 7.5), and less like positions far away, with ripples on the way down. Because the waves are sines and cosines, a shift by a fixed number of positions acts the same way everywhere in the text, which is the kind of regularity a query and key matrix can learn to use: "look three words back" becomes one fixed transformation.Many current models use other schemes, such as rotating each query and key by an angle that depends on its position, so the score between two words depends only on how far apart they are. The idea is the same: give attention a way to see order.

Learning to look

Every head so far was built by hand. In a real model nobody sets the matrices; they start as random numbers and are adjusted, a tiny step at a time, to make the model's predictions better. Let's watch that happen to one head on a task small enough to follow.

The task is a lookup. Each example holds four pairs of letters, a key and a value, like C→F, and then a question: one of the keys, say C?. The right answer is the value stored under it, F. Picture an earlier layer having already merged each pair into one position. The head has a query matrix and a key matrix, each turning 16 numbers (8 for the key letter, 8 for the value letter) into 4, so 128 numbers in all, all random to start. Its values are fixed to read out the value letter, so the only way to answer is to put the weight on the right pair. After each batch of 16 fresh examples, every one of the 128 numbers is nudged in the direction that would have made the right answers a little more likely.The nudge is gradient descent with momentum. The page computes the gradient exactly, by working backward through the softmax and the two matrix products by hand; the code is in the page source, in trainStep. Press play:

At first the weights are close to even, the head is guessing, and the loss sits near \ln 4 \approx 1.39, what you get from spreading your bet evenly over four pairs. Within a hundred or so steps the loss falls toward zero, and the weights on the example pile onto the pair whose key matches. The grid on the bottom right shows what the head learned: a bright red diagonal. Question letter A matches key letter A, B matches B, and every mismatch is pushed negative, blue. Nobody told the head to compare letters. It discovered that comparing them was the only way to lower the loss.

Try the step sizes. A small step learns slowly but steadily; a large one gets there fast here, though in big models too large a step makes training unstable. Reset starts again from the same random numbers, so runs with different step sizes are easy to compare.

Stacking it up

A transformer is this mechanism repeated. Each word keeps a running vector, often called the residual stream. Every layer reads from the streams and writes back into them, adding rather than replacing: first attention, whose heads move information between words, then a feed-forward network that transforms each word's vector on its own. Step through the layers:

Because each layer adds to the stream rather than overwriting it, information written by an early head is still there for later ones to read. That's how the induction head above could use what the previous-word head had copied. The last word's stream, at the top, is turned into a guess at the next word. Large models stack dozens of these layers, with dozens of heads in each.For example, the largest model described in the GPT-3 paper has 96 layers with 96 heads each, and vectors of 12,288 numbers .

The price of looking everywhere

Attention has one expensive habit. Every word scores every word before it, so a text twice as long needs about four times as many scores. Slide the length of the text:

A thousand words need about half a million scores per head per layer; a hundred thousand words need five billion, which would take 20 gigabytes to store at four bytes each, for one head in one layer. This square law is why models have a limit on how much text they can read at once, and why so much engineering goes into attention. One widely used method computes exactly the same result but never stores the whole grid, working through it in tiles small enough to stay in the fastest memory of the chip . Others change the recipe itself, letting each word look at only some of the others.

What attention is, and isn't

Go back to the sentence about the animal and the street. What settled the meaning of it was not understanding in any mysterious sense but a chain of ordinary arithmetic. Each word was turned into numbers that summarize how it's used. A query matrix turned "describes animals" into a direction to look in. A key matrix made animal point that way. Softmax turned a large dot product into most of the attention, and the value of animal flowed into it. In a trained model, nobody designs these directions; they are whatever made the predictions better, spread across thousands of numbers, dozens of heads, and dozens of layers, which is a large part of why it's so hard to say exactly what a big model is doing.

Attention was first added to translation systems in 2014, to let a model look back at the right part of the sentence it was translating ; three years later, a model built from little else, the transformer, beat the best translation systems of the time . Nearly every large language model since has been a stack of the mechanism on this page: a query, a key, a value, and a weighted average.

References

1. Hector J. Levesque, Ernest Davis, and Leora Morgenstern, “The Winograd Schema Challenge”, Proceedings of the Thirteenth International Conference on Principles of Knowledge Representation and Reasoning (2012).
2. Kenneth Ward Church and Patrick Hanks, “Word association norms, mutual information, and lexicography”, Computational Linguistics 16(1) (1990), 22–29.
3. Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N. Gomez, Łukasz Kaiser, and Illia Polosukhin, “Attention Is All You Need”, Advances in Neural Information Processing Systems 30 (2017).
4. Catherine Olsson, Nelson Elhage, Neel Nanda, and others, “In-context Learning and Induction Heads”, Transformer Circuits Thread (2022).
5. Nelson Elhage, Neel Nanda, Catherine Olsson, and others, “A Mathematical Framework for Transformer Circuits”, Transformer Circuits Thread (2021).
6. Tom B. Brown and others, “Language Models are Few-Shot Learners”, Advances in Neural Information Processing Systems 33 (2020).
7. Tri Dao, Daniel Y. Fu, Stefano Ermon, Atri Rudra, and Christopher Ré, “FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness”, Advances in Neural Information Processing Systems 35 (2022).
8. Dzmitry Bahdanau, Kyunghyun Cho, and Yoshua Bengio, “Neural Machine Translation by Jointly Learning to Align and Translate”, International Conference on Learning Representations (2015).