The Heart of Self-Attention: Queries, Keys, Values, and One Worked Example

The last article ended on a claim rather than a mechanism: attention lets any two tokens interact in one step, instead of pushing a signal through a chain of hidden states. This article fills in the mechanism, but only the smallest complete piece of it — the part that takes a handful of token vectors and produces a new vector for each one. By the end you should be able to run the whole computation by hand on three tokens, and say why each step is there rather than just what it looks like.

Every token gets three vectors

A Transformer receives the same thing the RNN did: a sequence of $n$ token vectors, each of length $d_{\text{model}}$. Call them $x_1,\dots,x_n$.

For each token, the model computes three vectors by multiplying $x_i$ with three separate weight matrices:

$q_i = W^Q x_i,\qquad k_i = W^K x_i,\qquad v_i = W^V x_i$

The matrices $W^Q$, $W^K$, $W^V$ are learned, and — exactly like the RNN's shared weights from last time — the same matrices are applied at every position. That is what makes this self-attention: every token is projected by the same rules, so all tokens end up in one shared query space and one shared key space, where they can be compared with each other.

The names come from database lookup. You make a query, each stored item advertises a key, and when the two match you get back the value. Treat this as a naming convention rather than a description: nothing here is matched discretely or retrieved intact. The match is a number, and the retrieval is a blend.

Why three matrices instead of one? If $q$ and $k$ were produced by the same projection, the score between tokens $i$ and $j$ would be symmetric — $i$'s interest in $j$ would equal $j$'s interest in $i$. Language is not symmetric. In "the cat sat on the mat", the verb "sat" has reason to care about "cat"; "cat" has much less reason to care about "sat". Separate projections make the relation directed, and let the model learn what "asking a question" and "advertising an answer" mean for its own data.

Two dimensions matter here. Queries and keys must have the same length, call it $d_k$, because we are about to take their dot product. Values can have their own length $d_v$, since they only ever get added together. In the original Transformer, $d_{\text{model}}=512$ and $d_k=d_v=64$ 1.

Scoring a pair: the dot product

The score between query $i$ and key $j$ is

$s_{ij} = q_i \cdot k_j = \sum_{t=1}^{d_k} q_{i,t}\,k_{j,t}$

This is large and positive when the two vectors point in the same direction, near zero when they are unrelated, and negative when they point in opposite directions. It is not a pure cosine — vector length matters too — but the ranking behaviour is what carries the meaning: a high score says "this key is relevant to this query".

Doing this for every query against every key produces an $n \times n$ table of scores. That $n^2$ is already visible, and it is where attention's quadratic cost comes from; we will come back to it.

Why the scores are divided by $\sqrt{d_k}$

This is the "scaled" in scaled dot-product attention, and it is the one step that is not obvious. The argument is statistical. Suppose the components of $q$ and $k$ are roughly independent, with mean 0 and variance 1 — which is what normalization layers in the network are there to encourage. Then each term $q_t k_t$ has mean 0 and variance 1, and $s_{ij}$ is a sum of $d_k$ such independent terms, so

$\mathbb{E}[s_{ij}] = 0, \qquad \operatorname{Var}(s_{ij}) = d_k$

A quantity with variance $d_k$ typically lands around $\pm\sqrt{d_k}$. So as $d_k$ grows, the scores get larger, not merely noisier. With $d_k=64$, typical scores sit around $\pm 8$ instead of $\pm 1$.

That matters because the next step is a softmax, which exponentiates. A gap of 8 between two logits means the higher one receives about $e^8 \approx 3000$ times the weight of the lower one. So a score table whose entries are spread over roughly $\pm 8$ produces rows that are nearly one-hot: one token takes essentially all the weight and the rest get almost none. Two things go wrong. First, attention stops being a blend and becomes a hard pick of a single token. Second — and this is what hurts training — when a softmax row is nearly one-hot, its gradient is near zero, so the model barely learns from it. The authors of the original paper put it plainly: for large $d_k$ the dot products "grow large in magnitude, pushing the softmax function into regions where it has extremely small gradients" 1.

Dividing every score by $\sqrt{d_k}$ brings the variance back to 1. That is the whole trick, and it explains the specific denominator: pulling a constant $c$ out of a variance squares it, so to turn a variance of $d_k$ into 1 you divide the values by $\sqrt{d_k}$, not by $d_k$. The paper also reports that this is not just theory — for small $d_k$, scaled and unscaled dot-product attention behave similarly, but for large $d_k$ the unscaled version performs worse, and scaling fixes it 1.

From scores to weights, and from weights to output

Softmax is applied row by row — for each query, across all the keys:

$a_{ij} = \frac{\exp\!\left(s_{ij}/\sqrt{d_k}\right)}{\sum_{l=1}^{n}\exp\!\left(s_{il}/\sqrt{d_k}\right)}$

Every $a_{ij}$ is positive, and each row sums to exactly 1. The row is therefore a probability distribution over the tokens: how much of this token's output should come from each position. Softmax never outputs an exact 0 for a finite input — it just makes small scores very small.

The output for token $i$ is the weighted sum of the value vectors:

$\text{out}_i = \sum_{j=1}^{n} a_{ij} v_j$

This form is worth dwelling on, because it limits what attention can express. The weights are non-negative and sum to 1, so the output is a convex combination of the value vectors — a point inside the shape they span. If the row is nearly uniform, the output is close to the plain average of all the values. If one weight is near 1, the output is close to that single value. Attention can only re-mix the value vectors it was handed; it cannot produce a direction that is not a blend of them.

The same thing in matrix form

Pack the vectors into rows: $Q$ is $n\times d_k$, $K$ is $n\times d_k$, $V$ is $n\times d_v$. Then everything above is one expression:

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

The shapes trace out the meaning. $QK^\top$ is $n\times n$, and entry $(i,j)$ is exactly $s_{ij}$. The softmax runs along the last axis, so it normalizes each row. Multiplying that $n\times n$ weight matrix by the $n\times d_v$ value matrix gives an $n\times d_v$ output — one new vector per input token. In code, skipping masking and dropout, this is about four lines 3:

1scores = Q @ K.transpose(-2, -1) / math.sqrt(d_k)   # (n, n)
2weights = scores.softmax(dim=-1)                    # row-wise, each row sums to 1
3out = weights @ V                                   # (n, d_v)

Doing it by hand

Take $n=3$ tokens, $d_k=2$, $d_v=2$, with simple integer rows:

$Q=\begin{bmatrix}1&0\\1&1\\0&1\end{bmatrix},\quad K=\begin{bmatrix}1&0\\1&1\\0&1\end{bmatrix},\quad V=\begin{bmatrix}1&0\\0&1\\1&1\end{bmatrix}$

Step 1 — raw scores $QK^\top$. Each entry is the dot product of one query row with one key row:

$\begin{bmatrix}1&1&0\\1&2&1\\0&1&1\end{bmatrix}$

For instance, entry $(2,2)$ is $q_2\cdot k_2 = 1\cdot 1 + 1\cdot 1 = 2$, and entry $(2,3)$ is $q_2\cdot k_3 = 1\cdot 0 + 1\cdot 1 = 1$.

Step 2 — divide by $\sqrt{d_k}=\sqrt{2}\approx 1.4142$:

$\begin{bmatrix}0.707&0.707&0\\0.707&1.414&0.707\\0&0.707&0.707\end{bmatrix}$

Step 3 — softmax each row. Using $\exp(0.707)\approx 2.028$, $\exp(1.414)\approx 4.113$, $\exp(0)=1$:

• row 1: $2.028/5.056,\ 2.028/5.056,\ 1/5.056 = [0.401,\ 0.401,\ 0.198]$
• row 2: $2.028/8.169,\ 4.113/8.169,\ 2.028/8.169 = [0.248,\ 0.503,\ 0.248]$
• row 3: $[0.198,\ 0.401,\ 0.401]$

Step 4 — multiply the weight matrix by $V$. For row 2:

$0.248\begin{bmatrix}1\\0\end{bmatrix}+0.503\begin{bmatrix}0\\1\end{bmatrix}+0.248\begin{bmatrix}1\\1\end{bmatrix}=\begin{bmatrix}0.248+0.248\\0.503+0.248\end{bmatrix}=\begin{bmatrix}0.497\\0.752\end{bmatrix}$

Doing the same for rows 1 and 3 gives $[0.599,\ 0.599]$ and $[0.599,\ 0.802]$.

Read the result back against the scores. $q_2$ matched $k_2$ best — score 2, the largest in the table — and token 2's output does lean on $v_2$. But nothing is sharp: the strongest weight is 0.503, so nearly half of token 2's output still comes from the other two tokens. Rows 1 and 3 happen to produce mirrored weight patterns, purely because $q_3$ and $k_3$ are the mirror images of $q_1$ and $k_1$ in these particular matrices. That is a property of the numbers I chose, not a rule of attention.

Why this is fast, and what it costs

Recall the second problem the previous article raised about RNNs: $h_5$ cannot be computed until $h_4$ exists, so a sequence of length $n$ costs $O(n)$ dependent steps, and a GPU spends most of that time waiting. Nothing in the recipe above has that shape. Row $i$ of the weight matrix depends only on $q_i$ and $K$; no output feeds into another. The entire $n\times n$ score table is a single matrix multiply, with only the row-wise softmax following it. The recurrence is gone.

The price sits in the shape of $QK^\top$: it is $n\times n$, so compute and memory grow with the square of the sequence length. Doubling the context length roughly quadruples the cost of the score matrix. Attention traded a long chain of dependent steps for a quadratic table of independent ones — a good trade on parallel hardware, and the reason long context windows are expensive rather than free.

What this piece does not include

Four things are deliberately missing, and it helps to know they are separate pieces so you do not go looking for them inside the formula.

• Position. $QK^\top$ treats the input as a set, not a sequence. Permute the tokens and the output rows are permuted in exactly the same way, so this computation cannot tell "dog bites man" from "man bites dog". Real Transformers add position information to the token vectors before this step.
• Multiple heads. The original model runs this whole computation $h=8$ times in parallel with eight different, smaller projections (each $d_k=d_v=64$ rather than $512$) and concatenates the results. Same formula, repeated, then re-projected 1.
• Causal masking. To generate text left to right, you set $s_{ij}=-\infty$ for every key $j$ that lies in the future of query $i$, before the softmax. Those positions then receive weight exactly 0, and each position can only see itself and the past. Nothing else changes 2.
• Self- vs cross-attention. Here $Q$, $K$, and $V$ all come from one sequence. Bahdanau attention in the previous article was cross-attention: queries from the decoder, keys and values from the encoder. Same arithmetic, different source of the three matrices.