Queries, keys and values, step by step

Key idea

Q = XW_Q, K = XW_K, V = XW_V; scores QKᵀ/√d; causal mask; softmax by row; output AV

Real world

A room of people at a networking event. Each person holds a card saying what they’re looking for (query) and wears a badge saying what they offer (key). Everyone walks towards the badges that match their card, and listens to those people (values) in proportion.

In an LLM

In “The cat sat on the”, the last “the” has a query that looks for a noun. The key of “cat” matches it best, so the output for “the” is mostly cat’s value: the model has carried information about the cat forward to the position that will predict the next word.

The recipe

The previous page described attention as a soft lookup. Here is the exact computation one attention head does. The input is a matrix $X$ with one row per token (5 tokens × 4 numbers in our toy model; thousands of numbers in a real one).

1. Project. Multiply $X$ by three learned matrices to get three new matrices, each with one row per token:
   • $Q = X W_Q$: each token’s query, what it is looking for.
   • $K = X W_K$: each token’s key, what it contains.
   • $V = X W_V$: each token’s value, what it passes on.
2. Score. $S = Q K^\top$: the dot product of every query with every key. Entry $(i, j)$ says how well token $i$’s query matches token $j$’s key.
3. Scale. Divide by $\sqrt{d}$, where $d$ is the length of a query.
4. Mask. Set every score where $j$ comes after $i$ to $-\infty$, so tokens can’t see the future.
5. Softmax each row, turning scores into weights $A$ that are positive and add up to 1.
6. Mix. $O = A V$: each token’s output is the weighted average of the value rows.

where $M$ is the mask: 0 where looking is allowed, $-\infty$ above the diagonal.

Try it

This is a real attention head with hand-picked weights so you can read it. The four embedding numbers mean thing, action, small word and first token. Press play, then click a token to follow its row, and hover any cell of a result to see exactly how it was computed.

Reading the matrices

Follow the last “the” through the steps:

• Its embedding is mostly small word. $W_Q$ maps small word onto the first query dimension, so its query says “I’m looking for a thing”.
• $W_K$ maps thing onto that same first dimension, so “cat” has a large first key number.
• Their dot product is the largest score in the row. After scaling, masking and softmax, “the” gives about half of its attention to “cat”.
• The output row for “the” is mostly cat’s value: a vector that says thing.

Nobody wrote these rules into a real model. $W_Q$, $W_K$ and $W_V$ start random and are learned by gradient descent, because moving the right information to the right place lowers the next-token loss. Real heads learn patterns like “look at the previous token”, “look at the subject of the sentence” or “look at what followed this token last time”: see real attention patterns.

Every token at once

The widget follows one row, but the matrices compute every token’s attention in one go. That is why attention runs fast on GPUs, which are built for big matrix multiplications, and why training can process all positions of a sequence in parallel. The mask is what makes this safe: every row only uses the rows above it.

Shapes

With $T$ tokens, model width $d_{\text{model}}$ and head size $d$:

The $T \times T$ score matrix is why attention gets expensive for long texts: double the length and it has four times as many entries. See what a forward pass costs.

import torch, math
 
def attention(X, Wq, Wk, Wv):
    Q, K, V = X @ Wq, X @ Wk, X @ Wv          # (T, d) each
    S = Q @ K.T / math.sqrt(Q.shape[-1])      # (T, T)
    T = X.shape[0]
    mask = torch.triu(torch.ones(T, T, dtype=torch.bool), diagonal=1)
    S = S.masked_fill(mask, float("-inf"))
    A = S.softmax(dim=-1)                     # each row sums to 1
    return A @ V                              # (T, d)