Self-Attention Deep Dive

Queries, keys, and values

From the same input matrix X, learned projections produce Q, K, and V. A query is compared with keys to compute attention scores; the resulting weights mix value vectors. The search analogy is helpful, but Q, K, and V are learned numerical features.

Scaled dot-product attention

Attention(Q,K,V) = softmax(Q K^T / sqrt(d_k)) V. Softmax operates across allowed key positions for each query. Scaling helps control the magnitude of dot products as the key dimension grows.

Causal masking

In a decoder, positions in the future are masked before softmax, typically using negative infinity for their scores. The probability assigned to a disallowed key becomes zero. An encoder can use a different visibility pattern.

Context and cost

Attention can connect distant tokens in one layer. Dense score matrices contain n squared entries for n tokens. Attention weights are not by themselves a faithful explanation of why a model produced an answer.

Worked example

A query [1,0] compares two keys. The more aligned key receives more weight; the output is a mixture of both value vectors.

Download lesson 05 Python example

from math import exp, sqrt

q = [1.0, 0.0]
keys = [[1.0, 0.0], [0.0, 1.0]]
values = [[10.0, 0.0], [0.0, 10.0]]
scores = [sum(a*b for a, b in zip(q, k))/sqrt(2) for k in keys]
e = [exp(s-max(scores)) for s in scores]
weights = [x/sum(e) for x in e]
out = [sum(w*v[j] for w, v in zip(weights, values)) for j in range(2)]
print("Weights:", [round(w, 3) for w in weights])
print("Output:", [round(x, 3) for x in out])

Expected output

Weights: [0.67, 0.33]
Output: [6.698, 3.302]

The weights are approximately 0.670 and 0.330. The output is approximately [6.698, 3.302]. If only the first key were allowed, the weights would be [1,0].

Practice and self-check

Common mistake

Without position information, unmasked self-attention is permutation-equivariant, not simply permutation-invariant. Reordering inputs reorders corresponding outputs.

Student tasks

1. Replace the query with [0,1] and predict the new output before running.
2. Implement a mask that permits only the first key.
3. For 100 and 200 tokens, compare the number of entries in a dense attention matrix.