Skip to content

Volume ITechnology

12.7.3 — Attention Worked by Hand ​

Read this support message: "I ordered a kettle and a toaster last week. The kettle arrived but it was broken." You know immediately that "it" means the kettle, not the toaster and not the week. To know that, you had to connect "it" with a word nine words earlier, and choose the right one among several candidates.

A language model has to make that kind of connection for every token, at every layer. When the vector for "it" moves up through the transformer stack of 12.7.1, it must end up carrying the meaning "the kettle", which means it must gather information from the token "kettle". The mechanism that does this gathering is attention. 12.5.6 introduced the idea with a single query; this page builds the full version used in every transformer, with real matrices, every number computed, and then explains each design choice.

The clearest way to think about attention is as a search over the other tokens. Every token produces three vectors, and their names come from the language of databases.

  • A query is what this token is looking for. The query from "it" says, roughly, "I am a pronoun; I need the noun I refer to".
  • A key is what this token offers, like a label on a filing drawer. The key from "kettle" says "I am a singular object noun".
  • A value is the information this token hands over if it is chosen. The value from "kettle" carries what "kettle" means.
"it"query: "my noun?"I · keyordered · keykettle · keytoaster · keyarrived · keymatch score0.020.030.710.200.04mix the values0.71 × kettle's value+ 0.20 × toaster's+ tiny amounts of the restnew "it"≈ "the kettle"
Attention as a soft search. The query from "it" is compared with the key of every earlier token. The match scores, turned into weights that add to 1, decide how much of each token's value flows into the new vector for "it". The weights here are illustrative.

A database lookup finds the one row whose key exactly equals your query. Attention is softer in two ways: the "match" is a similarity score, the dot product from 12.2.1, and the result is a weighted blend of every value, not one row. Most of the weight goes to the best matches. Nobody writes a rule saying "pronouns attend to nouns". The model learns to produce queries and keys that match in useful ways, because doing so lowers its prediction loss.

Where the three vectors come from ​

Each token starts as one vector \mathbf{x}, its embedding plus position, or the output of the previous block. Its query, key and value are made by multiplying that vector by three learned weight matrices:

\mathbf{q} = \mathbf{x} W_Q, \qquad \mathbf{k} = \mathbf{x} W_K, \qquad \mathbf{v} = \mathbf{x} W_V

Those three matrices are the learned part of attention. The rest of the mechanism, as you will see, has no weights at all: it is a fixed recipe of dot products, a softmax and a weighted sum. Using three different matrices matters: a token can look for one thing (its query), advertise something else (its key), and hand over something else again (its value). "It" looks for a noun but offers itself as a pronoun.

2. The whole computation, with real numbers ​

Here is attention for three tokens from an order-status message, "order 4471 shipped", small enough to compute every number by hand. Each token's starting vector has 4 numbers:

X = \begin{bmatrix} 1 & 0 & 1 & 0 \\ 0 & 1 & 0 & 1 \\ 1 & 1 & 0 & 0 \end{bmatrix} \begin{matrix} \leftarrow \text{order} \\ \leftarrow \text{4471} \\ \leftarrow \text{shipped} \end{matrix}

The three learned matrices each turn 4 numbers into 2 (in real models, into 64 or 128):

W_Q = \begin{bmatrix} 1 & 0 \\ 0 & 1 \\ 1 & 1 \\ 0 & 0 \end{bmatrix}, \quad W_K = \begin{bmatrix} 1 & 0 \\ 0 & 0 \\ 1 & 1 \\ 0 & 1 \end{bmatrix}, \quad W_V = \begin{bmatrix} 1 & 0 \\ 0 & 1 \\ 2 & 0 \\ 0 & 3 \end{bmatrix}

XQ = XW_QK = XW_KV = XW_V① QKᵀevery pair's score② ÷ √d_k, ③ mask,④ softmax rows⑤ weights × Vone new vector per token
The five steps of attention as matrix operations. Only the three weight matrices at the start are learned; everything after them is a fixed recipe.

Step 0: make the queries, keys and values. Multiply X by each matrix, using the matrix multiplication of 12.2.2. For "order", whose row is (1, 0, 1, 0), the query is row 1 plus row 3 of W_Q: (1, 0) + (1, 1) = (2, 1). Doing every token:

TokenQuery \mathbf{q}Key \mathbf{k}Value \mathbf{v}
order(2, 1)(2, 1)(3, 0)
4471(0, 1)(0, 1)(0, 4)
shipped(1, 1)(1, 0)(1, 1)

Step ①: score every query against every key. The table of all dot products is QK^\top: entry (row i, column j) says how well token i's query matches token j's key. The transpose turns the keys into columns so that one matrix multiplication computes all nine scores at once. For "shipped" against "order": (1, 1) \cdot (2, 1) = 3.

QK^\top = \begin{bmatrix} 5 & 1 & 2 \\ 1 & 1 & 0 \\ 3 & 1 & 1 \end{bmatrix}

Step ②: scale by \sqrt{d_k}. Divide every score by the square root of the key size, here d_k = 2, so by \sqrt{2} \approx 1.414. The reason for this step is section 3. The row for "shipped" becomes (2.121, 0.707, 0.707).

Step ③: hide the future. An LLM generates text left to right, so when it processes "4471" it must not be allowed to look at "shipped", which comes later; otherwise, during training, it could simply read the next word instead of predicting it. So every score where the key comes after the query is replaced by -\infty before the softmax. This is the causal mask, and it turns the score table into a lower triangle:

\begin{bmatrix} 3.536 & -\infty & -\infty \\ 0.707 & 0.707 & -\infty \\ 2.121 & 0.707 & 0.707 \end{bmatrix}

Step ④: softmax each row. Softmax from 12.2.6 turns each row into weights that add up to 1, and e^{-\infty} = 0, so masked positions get exactly zero weight. Row 1 has only one visible entry, so "order" puts all its weight on itself. Row 2 has two equal scores, so "4471" splits 0.5 and 0.5. Row 3: e^{2.121} \approx 8.34, e^{0.707} \approx 2.03 twice, adding to 12.40, so the weights are 8.34/12.40 \approx 0.673 and 0.164 twice.

attends to (keys)order4471shippedorder4471shipped1.00maskedmasked0.500.50masked0.670.160.16querieseach rowadds to 1
The attention weights after the causal mask and softmax. Each row says how a token divides its attention among itself and the tokens before it; the greyed cells are the future, which the mask hides. "shipped" pays two-thirds of its attention to "order".

Step ⑤: mix the values. Each token's output is its row of weights times the value vectors:

  • order: 1 \times (3, 0) = (3, 0), unchanged, since it can only see itself.
  • 4471: 0.5 \times (3, 0) + 0.5 \times (0, 4) = (1.5, 2), an equal blend of itself and "order".
  • shipped: 0.673 \times (3, 0) + 0.164 \times (0, 4) + 0.164 \times (1, 1) \approx (2.18, 0.82), mostly "order".

The new vector for "shipped" now contains information about what was shipped. That is attention's whole purpose: after this layer, each token's vector reflects not just itself but the parts of its context it found relevant. Written as one formula, the five steps are:

\text{Attention}(Q, K, V) = \text{softmax}\!\left(\frac{QK^\top}{\sqrt{d_k}} + M\right) V

Read it as "softmax of Q times K-transpose, divided by root d-k, plus the mask, all times V", where M holds 0 for allowed positions and -\infty for the future. It is called scaled dot-product attention, and it comes from the 2017 paper "Attention Is All You Need" by Ashish Vaswani and seven colleagues at Google.

python
import numpy as np

X  = np.array([[1, 0, 1, 0], [0, 1, 0, 1], [1, 1, 0, 0]], float)
WQ = np.array([[1, 0], [0, 1], [1, 1], [0, 0]], float)
WK = np.array([[1, 0], [0, 0], [1, 1], [0, 1]], float)
WV = np.array([[1, 0], [0, 1], [2, 0], [0, 3]], float)

Q, K, V = X @ WQ, X @ WK, X @ WV                       # step 0
scores = Q @ K.T / np.sqrt(K.shape[1])                 # (1) steps ① and ②
mask = np.triu(np.ones((3, 3), bool), k=1)             # (2) True above the diagonal
scores[mask] = -np.inf                                 # (3) step ③
w = np.exp(scores - scores.max(axis=1, keepdims=True))
w /= w.sum(axis=1, keepdims=True)                      # (4) step ④, row by row
print(w.round(3))      # → [[1. 0. 0.] [0.5 0.5 0.] [0.673 0.164 0.164]]
print((w @ V).round(2))   # (5) → [[3. 0.] [1.5 2.] [2.18 0.82]]

Line (1) computes all nine scaled scores in one matrix multiplication. Line (2) builds the mask: np.triu(..., k=1) marks every cell strictly above the diagonal, which is every "key comes after query" cell. Line (3) sets those to minus infinity. Line (4) is a numerically safe softmax for each row, subtracting the row's maximum first as in 12.2.6. Line (5) mixes the values. The printed weights and outputs match the hand calculation. PyTorch provides all of this, heavily optimised, as one function: torch.nn.functional.scaled_dot_product_attention(Q, K, V, is_causal=True).

3. Why divide by the square root of d_k? ​

The scaling in step ② looks like a detail, and it is a small numerical fix, but without it large models train badly. Here is the reason.

If the numbers in a query and a key are each roughly random with a typical size of 1 (variance 1), their dot product adds up d_k such products, and its variance grows to about d_k. In real models d_k is 64 or 128. With d_k = 128, typical scores have a spread of about \sqrt{128} \approx 11. Look at what softmax does to scores that far apart:

  • Scores (11, 0, -11) give softmax weights of about (0.99998, 0.00002, 0.0000): essentially all the weight on one token.
  • Scores (1, 0, -1) give (0.665, 0.245, 0.090): a real blend.

When scores are huge, softmax becomes almost a hard choice of one position, and, as 12.2.6 showed, its slope in that saturated region is nearly zero. The gradient that should teach the queries and keys to improve all but vanishes, and training stalls. Dividing by \sqrt{d_k} shrinks the spread back to about 1, whatever the size of the vectors, so softmax stays in its responsive middle range. This is the same "keep the signal near size 1" reasoning as the weight initialisation in 12.5.4.

4. What attention costs ​

Look at the shape of the score table: one row per token and one column per token. For a sequence of n tokens, attention computes n \times n scores in every head of every layer.

Tokens in contextScores per head per layer
1,0001 million
10,000100 million
100,00010 billion

Ten times longer means a hundred times more work in this part of the model. This quadratic growth is why long prompts are expensive, why context windows had to be engineered to grow, and why several important tricks exist. The most important, FlashAttention, computes exactly the same result without ever storing the whole table in slow memory, and is covered with the other speed-ups in Chapter 12.7.8. The causal mask also has a payoff for generation: because a token never attends to later tokens, its keys and values never change once computed, so they can be saved and reused for every later token. That saved store is the KV cache from 12.7.1, and Chapter 12.7.6 is about it.

5. Three things attention does not do on its own ​

It does not know the order of the tokens. Look back at the computation: nothing in it depends on where a token sits. Shuffle the rows of X and you get the same outputs, shuffled the same way. "Dog bites man" and "man bites dog" would be indistinguishable. Position information has to be added, which is part of the transformer block in 12.7.4.

It computes only one kind of relationship at a time. One set of W_Q, W_K, W_V learns one notion of "relevant". A real sentence needs several at once: which noun a pronoun refers to, which verb goes with which subject, what the previous token was. The answer is to run several attentions side by side, called heads, also in 12.7.4.

It only mixes; it does not transform. The output for each token is a weighted average of value vectors. The feed-forward layer that follows each attention layer is what then processes the gathered information, and 12.7.4 assembles the two.

A note on what attention is not: high attention weight on a token is often presented as "what the model was thinking about", but it is a weak explanation of a model's decision. A model has dozens of layers and heads, the values being mixed have already been transformed many times, and research has shown that very different attention patterns can lead to the same output. Attention maps are a useful debugging picture, not a reliable account of why the model said something.

What the interviewer will push on ​

"Explain attention." They want query, key and value as a soft lookup, the formula, and what each step does, ideally with the shapes: Q and K are n \times d_k, the score table is n \times n, and the output is n \times d_v. Saying "it lets the model focus on important words" with nothing underneath is the weak answer.

"Why divide by \sqrt{d_k}?" The tell of a real understanding: dot products of d_k-dimensional vectors have variance about d_k, large scores saturate softmax, and saturated softmax has near-zero gradients.

"What is the causal mask and why is it needed?" Future positions get -\infty before softmax so they receive zero weight; without it, a model trained on next-token prediction could read the answer. Bonus: it is what makes the KV cache possible.

"What is the complexity of attention?" O(n^2 d) in time and, naively, O(n^2) memory for the score table per head; a strong answer adds that FlashAttention removes the memory problem without changing the result.

Volunteer this: the only learned parts of attention are the projection matrices. The mechanism itself is a fixed recipe, so everything the model "knows" about what to attend to lives in W_Q, W_K and W_V.

Recall

  • Attention lets each token gather information from others. Each token makes a query (what it looks for), a key (what it offers) and a value (what it hands over), by multiplying its vector by three learned matrices.
  • Scaled dot-product attention: \text{softmax}(QK^\top/\sqrt{d_k} + M)V. Score every pair, scale, mask, softmax each row, mix the values. In the worked example "shipped" put 0.673 of its attention on "order" and its output became (2.18, 0.82).
  • The causal mask sets future positions to -\infty so they get zero weight; it stops training from cheating and makes the KV cache possible.
  • Dividing by \sqrt{d_k} keeps scores near size 1; without it softmax saturates and gradients vanish.
  • The score table is n \times n, so cost is quadratic in context length. Attention alone ignores order, captures one relationship per head, and only mixes; attention maps are a weak explanation of decisions.

Self-test: In the worked example, what weights would "4471" give if its key were (2, 0) instead of (0, 1)? Why does shuffling the input rows only shuffle the outputs? With d_k = 64, by what number are the scores divided?

Next: 12.7.4 The Transformer Block — attention is the heart of the transformer but not the whole of it; the next page adds multiple heads, positions, the feed-forward layer and the residual wiring, and assembles the block that every LLM stacks dozens of times.