Attention: queries, keys, values

A dot product scores how well a query matches a key

A token is one item's vector, such as the embedding of a word or of a point. Attention needs a score for "how well does this query match that key". It uses the dot product: multiply the two vectors entry by entry and add the results.

q · k = q₁k₁ + q₂k₂ + … + q_dk_d

For example, [1, 2] · [1, 1] = 1 × 1 + 2 × 1 = 3. The dot product is large when the two vectors point the same way, zero when they are at right angles, and negative when they point apart. It also grows with the length of either vector. Each entry of a matrix times a vector is a dot product of one row with the vector, so you have met this sum before.

One token's turn: score, scale, softmax, weighted sum

Follow one token as it attends to three tokens. Its query is q = [1, 2]. The three tokens have the keys [1, 0], [0, 1] and [1, 1], and the values [10, 0], [0, 10] and [5, 5]. Each key has 2 entries, so the key size is d_k = 2.

Predict first. Which key matches the query best? Which value will the result lean toward?

A longer query picks more sharply

The figure below runs the same four steps for four tokens, A to D. The left panel shows their keys and the query. The right panel shows their values and the result, which is the weighted average of the values. Every dot is a 2D vector, drawn from the center.

All tokens at once: one line of matrix algebra

In self-attention, every token computes its own query, key and value from its vector x, using three learned matrices:

q = W_Qx,   k = W_Kx,   v = W_Vx

Training sets these matrices, so the network decides what a query looks for and what a key offers. Every token then takes its turn, as in the worked example. Stack the n queries as the rows of a matrix Q, and the keys and values likewise as K and V. All the turns together are then exactly

Attention(Q, K, V) = softmax( QK^T / √d_k ) V

Here K^T is K turned on its side (its transpose), so QK^T is the table of every query's dot product with every key, with n rows and n columns. Softmax works on each row separately, so each token's weights add up to 1. Multiplying by V forms each token's weighted sum of values. The output has one row per token.

attention.py

import numpy as np

def softmax(s):
    # Subtracting the max first avoids overflow in exp.
    # axis=-1 works along each row. keepdims=True keeps the results in a column,
    # so that each row uses its own max and its own sum.
    e = np.exp(s - s.max(axis=-1, keepdims=True))
    return e / e.sum(axis=-1, keepdims=True)

def attention(Q, K, V):
    d_k = Q.shape[-1]
    weights = softmax(Q @ K.T / np.sqrt(d_k))       # one row of weights per query
    return weights @ V, weights

# The worked example: one query, three keys and values.
Q = np.array([[1.0, 2.0]])
K = np.array([[1.0, 0.0], [0.0, 1.0], [1.0, 1.0]])
V = np.array([[10.0, 0.0], [0.0, 10.0], [5.0, 5.0]])
out, w = attention(Q, K, V)
print(np.round(w, 3), np.round(out, 2))

# Self-attention: 5 tokens with 4 numbers each, as the rows of X.
rng = np.random.default_rng(0)
X = rng.normal(size=(5, 4))
Wq, Wk, Wv = rng.normal(size=(3, 4, 4))         # three random 4 x 4 matrices
def self_attention(X):
    return attention(X @ Wq, X @ Wk, X @ Wv)[0]  # [0]: the output, not the weights

perm = [3, 0, 4, 1, 2]                          # a new order for the tokens
# [perm] reorders the rows. np.allclose is True if two arrays match up to rounding.
print(np.allclose(self_attention(X)[perm], self_attention(X[perm])))

[[0.14  0.284 0.576]] [[4.28 5.72]]
True

In code, tokens are the rows of a matrix X, so the queries are X @ Wq, where Wq is W_Q turned on its side (its transpose). Each row of X @ Wq is then one token's query W_Qx. The last line is explained two sections below.

Dividing by √d_k keeps the softmax from saturating

Real models use keys of 32 to 128 numbers, not 2. A dot product adds up d_k products, so its typical size grows with d_k. If the entries of q and k are independent random numbers with mean 0 and spread (standard deviation) 1, their dot product has mean 0 and spread √d_k. Dividing by √d_k brings the spread back to 1, whatever the key size.

scaling.py

import numpy as np

rng = np.random.default_rng(0)
q = rng.normal(size=(10000, 64))     # 10,000 random queries, d_k = 64
k = rng.normal(size=(10000, 64))
scores = (q * k).sum(axis=1)         # 10,000 dot products
print(round(scores.std(), 2), round((scores / 8).std(), 2))

8.08 1.01

Why does the spread matter? Softmax of the scores [8, 0, −8] is [0.99966, 0.00034, 0.0000001]: almost all the weight goes to one token. Softmax of [1, 0, −1] is [0.665, 0.245, 0.090]. When softmax is that lopsided we say it saturates. Small changes to the scores then barely change the weights, so the training signal that should adjust the queries and keys becomes tiny, and learning stalls.

Without positions, attention treats its tokens as a set

Nothing in the formula looks at where a token sits in the list. Shuffle the input tokens, and each token still meets the same keys and values, so it gets the same output. The outputs come out shuffled in the same way. The last line of attention.py checks this. In the terms of the permutation invariance unit, self-attention is permutation equivariant.

For a point cloud, this is what we want: the points have no natural order, and each point's coordinates already say where it is. For a sentence, the order of the words matters. So language models give each token information about its position, for example by adding Fourier features of the word's position to its embedding.

Multi-head attention runs several attentions side by side

One set of weights per token can follow only one notion of "match". Multi-head attention runs h attentions, called heads, in parallel. Each head has its own W_Q, W_K and W_V, which make vectors with fewer entries: usually d_k = d / h, where d is the number of entries in a token. Each head computes its own weights and its own weighted sum. The h results are joined end to end into one vector of d entries and multiplied by one more learned matrix, W_O.

For example, the original transformer used tokens of 512 numbers and 8 heads with keys of 64 numbers. The total work is about that of one head with keys of 512 numbers. In a model of points, one head may learn to weight nearby points and another to weight points with similar features. Nothing forces a head to specialize; training finds whatever helps.