Worked numeric example of single-head self-attention

What's actually being asked

Walk through a tiny, fully worked numeric example of single-head scaled dot-product self-attention: compute Q, K, V, the raw scores, the scaled scores, the softmax probabilities, and the final output.

Answer

Setup

Take a toy sequence of 2 tokens with embedding dimension d_model = 2, and d_k = d_v = 2 (single head). Use one-hot token embeddings as rows of X:

X = [[1, 0], [0, 1]] (token 1 = [1,0], token 2 = [0,1])

Weight matrices: W_Q = [[1, 0], [0, 1]] (identity) W_K = [[1, 0], [0, 2]] W_V = [[1, 0], [0, 1]] (identity)

Step 1 — Q, K, V

Q = XW_Q, K = XW_K, V = XW_V. Since X = I (identity), each product equals the corresponding weight matrix:

Q = [[1, 0], [0, 1]] K = [[1, 0], [0, 2]] V = [[1, 0], [0, 1]]

(Row 1 of each is token 1's Q/K/V vector, row 2 is token 2's.)

Step 2 — raw attention scores, QK^T

scores = Q K^T

Row 1 (query = token 1 = [1,0]): dot with K's rows [1,0] and [0,2] → [1·1+0·0, 1·0+0·2] = [1, 0] Row 2 (query = token 2 = [0,1]): dot with K's rows → [0·1+1·0, 0·0+1·2] = [0, 2]

scores = [[1, 0], [0, 2]]

Step 3 — scale by √d_k

d_k = 2, so √d_k = √2 ≈ 1.4142.

scaled = [[1/1.4142, 0/1.4142], [0/1.4142, 2/1.4142]] = [[0.7071, 0], [0, 1.4142]]

Step 4 — softmax (row-wise)

Row 1: softmax([0.7071, 0]) e^0.7071 ≈ 2.0281, e^0 = 1, sum = 3.0281 → [2.0281/3.0281, 1/3.0281] = [0.6698, 0.3302]

Row 2: softmax([0, 1.4142]) e^0 = 1, e^1.4142 ≈ 4.1133, sum = 5.1133 → [1/5.1133, 4.1133/5.1133] = [0.1956, 0.8044]

Token 1 ends up attending mostly to itself (0.670 vs 0.330 to token 2). Token 2 also attends mostly to itself, but even more strongly (0.804 vs 0.196) — because W_K scaled token 2's key up, making it a stronger match for token 2's own query.

Step 5 — weighted output = softmax(scores) · V

V = [[1, 0], [0, 1]]

Output row 1 = 0.6698·[1,0] + 0.3302·[0,1] = [0.6698, 0.3302] Output row 2 = 0.1956·[1,0] + 0.8044·[0,1] = [0.1956, 0.8044]

Attention(Q,K,V) = [[0.6698, 0.3302], [0.1956, 0.8044]]

Each output row is a weighted blend of the two value vectors, weighted exactly by that row's softmax probabilities — this is the self-attention output: every token's new representation is a context-aware mix of all tokens' value vectors.

Exam tip

Show every intermediate matrix (Q, K, V, raw scores, scaled scores, softmax, final output) explicitly, in that order — examiners give partial credit per step, and skipping straight to softmax probabilities without showing the QK^T and /√d_k steps is the easiest way to lose marks on an otherwise-correct final answer.