Attention: self-attention from scratch

A word’s meaning depends on the words around it. Attention is the part of the machine that lets every word look at its neighbours, decide which of them matter, and take in a little of what they say. It is built from parts you already own: a mix, a dot product and softmax.

By the end of this chapter, you will be able to

Say why one fixed place for each word is not enough.
Change what a word means by mixing in some of its neighbours.
Explain why the amounts in that mix cannot be fixed knobs, and must be worked out fresh for every sentence.
Name the three things each word makes, a question, a label and a message, and say what each is for.
Work through attention for one word in five stages: ask, score, share, mix, add.

One word, two meanings

Read these two sentences.

the bat flew away

the bat was swung

You did not stop to wonder what “bat” meant. In the first it is a small animal. In the second it is a piece of wood. You knew at once, because of the words around it.

The machine of chapter 6 cannot do that. It gives “bat” one place on the map, the same in every sentence. The best it can do is put “bat” halfway between the animals and the sports gear, which is right for neither.

So the place of a word has to be able to move, depending on its sentence. And you already have a tool for moving a place towards other places: the mix from chapter 4. Take some of this and some of that, in set amounts, and add them up.

Try it by hand first. You decide how much “bat” should listen to each word.

Try itThe orange dot is what “bat” means in this sentence. Right now it only listens to itself, so it sits in the middle. Raise the amount for another word to mix some of that word in. Make “bat” an animal. Then switch to the second sentence, without touching the amounts.

sentence

how much “bat” listensshare of the mix

1. the00%

2. bat10100%

3. flew00%

4. away00%

“Bat” sits at (0.0, 2.0). It could still be the animal or the thing you hit a ball with. Which word in the sentence gives it away?

Mixing worked. A share of “flew” pulled “bat” to the animal side. A share of “swung” pulled it to the sport side. A word can pick up meaning from its neighbours by mixing a little of them into itself.

But the second sentence showed the catch. You had set the amounts by position: so much for word 3, so much for word 4. In a new sentence the helpful word was somewhere else, and the same amounts pointed at the wrong word.

In chapter 4 the amounts in a mix were knobs: learned once, then fixed. That cannot work here. There are endless sentences, and the helpful word can be anywhere.

The amounts have to be worked out fresh for each sentence, from the words that are in it. The machine has to decide for itself what to pay attention to.

The first idea: listen to words like yourself

Chapter 7 gave us a way to get shares out of lists of numbers. Score with a dot product. Turn the scores into shares with softmax. So here is a first try. Let “bat” score its own place against the place of every word in the sentence, and use the shares as the amounts.

Try itHere each word’s place is scored against every other word’s place, and the scores go through softmax. Tap a word to see who it would listen to. Try both sentences.

sentence

who is asking

place of “bat” · place of this wordshare

the0.0 × 0.0 + 2.0 × 0.2 = 0.43%

bat0.0 × 0.0 + 2.0 × 2.0 = 4.096%

flew0.0 × −2.0 + 2.0 × −1.5 = −3.00%

away0.0 × −0.5 + 2.0 × −0.4 = −0.81%

“bat” gives 96% to itself. Of course: nothing is more like a word than the same word. It learns nothing new from the sentence.

It fails, and the reason is worth a moment. The dot product is big when two arrows point the same way. No arrow points more the same way as “bat” than “bat” itself. So a word that looks for words like itself finds itself.

Think about what “bat” really needs. Not a word that is like it. It needs a word that tells it something: an action word, saying what is being done. “Flew” is nothing like “bat”. That is exactly why it is useful.

What a word is, and what a word is looking for, are two different things. One list of numbers cannot do both jobs.

A question and a label

Picture a room full of people, each wearing a name badge that says what they do: plumber, doctor, baker. You have a leaking tap. You do not look for people like you. You call out what you need, and look for the badge that matches.

Give every word the same two things.

A question: a list of numbers that says what this word is looking for.

A label: a list of numbers that says what this word has to offer, like a badge.

The score is now the dot product of one word’s question with another word’s label. A big score means “you are what I was looking for”.

Below is a map of labels. Notice that it is laid out differently from the map of meaning. On the map of meaning, “flew” and “swung” were far apart, one with the animals and one with the sports. Their labels say something else: what kind of word each one is.

Try itEach blue dot is a word’s label: what the word says about itself. The orange arrow is the question “bat” asks. Drag its tip until “bat” listens mostly to “flew” in the first sentence and to “swung” in the second. One arrow has to work for both.

“the bat flew away”

the25%

bat25%

flew25%

away25%

“bat” reads as: the animal

“the bat was swung”

the25%

bat25%

was25%

swung25%

“bat” reads as: could be either

The question has no length yet, so every word gets the same share. Drag the orange tip.

You found one arrow that works in both sentences. It does not point at “flew”. It does not point at “swung”. It points at what their labels share: both are action words. The question means “which word here says what is being done?”, and each sentence answers it with a different word.

That is the trick that position could not do. The question is fixed, but the answer depends on the sentence.

The message

One piece is left. Once “bat” has chosen to listen to “flew”, what does “flew” hand over?

Not its label. The label only says “I am an action word”. That was for matching, and its job is done.

The next guess is that “flew” hands over its whole place, as in the very first figure of this chapter. Look back at what happened there. “Bat” moved to the animal side, which is what we wanted. But it also sank halfway down the map, from 2.0 to 0.3 in its second number. That second number says whether a word is a thing or an action. By swallowing all of “flew”, “bat” was turning into an action word. It should stay a thing.

So a word should not pass on everything about itself. It should pass on the part that is useful to a listener. For “flew”, the useful part is “this is about animals”. The fact that “flew” is an action is its own business.

That useful part is the word’s message. And it is not a new thing to store. Like the question and the label, it is worked out from the word’s place.

Try itEvery word starts with one list: its place on the map of meaning. Its question, label and message are all made from that place. Tap each word and read down. Watch the message most closely.

word

placewhere the word sits on the map of meaningfirst number −2.0: on the animal side
second number −1.5: an action(−2.0, −1.5)

questionwhat am I looking for?(0.8 × −1.5, −0.8 × −1.5)An action asks: which word here is a thing?(−1.2, 1.2)

labelwhat kind of word am I?(0.2 × −2.0 + −1 × −1.5, 1 × −1.5)The badge says: I am an action.(1.1, −1.5)

messagewhat do I pass on?(1 × −2.0, 0 × −1.5)The message says: this is about animals.(−2.0, 0.0)

The message of “flew” is (−2.0, 0.0). It keeps the first number of the place, which says “animal side”. It drops the second, which says “I am an action”. A listener learns which side it is on, and is not turned into an action itself.

Read the four rows as one story. A word has a place. From the place, three small weighted sums make three lists. The orange numbers are their knobs.

The question and the label are both built from the second number of the place, the one that says thing or action. That is why they are good for finding the right kind of word.

The message is built from the first number, the one that says animal or sport. Its knobs are 1 and 0: keep the first number, drop the second. The message is the word’s place with everything taken out except what a listener can use.

I set these knobs by hand so that you can read them. In a real machine they are learned, and we come back to that below.

Now every part is in place. Here is the whole thing for the word “bat”, one stage at a time.

Try itThis is everything attention does for one word. Press Next to go through the five stages. Then switch sentence and go through them again.

sentence

stage

question of “bat”(1.6, −1.6)thebatflewaway

label(−0.2, 0.2)(−2.0, 2.0)(1.1, −1.5)(0.3, −0.4)

stage 1 of 5

“Bat” makes its question. Every word in the sentence, “bat” included, holds up its label.

Say the five stages once more, because this is the heart of the whole article.

Ask. The word makes a question. Every word holds up a label.
Score. Question times label, pair by pair, added up. Chapter 7’s dot product.
Share. The scores become shares of 100%. Chapter 7’s softmax.
Mix. The messages are mixed in those shares. Chapter 4’s weighted sum.
Add. The mix is added to the word’s own place. The word keeps who it was, and gains what it heard.

Attention is a mix where the words themselves work out the amounts.

Who sets those knobs?

You saw that the question, the label and the message are each made from a word’s place by a few weighted sums.

So there are three small sets of knobs: one that turns a place into a question, one that turns it into a label, one that turns it into a message. The same three sets are used for every word in every sentence. Nothing is stored for any one sentence. That is how a machine with a fixed number of knobs can handle sentences it has never seen.

In the figures, I set those knobs by hand so that things ask about actions, and actions ask about things. In a real machine nobody sets them. They are learned with the loop from chapter 3 and the blame from chapter 5, exactly as before. Each stage is made of multiplying and adding, and softmax has a slope, so blame can flow backwards through all five stages, into the three sets of knobs, and on into the places of the words.

Nobody tells the machine that “bat” should ask about actions. It finds questions that help it do its job, which is the subject of the last chapter.

In symbols

A set of knobs that turns one list into another is drawn as a grid, called a matrix. Each row of the grid is one weighted sum.

W_{Q} e = (00 0.8 - 0.8) (02) = (0 \times 0 + 0.8 \times 2 0 \times 0 + - 0.8 \times 2) = (1.6 - 1.6)

The question of “bat” in the figures. The place of “bat” is $(0, 2)$ . Each row of the grid makes one number of the question.

With that way of writing, the three lists for word number $i$ are:

q_{i} = W_{Q} e_{i} k_{i} = W_{K} e_{i} v_{i} = W_{V} e_{i}

Query, key and value. Three grids of knobs, each applied to the same place $e_{i}$ .

Then the five stages, for the word that is asking, $i$ , and each word it looks at, $j$ :

s_{ij} = q_{i} \cdot k_{j}

Score: the query of word $i$ with the key of word $j$ .

p_{ij} = \frac{e ^{s_{ij}}}{\sum _{m} e ^{s_{im}}}

Share: softmax over all the words in the sentence. Here $e$ is the number 2.718, not a place.

new e_{i} = e_{i} + j \sum p_{ij} v_{j}

Mix and add: the values, mixed by their shares, added to where the word started.

$e_{i}$: the place (embedding) of word $i$ .
$W_{Q}, W_{K}, W_{V}$: the three grids of knobs. They are the only knobs in attention.
$s_{ij}$: how well word $j$ answers the question of word $i$ .
$p_{ij}$: the attention weight: how much word $i$ listens to word $j$ .

In papers you will see all of this packed into one line:

softmax (\frac{Q K ^{⊤}}{d}) V

The same five stages, done for every word at once. $Q$ , $K$ and $V$ are the queries, keys and values of all the words, stacked up.

The only new piece is $d$ : the scores are divided by the square root of how many numbers are in each list. With long lists the scores get large, and chapter 7 showed that large scores make the split too sharp. Dividing keeps it soft.

Every word asks at once

We followed “bat”. But nothing about the five stages is special to “bat”. Every word in the sentence makes its own question, at the same moment, and gets its own mix.

Try itThe top row is each word asking. The bottom row is each word being listened to. A thick line means a big share. Tap the words in the top row, one by one.

sentence

“flew” gives 98% of its attention to “bat”. An action asks: who or what is doing this?

Each word came in as a fixed pin on the map. Each word goes out moved a little, towards whatever it found useful in the others. “Bat” now knows it is an animal. “Flew” now knows what was doing the flying.

Words go in alone and come out knowing about each other. That is what attention is for.

What is still missing

We have the main part. Three things are still missing before it becomes a machine that can write.

Order. Look back at the five stages. Nowhere did we use where a word is in the sentence. To this machine, “dog bites man” and “man bites dog” are the same bag of words.

Depth. One round of attention lets each word hear from the others once. Real meaning needs more than one round, and it needs the bends of chapter 5 as well.

A job. All these knobs need something to be wrong at, or they cannot learn. Chapter 6 gave a hint: guess the missing word.

The last chapter puts these in, and then lets the whole thing run.

You can now

Say why one fixed place for each word is not enough.
Change what a word means by mixing in some of its neighbours.
Explain why the amounts in that mix cannot be fixed knobs, and must be worked out fresh for every sentence.
Name the three things each word makes, a question, a label and a message, and say what each is for.
Work through attention for one word in five stages: ask, score, share, mix, add.