Attention is a lookup you can learn

Take the sentence "the river bank flooded". The word "bank" is ambiguous in isolation and completely unambiguous in context, and the thing that resolves it is one other word two positions to the left. A recurrent network has to have carried "river" forward in its state and not overwritten it. Attention does not carry anything forward. It goes and looks.

The mechanism is a soft dictionary lookup. Every position publishes a key describing what it offers and a value holding what it would contribute. Every position also issues a query describing what it wants. Instead of matching one key exactly, the query scores all of them, the scores become weights that sum to one, and the answer is the weighted mixture of every value. Nothing is retrieved exactly and nothing is missed entirely.

Below is that computation, in full, on four words and two dimensions.

The setup

Give each word a two-number vector. Call the first dimension water and the second dimension money. These labels are a teaching convenience: real learned dimensions are not individually interpretable, and the point here is the arithmetic, not the semantics.

• the: (0.1, 0.1)
• river: (1.0, 0.0)
• bank: (0.5, 0.5)
• flooded: (0.8, 0.2)

"Bank" sits exactly halfway between the two senses, which is the situation to be resolved.

Three learned matrices turn each vector into a query, a key and a value. Use these, which are made up but fixed:

• Query: q = (x1 + 0.5 x2, 0.5 x1 + x2)
• Key: k = (2 x1, x2)
• Value: v = x, unchanged

Applying them gives the keys straight away.

• k for the: (0.2, 0.1)
• k for river: (2.0, 0.0)
• k for bank: (1.0, 0.5)
• k for flooded: (1.6, 0.2)

And the query for "bank" is (0.5 plus 0.25, 0.25 plus 0.5), which is (0.75, 0.75).

Scoring

Score each key against that query with a dot product: multiply matching components and add.

• against the: 0.75 times 0.2 plus 0.75 times 0.1, which is 0.15 plus 0.075, or 0.225
• against river: 0.75 times 2.0 plus 0.75 times 0.0, which is 1.500
• against bank: 0.75 times 1.0 plus 0.75 times 0.5, which is 1.125
• against flooded: 0.75 times 1.6 plus 0.75 times 0.2, which is 1.350

A dot product is a similarity, in the same sense used in Dimensionality reduction without the hand-waving: it is large when two vectors point the same way and large in magnitude.

Divide each score by the square root of the key width, which is the square root of 2, or 1.414. The scaled scores are 0.159, 1.061, 0.795 and 0.955.

The division is not cosmetic. Vaswani and colleagues state the reason: for large values of the key width, the dot products grow large in magnitude, pushing the softmax function into regions where it has extremely small gradients. Their formula is attention equals the softmax of Q times K transposed, divided by the square root of the key width, all times V.

Softmax, and why the scaling was necessary

Exponentiate each scaled score and divide by the total.

• exp of 0.159 is 1.172
• exp of 1.061 is 2.889
• exp of 0.795 is 2.214
• exp of 0.955 is 2.599

Those sum to 8.875. Dividing gives the attention weights for "bank":

• the: 0.132
• river: 0.326
• bank: 0.250
• flooded: 0.293

They sum to 1.000, which is the property that makes this a mixture rather than a score.

Now see what happens without the division. Suppose the key width were 64 instead of 2, and the raw dot products were correspondingly larger, say eight times what they are here: 1.8, 12.0, 9.0 and 10.8. Run the same softmax and the weights become 0.00003, 0.740, 0.037 and 0.223. The distribution has collapsed onto one key. A softmax that has collapsed is nearly flat in its gradient, so the model can no longer learn to shift attention, which is exactly the saturation problem described in Why activation functions matter, arriving through a different door.

The answer

Mix the values with those weights. The values are the original vectors.

First dimension: 0.132 times 0.1, plus 0.326 times 1.0, plus 0.250 times 0.5, plus 0.293 times 0.8. That is 0.013 plus 0.326 plus 0.125 plus 0.234, which is 0.698.

Second dimension: 0.013 plus 0 plus 0.125 plus 0.059, which is 0.197.

The word "bank" arrived as (0.5, 0.5), balanced between the two senses. It leaves as (0.698, 0.197), decisively on the water side. Nothing was carried, nothing was gated, and no state was maintained. The word looked at its neighbours and rewrote itself.

The whole matrix, and something it reveals

Repeating the calculation for all four queries gives the full table.

Two things are worth reading off it.

The row for "the" is nearly flat: 0.222, 0.266, 0.252 and 0.260. A function word has no strong question to ask, so it takes a little of everything. The row for "bank" is not flat, and one third of what it reads comes from "river". A flat row is attention doing nothing, and a spiked row is attention doing something specific. Looking at how flat the rows are is a genuinely useful diagnostic.

The second thing is a defect. Every row puts its largest weight on "river". That is not because every word is asking about rivers. It is because the key matrix doubles the first dimension, and "river" is the token with the largest first dimension, so its key has the largest magnitude and scores well against any query. A token can attract attention simply by having a big key, regardless of what the query wanted. Attention weights are therefore evidence about the model, not an explanation of it, and reading them as an explanation is one of the more common mistakes made with these systems.

Hiding the future

A model that generates text one token at a time must not let a position read positions that come after it, or it will learn to predict a word by looking at the word. The fix is a mask applied before the softmax: set the scores for later positions to minus infinity, so their exponentials are zero and they take no weight.

Redo the row for "bank" with "flooded" masked out. The scaled scores for the first three positions are unchanged at 0.159, 1.061 and 0.795, their exponentials are 1.172, 2.889 and 2.214, and they now sum to 6.275 rather than 8.875. The weights become 0.187, 0.460 and 0.353, and they still sum to one.

The mixture is then 0.187 times (0.1, 0.1), plus 0.460 times (1.0, 0.0), plus 0.353 times (0.5, 0.5), which is (0.655, 0.195). Slightly less water than before, because "flooded" was contributing to that reading, and still decisively on the water side, because "river" alone was enough.

Removing a position does not merely delete its weight. The softmax redistributes it across everything that remains, so every other weight rose. That is worth remembering when padding is masked in a batch: the tokens that are left get more attention than they did, not the same amount.

Multiple heads

One query per position asks one question. Vaswani and colleagues run several in parallel: 8 heads, a model width of 512, and a key and value width of 64 each, so the heads together use the same total width as one full-width attention. Each head has its own query, key and value matrices, so one can learn to look at the previous word, one at the subject of the sentence, one at something with no name.

The outputs of the heads are concatenated and passed through a fourth matrix. The cost is the same as one wide head, and the behaviour is not, because eight separate softmaxes can hold eight different distributions where one cannot.

What it costs

Every position scores every position, so the number of dot products is the square of the sequence length. Four words gives 16. A 4,096 token context gives 16,777,216, per head, per layer, per forward pass.

Dao, Fu, Ermon, Rudra and Re state the constraint plainly: the time and memory complexity of self-attention are quadratic in sequence length. Their contribution is to attack the memory traffic rather than the arithmetic, using tiling to reduce the number of reads and writes between GPU high bandwidth memory and on-chip SRAM, which requires fewer high bandwidth memory accesses than standard attention. They report a 15 percent end-to-end speedup on BERT-large at sequence length 512, 3 times on GPT-2 at sequence length 1,000, and 2.4 times on long-range arena at 1,000 to 4,000, along with better-than-chance performance on Path-X at sequence length 16,000, at 61.4 percent accuracy.

Note what that work is and is not. It produces the same attention as before, faster, by respecting the memory hierarchy. The quadratic term is still there.

Two other costs are worth stating.

Attention has no idea what order the words are in. The calculation above would give the same weights for "flooded bank river the". Order has to be added to the inputs separately, because the mechanism itself is a set operation.

Attention alone cannot decide anything. Every output above is a weighted average of the values, so it lies inside the region the values span. Bahdanau, Cho and Bengio proposed this as a way past the fixed-length vector bottleneck described in Sequences, recurrence and its limits, not as a complete model, and it is not one. Averaging is a linear operation, so a stack of pure attention layers collapses in exactly the way the layers in A neuron is a weighted sum and a decision collapsed.

Something non-linear has to sit next to it. The next article assembles the full block, and shows that the part sitting next to attention holds twice as many parameters as attention does.

References

1. Neural Machine Translation by Jointly Learning to Align and Translate. Dzmitry Bahdanau, Kyunghyun Cho and Yoshua Bengio, arXiv, 2014.
2. Attention Is All You Need. Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N. Gomez, Lukasz Kaiser and Illia Polosukhin, arXiv, 2017.
3. FlashAttention, Fast and Memory-Efficient Exact Attention with IO-Awareness. Tri Dao, Daniel Y. Fu, Stefano Ermon, Atri Rudra and Christopher Re, arXiv, 2022.