Attention, one query at a time
Pick a token. Its query is compared with every key, the scores become weights, and the weights mix the values.
Drag the arrow tips. A key that points the same way as the query gets a high score.
| Key | q · k | ÷ √2 | softmax weight |
|---|---|---|---|
| -2.52 | -1.78 | 0% | |
| 5.88 | 4.16 | 76% | |
| 2.28 | 1.61 | 6% | |
| -0.96 | -0.68 | 1% | |
| -3.48 | -2.46 | 0% | |
| 3.78 | 2.67 | 17% |
qsat · kcat = 2.40 × 2.20 + 0.60 × 1.00 = 5.88
output = Σ weight × v = (1.54, 0.48)
All six queries at once: the attention weight matrix, softmax(QKᵀ/√d_k)
| query ↓ key → | The | cat | sat | on | the | mat |
|---|---|---|---|---|---|---|
| 0.01 | 0.62 | 0.11 | 0.02 | 0.01 | 0.23 | |
| 0.04 | 0.16 | 0.58 | 0.19 | 0.01 | 0.01 | |
| 0.00 | 0.76 | 0.06 | 0.01 | 0.00 | 0.17 | |
| 0.01 | 0.22 | 0.01 | 0.01 | 0.01 | 0.74 | |
| 0.01 | 0.31 | 0.03 | 0.01 | 0.02 | 0.62 | |
| 0.09 | 0.05 | 0.50 | 0.33 | 0.03 | 0.00 |
Shapes: Q, K, V are 6 × 2. QKᵀ is 6 × 6, one row per query. The output is 6 × 2, one mixed vector per token.
Query tokensat
Attends most tocat, 76%
Output vector(1.54, 0.48)
Challenges
0 of 4 complete
- With scaling on, make the chosen query give at least 95% of its weight to a single key. (not yet)
- Make the chosen query weigh all six tokens equally. (not yet)
- Make two different tokens get exactly the same weight from every query. (not yet)
- Without moving any query or key, make the output for sat point down and to the right (x above 0, y below 0). (not yet)
Try this
- Pick mat as the query. Which key does it favor, and what happens to the output when you drag that key’s value somewhere new?
- Drag two keys to exactly the same point. Their weights become equal for every query. (Same direction is not enough: a longer key gets a larger score.)
- Turn off scaling and compare the matrix with scaling on. Every row gets sharper.
Scaled dot-product attention in a few linesPython with NumPy
import numpy as np
tokens = ["The", "cat", "sat", "on", "the", "mat"]
# The demo's starting vectors: one row per token, two numbers each.
Q = np.array([[1.6, 0.4], [0.3, 2.0], [2.4, 0.6], [1.5, -1.2], [1.4, -0.8], [-0.5, 2.0]])
K = np.array([[-1.2, 0.6], [2.2, 1.0], [0.4, 2.2], [-0.8, 1.6], [-1.4, -0.2], [1.8, -0.9]])
V = np.array([[0.2, -1.0], [1.8, 0.8], [-0.4, 1.9], [-1.5, 0.3], [0.3, -1.1], [1.2, -1.4]])
def softmax(x):
e = np.exp(x - x.max(axis=-1, keepdims=True))
return e / e.sum(axis=-1, keepdims=True)
d_k = K.shape[1]
scores = Q @ K.T / np.sqrt(d_k) # 6 x 6: every query against every key
weights = softmax(scores) # each row sums to 1
output = weights @ V # each row is a weighted average of the values
i = tokens.index("sat")
print("sat attends to:", ", ".join(f"{t} {w:.0%}" for t, w in zip(tokens, weights[i])))
print("output for sat:", np.round(output[i], 2))
print("shapes:", scores.shape, output.shape)sat attends to: The 0%, cat 76%, sat 6%, on 1%, the 0%, mat 17%
output for sat: [1.54 0.48]
shapes: (6, 6) (6, 2)