Query, Key, Value

Step 3 of 7. In: \(X\) from step 2, \(5 \times 4\). Out: for each token a new vector that mixes the others, computed here for one head; step 4 adds the causal mask and the second head. ← Embeddings · next step: Masks and heads →

A soft dictionary

Recall the idea from the overview: every token builds a new vector by mixing the vectors of the others, with more weight on the ones that matter — "bank" pulls from "river" in one sentence and from "cash" in the other.

How do you decide how much each neighbour weighs? The answer comes from a structure you already know, the dictionary:

stock = {"apple": 3, "pear": 7}
stock["pear"]       # the key matches exactly one entry and returns its value: 7

Attention is a soft dictionary. Instead of demanding an exact match, it compares the question against every key, gives each one a similarity score, and returns a weighted average of all the values. Because the average is a smooth function of the scores, it can be differentiated — and the network learns, by gradient, what to ask and what to advertise.

The tool for "how similar" is the dot product, \(q \cdot k = \sum_i q_i k_i\): large and positive when the two vectors point the same way, near zero when they are unrelated, negative when they point apart. Everything below is that one operation, done for every pair of tokens.

Query, Key, Value

The soft dictionary needs three things from each token, and attention names each one:

• Query — what this position is looking for. It is the question.
• Key — what each position advertises about itself. It is the label compared with the question.
• Value — what each position actually delivers if selected. It is the content.

Score every key against the query with a dot product, turn the scores into a distribution with the softmax, and return the corresponding mixture of values:

Q, K and V are three different linear projections of the same input, \(Q = XW_Q\), \(K = XW_K\), \(V = XW_V\). The matrices \(W_Q\), \(W_K\) and \(W_V\) are the layer's only learned parameters. The three projections are why a token can ask for one thing, advertise another, and deliver a third — a pronoun can query for "the noun I refer to" without itself being a noun.

The panel walks through that chain one matrix at a time, with the shapes of a three-token sentence and the numbers of the worked example that follows. Pick a token to follow its row.

A worked example, by hand

Take the first of the two "bank" sentences of the overview and three of its words: "sat · bank · river" (the others are left out only so the arithmetic fits). We will compute the output for "bank". The vectors have two dimensions and were picked by hand; the values have a readable meaning, \([\text{river},\ \text{money}]\):

The query of "bank" is \(q = [1,\ 2]\) and \(d_k = 2\).

1. Scores. The dot product of \(q\) with each key: \(q \cdot k_\text{sat} = 1\), \(q \cdot k_\text{bank} = 1 + 2 = 3\), \(q \cdot k_\text{river} = 4\).
2. Scale. Dividing by \(\sqrt{2} \approx 1.41\): \([0.71,\ 2.12,\ 2.83]\).
3. Softmax. \(e^{0.71} \approx 2.0\), \(e^{2.12} \approx 8.3\), \(e^{2.83} \approx 16.9\); they sum to \(27.3\). Dividing each by the sum, the weights are \(a = [0.07,\ 0.31,\ 0.62]\). They are positive and sum to 1.
4. Mix. The output is the average of the values with those weights:

The "bank" that came in on the fence, \([0.5,\ 0.5]\), leaves the layer as \([0.85,\ 0.15]\): the river bank. No weight was fixed by position. The choice came from how much the query of "bank" resembles each neighbour's key.

The panel starts at exactly these numbers. Redo the sum following steps ①–④, then try:

• change the value of "river" to \([0,\ 1]\) — as if the sentence said "…bank… cash" — and watch the verdict flip without a single attention weight moving;
• move the query to \([1,\ 0]\) and watch the weight leave "river" and split between "sat" and "bank";
• set the query to zero: every score becomes 0 and attention becomes a plain average.

The same thing, in code

The whole function fits in three lines. The script redoes the sum above, then runs the same function for all three queries at once — the middle row of \(A\) and \(O\) is the "bank" you computed by hand (the script's comments use the Portuguese example, sentei · banco · praça, with the same numbers):

"""O exemplo feito à mão na página, refeito em código — e depois a mesma conta para todas as linhas.

Contexto: "sentei · banco · praça". A query é a de "banco"; os values têm duas dimensões,
[assento, dinheiro]. O "banco" sozinho é ambíguo ([0.5, 0.5]); o contexto decide.
"""
import numpy as np

def softmax(s):
    e = np.exp(s - s.max(axis=-1, keepdims=True))     # subtrair o máximo não muda o resultado
    return e / e.sum(axis=-1, keepdims=True)

def attention(Q, K, V):
    S = Q @ K.T / np.sqrt(K.shape[-1])                 # (n, n): score de cada par
    A = softmax(S)                                     # (n, n): cada linha soma 1
    return A @ V, A                                    # (n, d_v): mistura de values

K = np.array([[1.0, 0.0],      # sentei
              [1.0, 1.0],      # banco
              [0.0, 2.0]])     # praça
V = np.array([[1.0, 0.0],      # sentei -> assento
              [0.5, 0.5],      # banco  -> ambíguo
              [1.0, 0.0]])     # praça  -> assento
q = np.array([1.0, 2.0])       # a query de "banco"

np.set_printoptions(precision=2, suppress=True)
s = K @ q
print("q·k         =", s)
print("q·k / √2    =", s / np.sqrt(2))
print("a = softmax =", softmax(s / np.sqrt(2)))
print("Σ a·v       =", softmax(s / np.sqrt(2)) @ V)

# A mesma função, para as três queries ao mesmo tempo: uma linha de A por token.
Q = np.array([[2.0, 0.0],      # sentei
              [1.0, 2.0],      # banco (a mesma query de cima)
              [0.0, 1.0]])     # praça
O, A = attention(Q, K, V)
print("\nA =\n", A)
print("O =\n", O)

q·k         = [1. 3. 4.]
q·k / √2    = [0.71 2.12 2.83]
a = softmax = [0.07 0.31 0.62]
Σ a·v       = [0.85 0.15]

A =
 [[0.45 0.45 0.11]
 [0.07 0.31 0.62]
 [0.14 0.28 0.58]]
O =
 [[0.78 0.22]
 [0.85 0.15]
 [0.86 0.14]]

The running example: one head, by hand

The bank example picked \(q\), \(k\) and \(v\) by hand. In a model they are computed from \(X\), and that is where the running example continues: \(X\) from the embeddings page, five rows of four numbers, one per token of "the cat climbed the".

Step 0 — the same scale for every row. In a GPT-2 block, attention does not read \(X\) itself but a normalized copy, \(u = \text{LayerNorm}(X)\): from each row subtract its mean \(\mu\), then divide by its standard deviation \(\sigma\) (why, and the rest of the block, are on the block page). Every row of our \(X\) already has mean 0, so only the division is left:

For the last row: \(\sigma = \sqrt{(0.25 + 0.25 + 0.25 + 2.25)/4} = \sqrt{0.75} = 0.866\), and \(1.5 / 0.866 = 1.732\).

Step 1 — three projections. \(Q = uW_Q\), \(K = uW_K\) and \(V = uW_V\), each \(5 \times 4\). A trained model's matrices are dense; the toy's only pick columns of \(u\), so that every number can be traced:

The bar separates the two heads: columns 1–2 are head 1, the one on this page, and columns 3–4 are head 2, on the next page. For the last token, ␣the, \(u = [-0.577, -0.577, -0.577, 1.732]\), so its head-1 query is \(q = [u_4, u_2] = [1.732, -0.577]\). The head-1 keys \(k = [u_2, u_4]\) and values \(v = [u_1, u_2]\) of the five tokens:

Steps 2 to 4 — the four sums of the bank example, again. \(q = [1.732, -0.577]\) and \(d_k = 2\):

④ The output mixes the values with those weights:

The last token puts 94% of its attention on ␣climb and ed: it looks at the verb, which the tokenizer had cut in two. Nothing told it where the verb is; the weights came from how much its query resembles each key, and those came from \(W_Q\) and \(W_K\).

The panel runs the same chain for every row. Rows 1 to 4 look at tokens that come after them — the, at position 1, puts 0.66 on ␣climb — which a model that generates text must not do; the next page forbids it with the mask. The last row has nothing after it, so its numbers survive the mask unchanged.

The shape of every matrix

For a sentence of \(n\) tokens, each with \(d\) dimensions, this is how the shapes chain together. Half the bugs in a new implementation show up as a shape error in this table:

In the running example \(n = 5\), \(d = 4\) and \(d_k = d_v = 2\) per head: \(S\) and \(A\) are \(5 \times 5\) and \(O\) is \(5 \times 2\). In GPT-2 small, \(d = 768\) and \(d_k = d_v = 64\) per head.

The whole matrix, and why \(\sqrt{d_k}\) is not cosmetic

One query against three keys is the mechanism. What a layer actually computes is the full \(n \times n\) matrix: every position queries every position. The panel below builds it for a real sentence. Read it by row: a token's row shows where it looks, and every row sums to 1.

Each option under head is a different score function. The first three describe patterns that genuinely occur in trained models — a head that matches by content, a head that just looks one token back, a head that decays with distance — and the fourth mixes them. A layer runs many heads at once.

Now turn divide by \(\sqrt{d_k}\) off and drag \(d_k\) up. The matrix goes binary.

The reason is a two-line argument. If \(q\) and \(k\) have independent components with mean 0 and variance 1, then

because it is a sum of \(d_k\) independent terms, each with variance 1. The scores therefore have standard deviation \(\sqrt{d_k}\) — at \(d_k = 128\), about 11, and scores routinely differ by 20. A softmax over scores that differ by 20 is a hard argmax: \(e^{-20} \approx 2 \times 10^{-9}\), so one weight is 1, the rest are \(10^{-9}\), and the gradient through the softmax is zero. Dividing by \(\sqrt{d_k}\) restores unit variance and keeps the layer in the regime where it can still learn.

In the running example \(d_k = 2\), so the effect is mild but visible: without the division the last row's weights go from \([0.011, 0.035, 0.622, 0.319, 0.014]\) to \([0.002, 0.012, 0.707, 0.275, 0.003]\), already sharper. Untick "divide by √dₖ" in the running-example panel above to see it.

The backward pass through one head

This subsection is the most technical in the chapter. On a first read, the bold conclusion of the third paragraph is enough.

The sentence above — the gradient through the softmax is zero — falls out of the arithmetic, and writing it costs four lines. With \(S = QK^\top/\sqrt{d_k}\), \(A = \text{softmax}(S)\) and \(O = AV\):

The middle term is the softmax Jacobian, and it is where the lesson sits. Look at the \(A \odot\) factor: if a row of \(A\) is nearly one-hot — one weight at 1 and the rest at \(10^{-9}\) — then almost all of the bracket is multiplied by \(10^{-9}\); and at the one position where \(A\) is not tiny, the bracket itself is close to zero, because the weighted mean it subtracts is that very element. A saturated attention hands back a numerically zero gradient across the whole row. That is what \(\sqrt{d_k}\) protects against.

The same formula gives a cheap invariant for checking an implementation: every row of \(\partial L/\partial S\) sums to exactly zero, because adding a constant to a row of scores does not change its softmax.

"""A passagem reversa de uma cabeça de atenção, conferida por diferenças finitas.

O termo que importa é o do softmax: ele não é um fator por elemento, é uma jacobiana que
subtrai a média ponderada da linha — e é ela que zera o gradiente quando a atenção satura.
"""
import numpy as np

def softmax(s):
    e = np.exp(s - s.max(axis=-1, keepdims=True))
    return e / e.sum(axis=-1, keepdims=True)

def forward(Q, K, V):
    d = Q.shape[-1]
    S = Q @ K.T / np.sqrt(d)
    A = softmax(S)
    return A @ V, A

def backward(dO, Q, K, V, A):
    d = Q.shape[-1]
    dV = A.T @ dO
    dA = dO @ V.T
    dS = A * (dA - (dA * A).sum(axis=-1, keepdims=True))     # jacobiana do softmax, linha a linha
    return dS @ K / np.sqrt(d), dS.T @ Q / np.sqrt(d), dV, dS

rng = np.random.default_rng(0)
Q, K, V = (rng.normal(size=(3, 4)) for _ in range(3))
dO = rng.normal(size=(3, 4))

O, A = forward(Q, K, V)
dQ, dK, dV, dS = backward(dO, Q, K, V, A)

def num_grad(M):
    g, h = np.zeros_like(M), 1e-6
    for i in range(M.shape[0]):
        for j in range(M.shape[1]):
            old = M[i, j]
            M[i, j] = old + h; up = (dO * forward(Q, K, V)[0]).sum()
            M[i, j] = old - h; dn = (dO * forward(Q, K, V)[0]).sum()
            M[i, j] = old; g[i, j] = (up - dn) / (2 * h)
    return g

np.set_printoptions(precision=3, suppress=True)
print("pesos de atenção A (cada linha soma 1):\n", A)
print("dQ:\n", dQ)
print("erro máximo contra diferenças finitas — dQ:", f"{np.abs(num_grad(Q) - dQ).max():.1e}",
      " dK:", f"{np.abs(num_grad(K) - dK).max():.1e}",
      " dV:", f"{np.abs(num_grad(V) - dV).max():.1e}")
print("soma de cada linha de dS:", np.round(dS.sum(axis=-1), 12),
      "— sempre zero: somar a mesma constante a uma linha de scores não muda o softmax")

pesos de atenção A (cada linha soma 1):
 [[0.226 0.475 0.298]
 [0.144 0.598 0.259]
 [0.668 0.235 0.097]]
dQ:
 [[0.456 0.354 0.099 0.209]
 [0.407 0.404 0.013 0.129]
 [0.345 0.228 0.109 0.184]]
erro máximo contra diferenças finitas — dQ: 3.2e-10  dK: 6.1e-10  dV: 3.4e-10
soma de cada linha de dS: [ 0. -0. -0.] — sempre zero: somar a mesma constante a uma linha de scores não muda o softmax

← Embeddings · next step: Masks and heads →