Which earlier words matter for this one?
Attention is the step in a transformer where each token looks at the other tokens and takes in what it needs from them. It is why a language model can link "it" to the right noun many words back.
In "the bus is late", the word late is about the bus. Attention is how the model works that out with numbers.
By the end of this lesson
Note
I can work out the output of attention for one token by hand.
Key words
| Word | What it means | Example |
|---|---|---|
| Query | a vector saying what a token is looking for | |
| Key | a vector saying what a token offers | |
| Value | a vector holding what a token passes on | |
| Score | how well a query matches a key | 2 |
| Softmax | turns scores into shares that add up to 1 | |
| Weight | the share of a value that goes into the output | 0.665 |
The cells students fill in are left blank in this preview.
Every token has a query, key and value
In linear algebra, a table of numbers can multiply a vector to give a new vector. Attention multiplies each input vector by three learnt tables, which gives three new vectors.
The query says what this token is looking for. The key says what it offers to others. The value is what it passes on if chosen. Below, we start from the three vectors already made.
Worked example: scoring the query of late
| Token | Key | Query of late times key | Score |
|---|---|---|---|
| bus | 2, 1 | 1x2 + 0x1 | 2 |
| is | 0, 2 | 1x0 + 0x2 | 0 |
| late | 1, 1 | 1x1 + 0x1 | 1 |
The cells students fill in are left blank in this preview.
Softmax turns scores into weights
Scores can be any size, even below 0. Softmax turns them into weights that are all positive and add up to 1.
First take e to the power of each score, the exp button on a calculator. That makes every number positive, and a big score much bigger. Then divide each one by the total. Real models also divide every score by a fixed number first. We skip that here.
Worked example: softmax of 2, 0 and 1
| Token | Score | e to the score | Weight: divide by 11.11 |
|---|---|---|---|
| bus | 2 | 7.39 | 0.665 |
| is | 0 | 1.00 | 0.090 |
| late | 1 | 2.72 | 0.245 |
| Total | 11.11 | 1.000 |
The cells students fill in are left blank in this preview.
Scores and weights with NumPy
docker exec jupyter python3 -c "import numpy as np; K = np.array([[2, 1], [0, 2], [1, 1]]); s = K @ np.array([1, 0]); w = np.exp(s) / np.exp(s).sum(); print(s.tolist(), w.round(3).tolist())"Shown here, not run. Create a free account to press Run and type it into your own lab’s terminal.
The output is a weighted mix of values
Multiply each token's value by its weight, then add the results. That sum is the new vector for late.
Late took two thirds of its new vector from bus, so late now carries some of what bus means. Every token does the same with its own query at the same time.
Worked example: mixing the values
| Token | Weight | Value | Weight times value |
|---|---|---|---|
| bus | 0.665 | 10, 0 | 6.65, 0 |
| is | 0.090 | 0, 10 | 0, 0.90 |
| late | 0.245 | 5, 5 | 1.22, 1.22 |
| Output | about 7.88, 2.12 |
The cells students fill in are left blank in this preview.
Attention for one token in Python
import math
keys = [[2, 1], [0, 2], [1, 1]] # bus, is, late
values = [[10, 0], [0, 10], [5, 5]]
query = [1, 0] # the query of late
scores = []
for key in keys:
scores.append(query[0] * key[0] + query[1] * key[1])
powers = []
for s in scores:
powers.append(math.exp(s))
total = sum(powers)
weights = []
for p in powers:
weights.append(p / total)
mix = [0, 0]
for i in range(len(values)):
for j in range(2):
mix[j] = mix[j] + weights[i] * values[i][j]
print("scores:", scores)
print("weights:", [round(w, 3) for w in weights])
print("output:", [round(m, 2) for m in mix])
# Change the query to [0, 1] and press Run. Which token gets the most weight now?
A common mistake
Warning
Many people think attention picks one token and ignores the rest. It takes a share of every token's value. A weight is near 0 only when that key matches the query badly, and the weights always add up to 1.
Attention for all three tokens at once
docker exec jupyter python3 -c "import numpy as np; Q = np.array([[0, 1], [1, 1], [1, 0]]); K = np.array([[2, 1], [0, 2], [1, 1]]); V = np.array([[10, 0], [0, 10], [5, 5]]); S = Q @ K.T; W = np.exp(S) / np.exp(S).sum(axis=1, keepdims=True); print(W.round(3).tolist()); print((W @ V).round(2).tolist())"Shown here, not run. Create a free account to press Run and type it into your own lab’s terminal.