Queries, keys and values

Think of a sequence of n tokens, each stored as a vector of size d_model. For token i, the question is: which other tokens carry information I need right now, and how much of each? Self-attention answers this with no fixed pattern. The answer comes from the content.

The lookup picture

Each token plays three roles:

• It asks with a query vector: "what am I looking for?"
• It offers a key vector: "here is what I contain."
• It provides a value vector: "here is the content to pass along."

Token i compares its query with every key. Keys that match well get large scores, and softmax turns the scores into weights. The output for token i is the weighted sum of all the value vectors.

From input to Q, K and V

Start from the input matrix X of shape n × d_model. Three learned matrices project it:

Q = X W_Q    # n × d_k   (queries)
K = X W_K    # n × d_k   (keys)
V = X W_V    # n × d_v   (values)

For a single head, d_k = d_v = d_model. The standard formula, from Attention Is All You Need, is:

Attention(Q, K, V) = softmax( Q Kᵀ / √d_k ) V

The shapes flow like this:

Why divide by √d_k

If the query and key entries are independent with mean 0 and variance 1, their dot product has variance d_k. For d_k = 64, the typical score has a standard deviation of about 8.

Large scores make softmax very peaked. Take two keys with scores [8, 0]:

softmax([8, 0]) = [0.9997, 0.0003]

Almost all weight lands on one key, and the gradient for the other key nearly vanishes. Now divide by √64 = 8, so the scores become [1, 0]:

softmax([1, 0]) = [0.731, 0.269]

The result is still peaked toward the first key, but the second key keeps a useful share and still gets gradient. That is the whole job of the √d_k factor.

Common mistakes

• Scaling by the wrong number. Use the head dimension d_k, not d_model, when you have several heads.
• Softmax over the wrong axis. Normalize over the keys for each query, which is the last axis of QKᵀ.
• Confusing the roles. Queries and keys decide the weights. Values decide what gets copied.

Check yourself

1. In softmax(QKᵀ / √d_k) V, what is the shape of the weight matrix for a sequence of 10 tokens? It is 10 × 10, one row of weights per query.
2. Why does dividing by √d_k help? It keeps scores from growing with dimension, so softmax does not become a hard argmax.

Key takeaways

• Each token asks with a query, offers a key, and provides a value.
• Weights come from query-key dot products, normalized with softmax.
• The output is a weighted sum of values, so every position can see every other one in one step.
• Dividing by √d_k keeps softmax in a useful range as dimension grows.

Further learning

Masks decide which pairs are allowed at all. The next page covers them: attention masks.