Self-Attention

Your program will implement the operation the whole architecture is named after: given a query, key and value vector for each token, return for each token a blend of everyone's values, weighted by how well its query matches their keys.

A soft dictionary lookup

A Python dict returns the one value whose key equals your query. Attention returns all the values at once, each weighted by how well its key matches the query, and the weights add up to 1. Each token plays three roles, and three learned linear maps of its embedding produce the three vectors:

• Query: what I am looking for.
• Key: what I advertise to others.
• Value: what I hand over if someone attends to me.

In this lesson Q, K and V arrive as given matrices, each with T rows (tokens) and d columns. Producing them is three matrix multiplications that lesson 10 does with seeded weights.

Three steps per query row

For token i:

1. Score it against every key: score[i][j] = dot(Q[i], K[j]) / sqrt(d). A large dot product means "this key matches what I want".
2. Softmax the T scores of that row into weights that are positive and sum to 1.
3. Mix: out[i] = sum over j of weight[i][j] * V[j].

In matrix form the whole thing is softmax(Q K^T / sqrt(d)) V, where the softmax runs along each row.

Worked trace, T = 2, d = 2, Q = [[2,0],[0,2]], K = [[1,0],[0,1]], V = [[10,0],[0,10]]:

row 0 scores:  q0.k0 = 2, q0.k1 = 0   ->  divided by sqrt(2):  [1.4142, 0]
softmax:       e^1.4142 = 4.113, e^0 = 1   ->  weights [0.8044, 0.1956]
output row 0 = 0.8044 * [10,0] + 0.1956 * [0,10] = [8.0443, 1.9557]
row 1 by symmetry: weights [0.1956, 0.8044] -> output [1.9557, 8.0443]

Token 0's query lines up with key 0, so most of what it receives is value 0, but it still gets a share of value 1. Every output row is a weighted average of the V rows, so it always lies between them.

Why divide by sqrt(d)

If the entries of q and k are independent with variance 1, their dot product is a sum of d terms and has variance d. The scores therefore grow like sqrt(d). Large scores push softmax into saturation: one weight near 1, the rest near 0, and the gradient through softmax vanishes. Dividing by sqrt(d) brings the variance back to 1 whatever d is.

Stable softmax

exp overflows a double near 710, so scores of a few hundred make a naive softmax return infinity divided by infinity. Softmax does not change if every score in a row shifts by the same amount, so subtract the row maximum first. The largest term becomes exp(0) = 1; nothing overflows.

Common mistakes

• Softmax along the wrong axis. Row i of the weights must sum to 1 across the keys; normalising each column instead gives plausible numbers that are wrong whenever Q differs from K.
• Dividing by d instead of sqrt(d), or forgetting to divide at all.
• Skipping the max shift and watching a test with large scores return NaN.
• Mixing the keys instead of the values in the last step.

Your exercise

Implement the row softmax and the attention function over given Q, K, V matrices, no masking yet; lesson 8 adds that. The starter parses the Q, K and V lines and prints the result.