Transformers, lesson 2
Attention as a soft lookup
Queries, keys and values let each token gather what it needs from the others.
The last lesson mixed each word with its neighbours, using scores it made up. This lesson shows where real scores come from: every word asks for something, every word offers something to match, and a dot product measures the match. You need the dot product and softmax.
By the end, you can:
Part 1. A lookup that blends
Look up a value with a query vector, and get back a blend.
A Python dictionary looks things up exactly. With prices = {"apple": 3, "pear": 5, "plum": 4}, prices["apple"] gives 3, and prices["aple"] raises an error. Now make it soft: give each key a vector, ask with a vector too, and let every entry answer a little.
Example 1.1
A soft lookup
Apple's key is (1, 0), pear's is (0, 1) and plum's is (0.7, 0.7). Their values are 3, 5 and 4. You ask with the vector (2, 0). What comes back?
- Score apple and pear2 × 1 + 0 × 0 = 2, 2 × 0 + 0 × 1 = 0
A dot product with each key: multiply number by number, then add. Apple matches the query, and pear doesn't at all.
- Score plum2 × 0.7 + 0 × 0.7 = 1.4
Plum's key points partly the same way as the query.
- Softmax the scorese^2 = 7.389, e^0 = 1, e^1.4 = 4.055; total 12.444
Divide each by the total: the weights are 0.59, 0.08 and 0.33.
- Blend the values0.59 × 3 + 0.08 × 5 + 0.33 × 4 = 3.49
Mostly apple's value, with some plum mixed in. Nothing errors.
Answer3.49.
The dot product does the matching. A key pointing the same way as the question scores high, and one at right angles scores 0. Ask with a vector halfway between two keys and you get something between their values.
Query, key and value
The three roles in a soft lookup. The query is what you're asking for. Each key is what an entry offers to be matched against. Each value is what that entry hands back. Q, K and V for short.
answer = w₁ × v₁ + w₂ × v₂ + w₃ × v₃, with w = softmax(q · k)
A soft lookup never fails, and that's also its weak spot: it never says "not found". A query that matches no key well still gets an answer, a blend of everything.
Your turn
Same three entries, but the question is (0, 2). Roughly what comes back?
Part 2. Three vectors from one word
Make a word's query, key and value with three learned grids.
In a sentence, every word plays all three roles at once. It sends out a query (what it's looking for), offers a key (what it can be found by) and carries a value (what it passes on). All three come from the word's own vector, each through its own grid of learned weights.
Example 2.1
"it" makes its query, key and value
The word "it" has the vector x = (0.6, 0.8). The query grid W_Q has the rows (0, 2) and (2, 0). The key grid W_K has the rows (2, 0) and (0, 1). The value grid W_V leaves a vector as it is. What are the query, key and value?
- Make the query(0.6 × 0 + 0.8 × 2, 0.6 × 2 + 0.8 × 0) = (1.6, 1.2)
Each number is x dotted with one column of the grid. W_Q swaps the two numbers and doubles them.
- Make the key(0.6 × 2 + 0.8 × 0, 0.6 × 0 + 0.8 × 1) = (1.2, 0.8)
W_K doubles the first number and keeps the second.
- Make the value(0.6, 0.8)
W_V leaves the vector as it is, to keep the example small.
AnswerQuery (1.6, 1.2), key (1.2, 0.8), value (0.6, 0.8).
Why not use the word's own vector for all three? A vector dotted with itself is its length squared, and it points exactly its own way, so a word would often match itself best. Separate grids let a word ask for something unlike itself: a verb can look for its subject, and "it" can look for a noun.
Projection
Multiplying a vector by a grid of weights to get a new vector, as in "the query projection". W_Q, W_K and W_V are learned in training like every other weight, and every word goes through the same three grids.
q = x W_Q, k = x W_K, v = x W_V
The grids belong to the layer, not to the word. What makes two words' queries differ is only their vectors. A real model's grids are much bigger, 768 by 768 in GPT-2 small, but each number comes out the same way.
Your turn
Another word has the vector (1, 0.5). It goes through the same W_Q, with rows (0, 2) and (2, 0). What is its query?
Part 3. Every word looks at every word
Work out one word's row of attention weights in a sentence.
Now run the lookup inside a sentence, "cat sat it". Every word's query is scored against every word's key, its own included. Softmax turns each word's scores into weights, and each word's new vector is the weighted blend of the values. Here's the row for "it".
Example 3.1
The row for "it"
The keys are cat (2, 0), sat (0, 1) and it (1.2, 0.8). The values are the word vectors: cat (1, 0), sat (0, 1) and it (0.6, 0.8). The query for "it" is (1.6, 1.2). What new vector does "it" get?
- Score cat and sat1.6 × 2 + 1.2 × 0 = 3.2, 1.6 × 0 + 1.2 × 1 = 1.2
The query for "it" against the keys of cat and sat.
- Score itself1.6 × 1.2 + 1.2 × 0.8 = 2.88
A word scores its own key too.
- Softmax the scorese^3.2 = 24.533, e^1.2 = 3.320, e^2.88 = 17.814
The total is 45.667, so the weights are 0.54, 0.07 and 0.39.
- Blend the values(0.54 + 0.39 × 0.6, 0.07 + 0.39 × 0.8) = (0.77, 0.38)
Each value times its weight, one number at a time. Cat's (1, 0) only adds to the first number, sat's (0, 1) to the second.
Answer(0.77, 0.38), more than half cat.
Each row is one word's soft lookup over the whole sentence. The rows add up to 1, so every new vector is a blend of the values, weighted by how well each key answered that word's query. All the rows come out of a few grid multiplications at once.
Self-attention
Attention where the queries, keys and values all come from the same sentence, so every word looks at every word, itself included. The weights for every pair of words form a square grid, the attention matrix, whose rows each add up to 1.
n words → an n × n grid of weights
In a trained model the grids are learned, and some end up doing jobs you can name, such as linking "it" to the noun it stands for. The grids here were picked by hand to show that. "What trained heads learn", later in this module, shows real ones.
Your turn
Work out the row for "sat". Its query is (2, 0), and the keys are cat (2, 0), sat (0, 1) and it (1.2, 0.8). What weights does it give cat, sat and it?
Where you'll meet it
A common coding task for ML roles: write one attention head in NumPy. It takes a few lines: Q, K and V from three matrix products, the scores Q @ K.T, a softmax along each row, then the weights times V. Be ready to say why W_Q and W_K are separate: so a word can look for something unlike itself.
Open nanoGPT, Andrej Karpathy's small GPT, and you'll find the three grids stored as one: c_attn = nn.Linear(n_embd, 3 * n_embd). One matrix product makes a long vector for each token, and .split(n_embd, dim=2) cuts it into q, k and v. It's the same arithmetic as three grids, in one call.
import numpy as np
X = np.array([[1.0, 0.0], [0.0, 1.0], [0.6, 0.8]]) # cat, sat, it
Wq = np.array([[0.0, 2.0], [2.0, 0.0]])
Wk = np.array([[2.0, 0.0], [0.0, 1.0]])
Q, K, V = X @ Wq, X @ Wk, X # W_V leaves X as it is
scores = Q @ K.T # every query · every key
weights = np.exp(scores) / np.exp(scores).sum(axis=1, keepdims=True)
print(weights.round(2))
print((weights @ V)[2].round(2)) # the new vector for "it"
# [[0.07 0.55 0.37]
# [0.82 0.02 0.17]
# [0.54 0.07 0.39]]
# [0.77 0.38]Q @ K.T is every dot product at once, a 3 × 3 grid of scores, and keepdims=True makes each row divide by its own total. The next lesson adds one division to the scores.
Check yourself
Three questions from across the lesson. Try each one before you look back.
1
You skip the grids and use each word's own vector as its query and its key. In "cat sat it", which word does each one give the most weight?
2
A sentence has 6 tokens. How many weights are in its attention matrix, and what does each row add up to?
3
A soft lookup has the values 3, 5 and 4, and the query scores 0 against every key. What comes back?
What to remember
1A lookup that blends
A soft lookup scores a query against every key, softmaxes the scores and blends the values: the query (2, 0) gets 3.49 from the values 3, 5 and 4.
2Three vectors from one word
Each word makes its query, key and value with the same three learned grids, as q = x W_Q: the vector (0.6, 0.8) gives the query (1.6, 1.2).
3Every word looks at every word
In self-attention every word gives a row of weights to every word, itself included, so n words make an n × n grid whose rows each add up to 1.
Recap and going further
The lesson on one page
- Score the query against every key
- Softmax the scores, then blend the values
- Query: what a word looks for
- Key: what it can be found by
- Value: what it passes on
- q = x W_Q, k = x W_K, v = x W_V
- So a word can ask for something unlike itself
- Every word scores every word, itself included
- An n × n grid, each row adding up to 1
- Each word's new vector is a weighted blend
- "it" gets pulled towards cat: (0.77, 0.38)
Try this
Write attention and check it against PyTorch
In PyTorch, give a 4-word sentence random vectors of 8 numbers, and make three random 8 × 8 grids for W_Q, W_K and W_V. Work out the scores, the softmax weights and the blended outputs yourself. Then call F.scaled_dot_product_attention on the same queries, keys and values and compare. Last, use each word's own vector as its query and key, with no grids, and look at the weights again.
You need: Python with PyTorch (pip install torch), or a free notebook such as Google Colab.
Look for
- Every row of weights adding up to 1.
- Whether your output matches PyTorch's, and the one change it takes to make them match.
- How much weight each word gives itself when queries and keys are the word vectors themselves.
- How the weights shift when you double one word's key.
The words, in plain language
- Query, key and value
- The three roles in a soft lookup. The query is what you're asking for. Each key is what an entry offers to be matched against. Each value is what that entry hands back. Q, K and V for short.
- Projection
- Multiplying a vector by a grid of weights to get a new vector, as in "the query projection". W_Q, W_K and W_V are learned in training like every other weight, and every word goes through the same three grids.
- Self-attention
- Attention where the queries, keys and values all come from the same sentence, so every word looks at every word, itself included. The weights for every pair of words form a square grid, the attention matrix, whose rows each add up to 1.
Read next
- The Illustrated TransformerJay AlammarQueries, keys and values drawn as vectors and grids, one step at a time.
- Neural Machine Translation by Jointly Learning to Align and TranslateBahdanau, Cho and Bengio, 2014The paper that introduced attention, as a way for a translator to look back at the source sentence.
- Attention? Attention!Lilian WengA tour of attention's versions, from translation to transformers.
Transformers overviewSpotted a mistake? Tell us.