The AI Reader.
The Mechanics

An attention head is a weighted average that computes its own weights

Nine dot products across three tokens compute one attention head's output, and the same nine show why long context costs what it does and why the softmax weights are not, by themselves, an explanation.

Prior Context

Everything a transformer does above its input layer is attention run many times over, so this one operation, understood once, unlocks the rest of the stack. Three facts fall out of its arithmetic and nowhere else, two of them deciding real money and real trust. Every token's query is scored against every token's key, and that all-pairs count sets the price of a long context window; it returns whenever serving cost or a token limit comes up. The scores are then squeezed into weights that read like the model reporting what it looked at, and whether that reading holds is a live argument in interpretability. The mixing carries no sense of order either, which is why position must be supplied from outside, a thread the next part of the course picks up.

The whole operation is an average

A self-attention head takes a set of token vectors and, for each one, returns a weighted average of value vectors drawn from all of them. The weights are not fixed. They are computed on the spot from how well the token doing the asking matches each of the others. That is the entire operation, written by Vaswani and colleagues as one line, softmax(QKT/√dk)V: score every query against every key, turn the scores into weights with a softmax, and mix the values by those weights.1

Two consequences already sit in that sentence. Scoring every query against every key means the work climbs with the square of the number of tokens, and that growth is the real content of a context-window limit. The weights, once computed, look like a confession of what the model paid attention to, and that appearance does not survive a careful test.

In this article
  • Project three token vectors into queries, keys, and values by hand.
  • Fill a three-by-three score grid, scale it, and mix the values into one output.
  • Count the grid to price context length; reorder the tokens to expose the missing position.

A weighted average, before the weights are learned

The idea predates transformers. Bahdanau, Cho, and Bengio built a translation model whose context vector was a weighted sum of the source annotations, with weights produced by a softmax over alignment scores.2 They read those weights as a probability that a target word aligns to a source word, the earliest version of the tempting idea that the weights say what the model is looking at.

The only open question is where the weights come from. They come from a similarity score between the token asking and each token available. The dot product is that score: two vectors pointing the same way have a large one, two at right angles have zero. Bahdanau computed it with a small trained network. Luong, Pham, and Manning showed the simpler alternative, a plain dot product, set beside the additive form as the multiplicative one.3 The transformer uses that dot product with one addition, a division by the square root of the vector's dimension, which is why the mechanism is named scaled dot-product attention.1

Three small projections turn a token vector into the parts of this score: the query is what it asks for, the key is what each token offers to a match, the value is the content that gets averaged. Query against key gives the weight; the weight multiplies the value. The projections are the learned part, and they are the last thing to add, not the first.

One head, three tokens, worked

Shrink everything to two dimensions. Take three token vectors, t1 = 1 0, t2 = 0 1, and t3 = 1 1, standing in for three embedding rows; the digits are invented to be legible, not produced by any model. Give the head three weight matrices, one each for query, key, and value: WQ with rows 1 0 and 1 1, WK with rows 1 1 and 0 1, and WV with rows 0 1 and 1 0. These are given, already trained; where a real head's numbers come from is gradient descent over a corpus, a step that has already happened. This is one forward pass through the finished weights.

Each projection is a single row-vector-times-matrix multiply. Token t3 = 1 1 projects to a query of 2 1, the first matrix row plus the second. All three tokens through all three matrices fill one table.4

Each token vector projected into a query, key, and value by the three given matrices.4
tokenvectorquerykeyvalue
t11 01 01 10 1
t20 11 10 11 0
t31 12 11 21 1

Now score. The recipe every tutorial repeats is the dot product of a query with a key.5 Dot each query against all three keys to fill a three-by-three grid. The query for t3, 2 1, against the key for t3, 1 2, scores 2×1 + 1×2 = 4. Every cell is one such sum.

scorekey t1 = 1 1key t2 = 0 1key t3 = 1 2
query t1 = 1 0101
query t2 = 1 1213
query t3 = 2 1314

Follow the last row, the query for t3, the rest of the way. First divide each score by the square root of the vector's dimension, here √2 ≈ 1.41, so 3 1 4 becomes 2.12 0.71 2.83. The division is not decoration. Vaswani's team added it because, for long vectors, raw dot products grow large enough to push the softmax into a region where its gradients nearly vanish, which stalls learning.1 Weng restates the reason plainly: without the scale, a large input can leave the softmax with an extremely small gradient.6

The softmax turns those scaled scores into weights that sum to one, exponentiating each and dividing by the total. This is the one step a pencil cannot finish, because ex is irrational; a calculator gives 8.34, 2.03, and 16.92, summing to 27.29. Dividing through gives weights of 0.31, 0.07, and 0.62, which add to 1.00. The output for t3 is those weights times the value vectors, 0.31×(0 1) + 0.07×(1 0) + 0.62×(1 1), which comes to 0.69 0.93.

The output is a blend of the three value vectors, every weight set by content similarity, the query-key scores, and nothing else. The value for t3 dominates because t3's query matched t3's key best, the top score of four. Phuong and Hutter write this same procedure as a formal algorithm, projection to score to softmax to weighted sum; the numbers above are that algorithm run once by hand.4

Nine scores, then thirty-six

The score grid has nine cells, because each of three queries met each of three keys. Feed the head six tokens and the grid is six by six, thirty-six cells: doubling the tokens quadrupled the scoring work. The number of dot products is the square of the sequence length, and Vaswani's complexity table lists the per-layer cost of self-attention as O(n2·d) for a sequence of length n.1

This square is the concrete thing behind a context-window number. A model advertised for a longer context must hold that same grid at, say, a hundred thousand tokens by a hundred thousand, and computing it grows as the second power of the length. That growth is the bottleneck two well-known lines of work were built to cut. Child and colleagues introduced sparse factorizations precisely because a dense transformer needs time and memory that grow quadratically with sequence length, reducing the attention term toward O(n√n).7 Kitaev, Kaiser, and Levskaya attacked the same term with locality-sensitive hashing, changing it from O(L2) to O(L·log L) in the length L.8 Both papers spend their length reducing exactly this grid.

The weights look like a reason but the arithmetic alone does not make one

The three weights for t3, 0.31, 0.07, and 0.62, invite a story: the token mostly read t3, glanced at t1, ignored t2. Bahdanau's alignment reading encouraged exactly this, treating the weights as where the model looks.2 But nothing in the arithmetic ties the size of a weight to how much that token changed the answer. The weight is a normalized similarity score. Whether it is also a faithful account of the model's reasoning is a separate question the mixing does not settle, and it is genuinely disputed.

Jain and Wallace The weights do not explain the prediction.

They found learned attention weights frequently uncorrelated with gradient-based measures of feature importance, and built very different distributions that yielded equivalent predictions.9

Serrano and Smith The weights carry a noisy signal, not a reliable one.

Higher weights correlated with greater impact in some cases and not others, and gradient-based rankings predicted a weight's effect better than its magnitude did; attention is no fail-safe indicator of importance.10

Wiegreffe and Pinter The strong dismissal is too strong.

The claim depends on one's definition of explanation, and one adversarial distribution does not prove the trained weights explain nothing; they separate a plausible explanation from a faithful one and keep the door open for the latter.11

The honest reading is the narrow one all three sides support. The arithmetic that produces the softmax weights does not, by itself, license treating them as the model's reasoning. It computes a similarity-weighted average; a faithful account of what drove the output is an extra claim the numbers neither make nor forbid. An attention map is a picture of that similarity, not yet a picture of a reason.

Reorder the tokens and nothing moves but the labels

One more fact hides in the same nine numbers. Feed the head those three tokens in a different order, say t3, t1, t2. Each still projects to the same query, key, and value, because the projection depends on the token, not its slot, so every pairwise score is one of the same nine. Each token's output vector is identical to before; only the row it lands in has moved. The head cannot tell the two orderings apart.

That invariance is a property, not an accident of the toy numbers. Yun and colleagues state it formally: without positional information the operation is permutation equivariant, and positional encodings are what circumvent that restriction.12 Shaw, Uszkoreit, and Vaswani put the plain version alongside it: the mechanism does not model position in its structure and instead requires adding representations of position to its inputs.13 Vaswani injects that signal at the bottom of the stack for the same reason, since a model with no recurrence and no convolution has no other way to use word order.1 Where the position vectors come from is the next part of the course.

One weighted average carried all of it. Its weights, read off the content by dot products, produced the output. Counting those dot products priced the context window. Reordering the inputs showed the head blind to sequence until position is added from outside. Inspecting the weights showed a distribution of similarity, not a record of reasoning. The equation is short and the head is one of hundreds in a real model, but each consequence is visible in the arithmetic and invisible in the slogan.

Sources

  1. Vaswani et al. · Attention Is All You Need (NeurIPS 2017)
  2. Bahdanau, Cho & Bengio · Neural Machine Translation by Jointly Learning to Align and Translate (2014)
  3. Luong, Pham & Manning · Effective Approaches to Attention-based Neural Machine Translation (2015)
  4. Phuong & Hutter · Formal Algorithms for Transformers (2022)
  5. Jay Alammar · The Illustrated Transformer
  6. Lilian Weng · Attention? Attention! (2018)
  7. Child, Gray, Radford & Sutskever · Generating Long Sequences with Sparse Transformers (2019)
  8. Kitaev, Kaiser & Levskaya · Reformer: The Efficient Transformer (2020)
  9. Jain & Wallace · Attention is not Explanation (2019)
  10. Serrano & Smith · Is Attention Interpretable? (2019)
  11. Wiegreffe & Pinter · Attention is not not Explanation (2019)
  12. Yun et al. · Are Transformers Universal Approximators of Sequence-to-Sequence Functions? (2020)
  13. Shaw, Uszkoreit & Vaswani · Self-Attention with Relative Position Representations (2018)