Attention, worked through with three tokens

The idea

In the sentence “the cat sat”, what “sat” means depends on who sat. Attention gives every token a way to look up information from the other tokens. Each token produces three vectors:

a query $q$ : what this token is looking for;
a key $k$ : what this token offers to be found by;
a value $v$ : the information it hands over when found.

A token compares its query with every key using a dot product, turns the scores into weights with softmax, and takes the weighted average of the values. It is a soft, differentiable version of a dictionary lookup: instead of returning the one value whose key matches, it returns a blend, weighted by how well each key matches.

The formula

Stack the token embeddings as the rows of a matrix $X$ . Three learned weight matrices produce the queries, keys and values for all tokens at once:

Q = XW_Q, \qquad K = XW_K, \qquad V = XW_V.

Then scaled dot-product attention is

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

where softmax is applied to each row separately and $d_k$ is the length of the key vectors. Entry $(i, j)$ of $QK^\top$ is $q_i \cdot k_j$ , the score of token $i$ attending to token $j$ .

Why divide by $\sqrt{d_k}$ ?

If the entries of $q$ and $k$ are roughly independent with mean 0 and variance 1, their dot product has variance $d_k$ . With $d_k = 64$ , typical scores would be around ±8, and softmax of such spread-out numbers is almost one-hot: one weight near 1, the rest near 0, with gradients near zero. Dividing by $\sqrt{d_k}$ brings the variance back to 1 and keeps the softmax in a range where it can learn.

Worked example

Step 1: queries, keys and values

The tokens “the”, “cat” and “sat” have embeddings $(1,0)$ , $(0,1)$ and $(1,1)$ . With the weight matrices

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

we get, multiplying row by row,

Q = \begin{pmatrix} 1 & 0 \\ 1 & 2 \\ 2 & 2 \end{pmatrix},\quad K = \begin{pmatrix} 2 & 0 \\ 0 & 1 \\ 2 & 1 \end{pmatrix},\quad V = \begin{pmatrix} 1 & 0 \\ 0 & 2 \\ 1 & 2 \end{pmatrix}.

For example, the query for “sat” is $(1, 1)\,W_Q = (1 + 1,\; 0 + 2) = (2, 2)$ .

Step 2: scores

Every query dotted with every key gives $QK^\top$ ; then divide by $\sqrt{d_k} = \sqrt 2 \approx 1.414$ :

QK^{\top} = \begin{pmatrix} 2 & 0 & 2 \\ 2 & 2 & 4 \\ 4 & 2 & 6 \end{pmatrix} \quad\longrightarrow\quad \frac{QK^{\top}}{\sqrt 2} = \begin{pmatrix} 1.414 & 0.000 & 1.414 \\ 1.414 & 1.414 & 2.828 \\ 2.828 & 1.414 & 4.243 \end{pmatrix}.

Row 3 is “sat”: $q_{\text{sat}} \cdot k_{\text{the}} = 2 \cdot 2 + 2 \cdot 0 = 4$ , $q_{\text{sat}} \cdot k_{\text{cat}} = 2$ , $q_{\text{sat}} \cdot k_{\text{sat}} = 6$ .

Step 3: softmax each row

A = \begin{pmatrix} 0.446 & 0.108 & 0.446 \\ 0.164 & 0.164 & 0.673 \\ 0.187 & 0.045 & 0.768 \end{pmatrix}

For “sat”: $e^{2.828}, e^{1.414}, e^{4.243} \approx 16.92, 4.11, 69.6$ , which sum to about 90.6, giving weights 0.187, 0.045, 0.768. “sat” attends mostly to itself and somewhat to “the”. Each row sums to 1.

Step 4: weighted sum of values

Z = AV = \begin{pmatrix} 0.892 & 1.108 \\ 0.836 & 1.673 \\ 0.955 & 1.627 \end{pmatrix}

For “sat”: $0.187\,(1, 0) + 0.045\,(0, 2) + 0.768\,(1, 2) \approx (0.955, 1.627)$ . That vector, a blend of the three tokens’ values, is the new representation of “sat”.

Try it

The widget starts from the same numbers. Edit any embedding or weight and every matrix updates. Click a row of the heatmap to see that token’s calculation written out.

Causal masking

A model that generates text one token at a time must not let a token look at tokens that come after it, or training would let it cheat by reading the answer. The fix is a causal mask: before the softmax, set every score with $j > i$ to $-\infty$ , so its weight becomes exactly 0. In the example, “the” can then only attend to itself (weight 1), “cat” splits its attention between “the” and itself (0.500 and 0.500), and “sat” is unchanged because it was already last. Tick the checkbox in the widget to see it.

From one head to a transformer

Multi-head attention. Instead of one set of $W_Q, W_K, W_V$ , a layer runs several in parallel (often 8 to 128 heads), each with smaller vectors, and concatenates the results. Different heads can specialise: one may track which noun a pronoun refers to, another the previous token.
Position. Attention by itself ignores order: shuffle the tokens and the outputs shuffle the same way. Transformers add positional information to the embeddings, or build it into the scores, so that order can matter.
Cost. The score matrix has an entry for every pair of tokens, so attention over a sequence of length $n$ takes $O(n^2 d)$ time and $O(n^2)$ memory. That quadratic cost is why long context windows are expensive, and the target of a great deal of research.

Where it came from Attention was introduced for machine translation in 2014, letting a decoder look back at the relevant words of the source sentence. The 2017 transformer paper, “Attention is all you need”, dropped recurrence entirely and built the whole model from attention and simple feed-forward layers.

References

A. Vaswani, N. Shazeer, N. Parmar, J. Uszkoreit, L. Jones, A. N. Gomez, Ł. Kaiser and I. Polosukhin, “Attention is all you need”, NeurIPS 2017. arXiv:1706.03762
D. Bahdanau, K. Cho and Y. Bengio, “Neural machine translation by jointly learning to align and translate”, ICLR 2015. arXiv:1409.0473
3Blue1Brown, Neural networks, whose later chapters explain transformers and attention visually.

The idea

The formula

Why divide by dk\sqrt{d_k}dk​​?

Worked example

Step 1: queries, keys and values

Step 2: scores

Step 3: softmax each row

Step 4: weighted sum of values

Try it

Scaled dot-product attention, three tokens

Causal masking

From one head to a transformer

References

Keep going

Softmax and cross-entropy

Vectors and embeddings

Gradient descent

Why divide by $\sqrt{d_k}$ ?