Attention Is a Heatmap

Attention begins with a sequence of vectors, one per token — a token being a word, a piece of a word, or any other unit the model works with. Each vector describes what that position means so far. The question attention answers is relational: for this position, which other positions carry information I need?

The answer is a weighted average of the other positions' vectors, where the weights are learned from data rather than fixed by distance. Each position gets its own private set of weights. That is the whole idea. Everything else in this article exists to make that average expressive and stable.

Each token's vector is projected three ways, 'projected' meaning multiplied by a matrix the model learned during training.

- The query q = x·W_Q is what this position is looking for. - The key k = x·W_K is what this position offers to others. - The value v = x·W_V is the content that gets passed along if this position is chosen.

Here x is the token's vector, and W_Q, W_K, and W_V are the three learned matrices. A query and a key are compared with a dot product: multiply matching entries and add the products. A large positive q·k means the query is looking for something this key advertises.

Gather all the queries into a matrix Q, all the keys into K, and all the values into V. Then every comparison at once becomes one matrix multiplication, Q·Kᵀ, where the small T means rows and columns are swapped. Entry (i, j) of the result says how much query i likes key j.

The scores are divided by √d_k, where d_k is the number of entries in each query and key vector. The square root keeps the numbers from growing with dimension.

Suppose the entries of q and k are independent random numbers with average zero and spread one. Their dot product is then a sum of d_k terms, so it gets bigger as d_k grows. Feed huge scores into the step that turns them into weights and it saturates: one weight becomes 1 and the rest become 0. The slope that learning needs dies, and training stalls. Dividing by √d_k holds the scores at a workable scale no matter how wide the vectors are.

The step that turns scores into weights is called softmax. It is applied to each row of scores separately, and each row becomes weights between 0 and 1 that add up to exactly 1 — a budget of attention spent across the keys.

A peaked row means the query is locked onto one position. A flat row means it is spreading its attention widely. The entropy of the row, a number that measures how spread out the weights are, summarizes which case you are in.

In the demo, flip between the raw scaled scores and the attention weights, and hover a cell to see one query-key pair. Then follow the output. It is the weighted average of the values: softmax(QKᵀ/√d_k)·V. Row i of the output mixes all the value vectors, using the weights from row i of the attention table.

Notice what the shapes do: one output vector per input token, each of the same length as the values. That is what lets attention slot into a larger network. Around it sit residual connections (which add a layer's input back to its output), normalization steps, and a small feed-forward network, but the mixing itself is exactly this one multiplication.

A real transformer runs several attention operations side by side, each with its own W_Q, W_K, and W_V. These parallel copies are called heads.

One head might track grammar, another might track which pronoun refers to which noun, another might just look at nearby words. Each head has its own learned subspace and can develop its own notion of similarity.

The heads' outputs are stitched together and passed through one more learned matrix, W_O, so the model can combine what the heads found. Splitting the model's width d_model across h heads keeps the parameter count and the arithmetic roughly constant while giving the model several ways to look at the same sequence.

Two practical facts complete the picture. First, masks. If a position must not be looked at — say the future in a language model, or padding at the end of a batch — add a very large negative number (negative infinity in theory) to its score before the softmax. The exponential of that number is zero, so the forbidden position receives no attention.

Second, cost. For a sequence of length n, the table of scores has n² entries. Doubling the sequence length quadruples the table. That single fact is why long context windows are expensive and why so much engineering goes into sparse, sliding-window, and memory-saving variants.

The problems below ask you to scale the scores, soften the rows, take the weighted sum, and reproduce the whole computation end to end.