Queries, keys and values

Scoring and mixing

Now each token has a query, a key and a value. Attention runs the same three steps as in the last lesson: score, normalize, mix. What changes is which vector goes where.

Two letters keep track of positions. $i$ is the token doing the looking. $j$ is any token it might look at, itself included. Token $i$'s score for token $j$ is written $s_{ij}$:

In plain words: to score token $j$, take $i$'s query and $j$'s key. Multiply the matching numbers and add them up. That's the dot product. A high score means $j$ is what $i$ is looking for.

Next, a softmax turns the scores into weights. They're all positive, and they add up to 1. The output for token $i$ is a blend of every token's value, mixed by those weights:

The symbol $\sum$ (a Greek capital sigma) means "add up". Here it adds over every token $j$. The little $j$ on the softmax means it runs across all of token $i$'s scores together.

In plain words: turn the scores into shares that add up to 100%. Then build the output by taking that share of each token's value. Queries and keys decide who to listen to. Values decide what gets passed on.

Back to "warm jacket". To keep the numbers easy, let $W_K$ and $W_V$ leave vectors unchanged. (In a real model they reshape them too.) So "warm" has key [1, 0] and value [1, 0]. "jacket" has key [0, 1] and value [0, 1].

"warm" has the query [0, 2]. Its two scores are:

• against its own key [1, 0]: 0 × 1 + 2 × 0 = 0;
• against jacket's key [0, 1]: 0 × 0 + 2 × 1 = 2.

The softmax raises $e$, a fixed number about 2.718, to the power of each score. $e^0 = 1$ and $e^2$ is about 7.39. The total is 8.39. Divide each by the total to get the weights: 1 ÷ 8.39 ≈ 0.12 and 7.39 ÷ 8.39 ≈ 0.88. The output is 0.12 × [1, 0] plus 0.88 × [0, 1], which is [0.12, 0.88]. The new vector for "warm" is now mostly about the jacket.

"jacket" has a blank query, so both of its scores are 0. Equal scores give equal weights: 0.5 and 0.5.

import math

q_warm = [0, 2]                          # warm's query: "I want a noun"
keys = {"warm": [1, 0], "jacket": [0, 1]}
values = {"warm": [1, 0], "jacket": [0, 1]}

scores = {w: sum(a * b for a, b in zip(q_warm, k)) for w, k in keys.items()}
total = sum(math.exp(s) for s in scores.values())
weights = {w: math.exp(s) / total for w, s in scores.items()}
out = [sum(weights[w] * values[w][i] for w in values) for i in range(2)]

print("scores: ", scores)
print("weights:", {w: round(p, 2) for w, p in weights.items()})
print("output: ", [round(o, 2) for o in out])

Why separate projections matter

The last lesson used plain similarity: queries, keys and values were all the raw embeddings. That has two limits.

• It's symmetric. A's score for B always equals B's score for A. A dot product gives the same answer either way around.
• Tokens favor themselves. When vectors are of similar length, a token usually scores itself highest.

Language isn't symmetric. An adjective looks for the noun it describes. The noun doesn't need the adjective in the same way. Separate query and key projections allow that.

Here's a hand-built example with four words. Each word gets three features: is it an adjective, a noun, or neither?

import numpy as np

words = ["warm", "jacket", "for", "winter"]
X = np.array([[1.0, 0.0, 0.0],      # adjective
              [0.0, 1.0, 0.0],      # noun
              [0.0, 0.0, 1.0],      # neither
              [0.0, 1.0, 0.0]])     # noun

def attention(Q, K):
    s = Q @ K.T
    e = np.exp(s - s.max(axis=1, keepdims=True))
    return e / e.sum(axis=1, keepdims=True)

print("plain similarity (Q = K = X):")
print(np.round(attention(X, X), 2))

W_Q = np.array([[0, 3, 0],          # an adjective's query asks for "noun"
                [0, 0, 0],          # nouns and other words ask for nothing in particular
                [0, 0, 0]], dtype=float)
W_K = np.eye(3)                     # keys simply advertise each word's type
print("\nwith projections (Q = X W_Q, K = X W_K):")
print(np.round(attention(X @ W_Q, X @ W_K), 2))

Each printed table has one row per word, in order: warm, jacket, for, winter. A row shows how that word spreads its attention over the four words, and each row adds up to 1.

With plain similarity, "warm" gives its biggest share, 0.48, to itself. With the projections, "warm" sends 0.48 to each noun, "jacket" and "winter". Only 0.02 is left for itself. The other words have blank queries, so they spread their attention evenly: 0.25 each.

Nobody writes these matrices by hand in a real model. Training finds projections like this, because they help predict the next token.

All tokens at once

A real model handles every token, and several texts, in one go. The projections become linear layers. nn.Linear with bias=False is a learned matrix and nothing else.

(torch here is the course's in-browser teaching subset of PyTorch. The code is the same as real PyTorch. It runs on the CPU only, much more slowly, and its seeded random numbers differ.)

The sizes in the code are:

• B = 2: two texts processed together, called a batch;
• T = 5: five tokens in each text;
• d_model = 16: the numbers in each token's vector;
• d_k = 8: the numbers in each query, key and value.

import torch
import torch.nn as nn

torch.manual_seed(0)
B, T, d_model, d_k = 2, 5, 16, 8                 # 2 sequences, 5 tokens each, 16 features per token
x = torch.randn(B, T, d_model)

W_q = nn.Linear(d_model, d_k, bias=False)
W_k = nn.Linear(d_model, d_k, bias=False)
W_v = nn.Linear(d_model, d_k, bias=False)
q, k, v = W_q(x), W_k(x), W_v(x)                 # each (2, 5, 8)

scores = q @ k.transpose(-2, -1)                 # (2, 5, 8) @ (2, 8, 5) → (2, 5, 5): every query vs every key
weights = torch.softmax(scores, dim=-1)          # each row sums to 1
out = weights @ v                                # (2, 5, 5) @ (2, 5, 8) → (2, 5, 8)
print("q, k, v:", tuple(q.shape), "  scores:", tuple(scores.shape), "  out:", tuple(out.shape))
print("rows sum to 1:", torch.allclose(weights.sum(dim=-1), torch.ones(B, T)))
print("parameters in the three projections:", 3 * d_model * d_k)

Each projection turns 16 numbers per token into 8. So q, k and v each have the shape (2, 5, 8). k.transpose(-2, -1) turns each text's keys on their side, like .T in NumPy. That lets every query meet every key.

The scores come out as (2, 5, 5). For each text, that's a 5 × 5 grid with one score for every (query token, key token) pair. Its size depends on the number of tokens, not on the vector length. Each projection matrix holds 16 × 8 = 128 weights, so the three hold 384.

A soft search inside the model

This is the retrieval idea from Phase 1's RAG lessons, built into the model. RAG compares a query with an index of documents and returns the top few matches. Attention changes that in two ways.

• It's soft. It doesn't pick the top results. It blends every value, weighted by how well its key matched.
• It's differentiable. Every step is smooth math, so training can use gradients. That's how it tunes what gets searched for and what gets returned.

Still missing

Two details come in the next lessons. First, the scores get divided by $\sqrt{d_k}$, the square root of $d_k$. This stops the softmax from saturating, which means piling almost all the weight on one token, as the vectors get longer. Second, a mask stops tokens in a language model from attending to the future.