Self-Attention Workbook: 10 Exercises

Reading the attention equation and executing it are different skills, and only the second one sticks. This workbook makes you compute self-attention by hand on a three-token sequence with numbers chosen to be trustworthy, then verify your arithmetic in NumPy. Part one warms up on scores and asymmetry, part two runs the full mechanism end to end, and part three covers the two things production models add: masking and the cost that shapes the whole field.

Every exercise carries its full solution, so this works for self-study. A calculator is enough for exercises 1 to 8; Python is needed only for the last two. For the derivations behind what you are computing, read the mathematics behind self-attention first.

The setup, used throughout

A sequence of three tokens with key dimension $d_k = 4$ and value dimension 2. The projections have already been applied, so you are handed the queries, keys, and values directly:

Row i of Q is token i’s query, row j of K is token j’s key, and row j of V is its value. Keep this page open; every exercise refers back to it.

Part one: scores

Exercise 1: the score matrix

Compute all nine raw $scores s_{ij} = q_i \cdot k_j$scores and assemble $S = QK^\top$.

Solution. Each entry is a four-term dot product, and the zeros do most of the work:

For instance $s_{13} = q_1 \cdot k_3 = 2 \times 2 = 4, and s_{32} = q_3 \cdot k_2 = 2 \times 1 + 2 \times 1 = 4$. Row i reads as token i’s raw relevance judgement over the sequence: token 1 finds token 3 most relevant, while tokens 2 and 3 both point at token 2.

Exercise 2: asymmetry

Compare $s_{13}$ with $s_{31}$. What does the pair say about tokens 1 and 3, and why would this be impossible if scores were computed directly between embeddings as $x_i \cdot x_j$?

Solution. From the matrix, $s_{13} = 4$ while $s_{31} = 0$: token 1 attends strongly to token 3, and token 3 ignores token 1 completely. Relevance runs one way here, like “keys” needing “piano” while “piano” needs nothing back. Raw embedding dot products cannot express this, because $x_i \cdot x_j = x_j \cdot x_i$ forces symmetry on every pair. The separate $W_Q$ and $W_K$ projections are precisely what break it, since $q_1 \cdot k_3$ and $q_3 \cdot k_1$ involve different vectors and are free to disagree.

Exercise 3: scale the scores

Divide S by $\sqrt{d_k}$. Then state what the variance of a raw score would be if query and key components were independent with mean 0 and variance 1, and what the division does to it.

Solution. With $d_k = 4$ the divisor is 2:

A raw score is a sum of $d_k$ independent component products, and variances of independent terms add, so $\operatorname{Var}(q_i \cdot k_j) = d_k$. Dividing the score by $\sqrt{d_k}$ divides its variance by $d_k$, returning it to 1. The scaling is a variance correction, which is why it is a square root and why $d_k$ and nothing else appears in it.

Part two: the mechanism

Exercise 4: softmax by hand

Apply the softmax to row 1 of the scaled matrix, using $e^0 = 1$, $e^1 \approx 2.718$, $e^2 \approx 7.389$. Check that the weights sum to 1.

Solution. Row 1 is $[1, 0, 2]$. Exponentiate to get [2.718, 1, 7.389], which sums to 11.107. Dividing through:

The weights sum to 1.000, so they form a probability distribution over the three tokens. Notice the amplification: a scaled-score gap of 2 between tokens 3 and 2 became a weight ratio above 7 to 1. Softmax sharpens moderate score differences into decisive weights while keeping every weight positive and every gradient alive.

Exercise 5: the output vector

Using those weights, compute token 1’s output $z_1 = \sum_j w_{1j} v_j$.

Solution. A weighted average of the three value rows:

Token 1’s new representation is two-thirds token 3’s value, a quarter its own, and a sliver of token 2’s. That is the mechanism’s entire output: a context-aware vector built as a content-weighted blend of what the other tokens offer.

Exercise 6: complete the other rows

Softmax scaled rows 2 and 3 and compute $z_2$ and $z_3$.

Solution. Row 2 is [0, 1, 0]: exponentials [1, 2.718, 1] sum to 4.718, giving weights [0.212, 0.576, 0.212] and

Row 3 is [0, 2, 0]: exponentials [1, 7.389, 1] sum to 9.389, giving weights [0.107, 0.787, 0.107] and

Compare the two rows. Both attend mainly to token 2, but token 3’s higher score produces a far harder commitment, 0.787 against 0.576, the exponential again converting a score gap into a weight gap.

Exercise 7: saturation, or why the scaling matters

Multiply row 1’s scaled scores by 5, giving [5, 0, 10], and softmax again with $e^5 \approx 148.4$ and $e^{10} \approx 22026$. What happens, and why does it break learning?

Solution. The exponentials are [148.4, 1, 22026], summing to 22175, so the weights are

Essentially all mass sits on token 3: the softmax has saturated into a hard, one-token selection. This is what unscaled scores do at realistic dimensions, since raw scores grow like $\sqrt{d_k}$. Two failures follow. The token loses its blended context, and the gradient of a saturated softmax is near zero, so the attention weights stop updating. The $\sqrt{d_k}$ division exists to keep the mechanism out of exactly this regime.

Part three: what production adds

Exercise 8: the causal mask

A language model must not let position i see positions after it. Apply the causal mask by setting $s_{ij} = -\infty$ for $j > i$, then recompute all three rows of weights. What happens to row 1?

Solution. The masked scaled matrix is

Since $e^{-\infty} = 0$, masked positions receive exactly zero weight. Row 1 has only itself available, so its weights are [1, 0, 0] and $z_1 = v_1 = [1, 0]$: the first token can attend to nothing but itself, by construction. Row 2 softmaxes over [0, 1], giving [0.269, 0.731] and $z_2 = [0.269, 0.731]$. Row 3 has nothing after it, so it keeps its Exercise 6 result. Note how far $z_1$ and $z_2$ have moved from their unmasked versions: the mask genuinely changes what every non-final token knows.

Exercise 9: implement and verify

Write scaled dot-product attention as a function and confirm it reproduces your hand results from Exercises 1 to 6.

Solution.

import numpy as np
def attention(Q, K, V):
    d_k = Q.shape[1]
    scores = Q @ K.T / np.sqrt(d_k)                          # scaled scores
    weights = np.exp(scores)
    weights = weights / weights.sum(axis=1, keepdims=True)   # row-wise softmax
    return weights @ V, weights
Q = np.array([[2., 0, 0, 0], [0, 2, 0, 0], [0, 2, 2, 0]])
K = np.array([[1., 0, 0, 0], [0, 1, 1, 0], [2, 0, 0, 0]])
V = np.array([[1., 0], [0, 1], [1, 1]])
Z, W = attention(Q, K, V)
print("weights:\n", W.round(3))
print("outputs:\n", Z.round(3))

The output matches Exercises 4 to 6 exactly: weights [0.245, 0.090, 0.665], [0.212, 0.576, 0.212], [0.107, 0.787, 0.107], and outputs [0.910, 0.755], [0.424, 0.788], [0.214, 0.894]. The whole mechanism is four lines: a matrix product, a scale, a row-wise softmax, and another matrix product.

Exercise 10: mask it, then count the cost

Add a causal mask to the function and verify Exercise 8. Then answer: how many entries does the score matrix hold for 1,000 tokens, and for 2,000?

Solution.

import numpy as np
def causal_attention(Q, K, V):
    n, d_k = Q.shape[0], Q.shape[1]
    scores = Q @ K.T / np.sqrt(d_k)
    mask = np.triu(np.ones((n, n)), k=1).astype(bool)   # True above the diagonal
    scores[mask] = -np.inf                               # the future -> -infinity
    weights = np.exp(scores)
    weights = weights / weights.sum(axis=1, keepdims=True)
    return weights @ V, weights
Q = np.array([[2., 0, 0, 0], [0, 2, 0, 0], [0, 2, 2, 0]])
K = np.array([[1., 0, 0, 0], [0, 1, 1, 0], [2, 0, 0, 0]])
V = np.array([[1., 0], [0, 1], [1, 1]])
Z, W = causal_attention(Q, K, V)
print(W.round(3))    # row 1: [1, 0, 0];  row 2: [0.269, 0.731, 0]
print(Z.round(3))

The weights match Exercise 8 exactly. On the cost question: the score matrix holds $n^2$ entries, so 1,000 tokens produce 1,000,000 scores and 2,000 tokens produce 4,000,000. Doubling the context quadruples the attention cost, which is the quadratic constraint that makes long-context models an engineering achievement rather than a parameter change.

What you have now

Ten exercises in, the attention equation is no longer a formula you recognise but one you have executed: you scored a sequence, saw the asymmetry the projections buy, scaled by a constant you can derive, softmaxed by hand, blended values into context-aware vectors, masked the future, watched saturation kill the gradients, and reproduced all of it in four lines of NumPy.

See you soon.