Attention: a soft lookup

Attention lets each token pull in information from other tokens. It works like a dictionary lookup that never fails: instead of returning the one entry whose key matches, it returns a weighted average of every entry, weighted by how well each key matches the question.

Key idea: Score every key against the query, softmax the scores into weights, return the weighted average of the values
Real world: An estate agent prices a 3-bedroom house nobody has sold before by looking at houses that have sold, and averaging their prices, giving the most similar houses the most say.
In an LLM: In “The cat sat on the mat because it was tired”, the word “it” has to find out what it refers to. A GPT-2 attention head gives 84% of “it”’s attention to “cat” and copies information about the cat to the position of “it”.

Why tokens need to talk

After the embedding step, every token is a vector, but each vector only knows about its own token. The vector for “it” is the same in every sentence: it has no idea whether “it” means a cat, a cup or the weather.

To predict the next word, the model needs context. Something has to move information from one position to another: from “cat” to “it”, from “France” to the position that will say “Paris”. is that step. It is the only place in a transformer where tokens exchange information; every other part works on each position separately.

A lookup that never fails

A normal dictionary (a Python dict, a phone book) is a hard lookup. You bring a key, it finds the entry with exactly that key, and returns its value. Ask for a key that isn’t there and you get nothing.

Attention is a soft lookup:

Compare the question, called the query, with every entry’s key and give each a score.
Turn the scores into weights that are positive and add up to 1 (with softmax).
Return the weighted average of all the values.

The answer is never “not found”. Close matches dominate, poor matches contribute a little. Play with the three kinds of lookup:

sharpness1.5

query: 3.00 bedrooms

keyvaluescoreweightw × v

1 bed£150k-6.001%1

2 bed£210k-1.5049%104

4 bed£340k-1.5049%168

5 bed£420k-6.001%2

answer for 3.00 bedrooms£275k

Every house gets a weight from how close its key is to the query; the weights add up to 100%. The answer, £275k, is the weighted average of the prices, 49% of it from the 2-bedroom house. The purple curve is smooth: a small change in the query gives a small change in the answer, which is what lets a network learn it.

Four known houses: bedrooms are the keys, prices the values. Press play to sweep the question (the query) from 1 to 5 bedrooms, or drag the slider. Switch between an exact lookup, the nearest match, and attention’s soft lookup.

Two things to notice. First, the soft answer changes smoothly as the query moves. A hard lookup jumps from one entry to the next, and you can’t learn by gradient descent through a jump; you can through a smooth curve. Second, the sharpness slider: when the scores are very spread out, softmax puts nearly all the weight on the best match and the soft lookup behaves like the nearest match. How spread out the scores are matters a lot, and the scaling lesson is about keeping that under control.

example · pricing a house

An agent has records of four sold houses: 1, 2, 4 and 5 bedrooms. A seller asks about a 3-bedroom house. A strict database query (WHERE bedrooms = 3) returns nothing. A sensible agent instead looks at all the records, trusts the 2- and 4-bedroom houses most, the others less, and quotes a weighted average. That is attention: query = the house asked about, keys = the bedrooms on file, values = the prices.

Queries, keys and values as vectors

In a transformer the query, the keys and the values are all vectors, and the score is a dot product: multiply the two vectors number by number and add up. A dot product is large when two vectors point the same way, near zero when they are at right angles, and negative when they point apart.

Each token makes three vectors from its embedding:

a : what am I looking for?
a : what do I contain?
a : what do I hand over if someone attends to me?

Below is a two-number version of our running sentence, “The cat sat on the”. The horizontal axis means is a thing, the vertical one is an action. The query belongs to the last “the”, which needs a noun.

query (blue) and keys: drag the query

values, and the output (blue) = their weighted mix

keyscore q·kweight (softmax)

The1.40×−0.9 + 0.30×−0.9 = −1.531%

cat1.40×1.7 + 0.30×0.2 = 2.4478%

sat1.40×0.2 + 0.30×1.7 = 0.7915%

on1.40×−0.9 + 0.30×0.7 = −1.052%

the1.40×−0.2 + 0.30×−1.2 = −0.644%

The query points most towards “cat”, so that dot product is the largest and softmax gives it 78% of the weight. The output (blue) is mostly cat’s value, with a little of the others mixed in. Drag towards another key and the output follows.

The query belongs to the last “the”. Drag it anywhere in the left plane. Each key’s arrow gets thicker as its softmax weight grows; on the right the output is built from the value vectors, each shrunk by its weight and laid tip to tail.

Drag the query towards “sat” and the weight moves to “sat”; the output moves towards sat’s value. Make the query short and all scores are near 0, so the weights are nearly equal and the output is a bland average. Make it long and the best score wins almost everything. The length of the query (and of the keys) acts like the sharpness slider above.

Notice that the key and the value of a token are different vectors. The key decides who gets attention; the value decides what information moves. A token can be easy to find for one reason and pass on something else entirely.

example · a library catalogue

You type “French cookery” into a library search (query). It is compared with every book’s subject labels (keys). Instead of handing you one book, imagine it hands you a summary built from all of them: mostly from the cookery books, a little from the French history books, almost nothing from the physics shelf. The labels decided how much each book counts; the books themselves (values) supplied the content.

The real thing: “it” finds “cat”

Here is a head from GPT-2 small, a real 124-million-parameter model, reading a real sentence. Nobody programmed it to resolve pronouns: its query and key matrices were learned during training because moving information this way helped predict the next word.

Loading GPT-2 data…

The output of attention for “it” is mostly cat’s value vector. That output is added to the vector at the position of “it”, so from this layer on, the model’s vector for “it” carries information about the cat. Later layers can use that: a cat can be tired, sleepy or hungry; a mat can’t.

example · “it” in GPT-2

At layer 4, head 3 (GPT-2 numbers its layers and heads from 0), the query of “it” matches the key of “cat” better than any other key, and softmax gives “cat” 84% of the weight and “mat” 4%. Averaged over all 12 heads of the layer, “it” sends only 10% to “cat” and 51% to the first token: each head does its own job, and the multi-head lesson shows why a layer has several.

What comes next

The next page makes this exact. The queries, keys and values come from three learned matrices; all the dot products are done at once as one matrix product; and a mask stops tokens from looking at words that haven’t been written yet.

going deeper

The soft lookup as a formula

For one query $q$ , keys $k_1, \dots, k_n$ and values $v_1, \dots, v_n$ :

\operatorname{attend}(q) = \sum_{j=1}^{n} a_j\, v_j, \qquad a_j = \frac{\exp(q \cdot k_j)}{\sum_{j'} \exp(q \cdot k_{j'})}

The weights $a_j$ are a softmax of the scores, so they are positive and sum to 1, and the output is a convex combination of the values: it always lies inside the region the values span. Transformers also divide the scores by $\sqrt{d}$ ; see scaling.

The house-price widget uses a negative squared distance, $-s\,(q - k_j)^2$ , as its score instead of a dot product. With that score, the soft lookup is the classic Nadaraya–Watson kernel regression from 1964 statistics; attention is the same idea with a learned dot-product similarity.

Hard versus soft

As the scores are multiplied by a growing factor, softmax tends to a one-hot vector on the largest score, and the soft lookup becomes the nearest-match lookup. Soft is differentiable everywhere; hard is not. That is the reason attention is soft.

In papers

Attention for translation, where a decoder softly looks up encoder states: Bahdanau et al., 2014.
The transformer, attention as the only way tokens communicate: Vaswani et al., 2017.
Attention as a differentiable key–value memory: Graves et al., Neural Turing Machines, 2014.