Implementing One Attention Head in NumPy

Last week you built up the pieces: dot products as similarity scores, softmax as a "turn scores into a probability distribution" operator, and the idea that a token's new representation is a weighted average of other tokens' vectors. Today you turn that into code you can run, poke, and break.

The whole head is one formula:

$\text{Attention}(Q,K,V) = \operatorname{softmax}\!\left(\frac{QK^\top}{\sqrt{d_k}}\right)V$

Everything below is bookkeeping around that line.

1. Shapes first

Shape discipline is the single most useful habit here — and it's also the best way to explain attention to a colleague, because the shapes are the explanation.

Let a sequence be $n$ tokens, each already embedded into $d_{\text{model}}$ dimensions, stacked as rows of a matrix $X$.

| Object | Shape | What a row means | |---|---|---| | $X$ | $(n, d_{\text{model}})$ | token $i$'s current representation | | $W_Q, W_K$ | $(d_{\text{model}}, d_k)$ | learned projections | | $W_V$ | $(d_{\text{model}}, d_v)$ | learned projection | | $Q = XW_Q$ | $(n, d_k)$ | "what token $i$ is looking for" | | $K = XW_K$ | $(n, d_k)$ | "what token $j$ offers as a label" | | $V = XW_V$ | $(n, d_v)$ | "what token $j$ contributes if attended to" | | $QK^\top$ | $(n, n)$ | raw score of query $i$ against key $j$ | | $A = \operatorname{softmax}(\cdot)$ | $(n, n)$ | row $i$ = a probability distribution over all tokens | | $AV$ | $(n, d_v)$ | token $i$'s output = weighted average of value rows |

Two facts worth memorizing, because they answer 80% of a colleague's questions:

• The $(n,n)$ matrix is the only place tokens interact. Everything else is applied to each token independently.
• $A$'s rows sum to 1 by construction (softmax along the last axis), so each output row is a convex combination of value vectors — an average, never an extrapolation.

Why divide by $\sqrt{d_k}$?

If $q$ and $k$ have roughly independent, zero-mean, unit-variance entries, then $q\cdot k = \sum_{t=1}^{d_k} q_t k_t$ has variance $\approx d_k$, so typical magnitude $\approx \sqrt{d_k}$. For $d_k = 64$ that's scores spread over $\pm 8$ or more, and softmax of that is nearly one-hot: the head becomes a hard lookup with near-zero gradient for every non-winning token. Dividing by $\sqrt{d_k}$ puts the logits back on an $O(1)$ scale. It's variance control, nothing deeper.

2. The code

import numpy as np

def softmax(x, axis=-1):
    x = x - x.max(axis=axis, keepdims=True)      # stability: exp never overflows
    e = np.exp(x)
    return e / e.sum(axis=axis, keepdims=True)

def attention(Q, K, V, mask=None):
    d_k = Q.shape[-1]
    scores = (Q @ K.T) / np.sqrt(d_k)            # (n, n)
    if mask is not None:
        scores = np.where(mask, scores, -np.inf)  # mask==False -> weight 0
    A = softmax(scores, axis=-1)                 # (n, n), rows sum to 1
    return A @ V, A                              # (n, d_v), (n, n)

def self_attention(X, Wq, Wk, Wv, mask=None):
    return attention(X @ Wq, X @ Wk, X @ Wv, mask)

def causal_mask(n):
    return np.tril(np.ones((n, n), dtype=bool))  # token i may see j <= i

Fifteen lines of substance. Notes on the fiddly bits:

• x - x.max(...) shifts every logit in a row by the same constant. Softmax is invariant to that ($e^{x_i - c}/\sum_j e^{x_j - c}$ cancels $e^{-c}$), so it changes nothing mathematically but prevents exp from overflowing to inf.
• -np.inf in masked slots gives $e^{-\infty} = 0$, exactly zero weight. Combined with the max-subtraction this is safe as long as every row has at least one unmasked entry — a causal mask always keeps the diagonal, so you're fine. A fully masked row yields nan; that's your signal you built the mask wrong.
• No batch dimension, on purpose. To batch, Q @ K.T becomes np.einsum('bqd,bkd->bqk', Q, K) (or Q @ K.transpose(0, 2, 1)), and everything else is unchanged because the softmax and matmuls already act on the last axes.

3. Worked example you can check by hand

Skip the projections and feed $Q, K, V$ directly, with $n = 3$, $d_k = d_v = 2$:

$Q = K = \begin{bmatrix} 1 & 0 \\ 0 & 1 \\ 1 & 1 \end{bmatrix}, \qquad V = \begin{bmatrix} 10 & 0 \\ 0 & 10 \\ 5 & 5 \end{bmatrix}$

Step 1 — raw scores. $QK^\top$ is just all pairwise dot products:

$QK^\top = \begin{bmatrix} 1 & 0 & 1 \\ 0 & 1 & 1 \\ 1 & 1 & 2 \end{bmatrix}$

Row 3 has the largest self-score ($2$) because $q_3 = [1,1]$ has the largest norm — a reminder that these scores mix "direction agreement" and "magnitude."

Step 2 — scale by $1/\sqrt{2} \approx 0.7071$. Row 1 becomes $[0.7071,\ 0,\ 0.7071]$.

Step 3 — softmax row 1. $e^{0.7071} = 2.0281$, $e^{0} = 1$, sum $= 5.0562$:

$A_{1,\cdot} = \left[\tfrac{2.0281}{5.0562},\ \tfrac{1}{5.0562},\ \tfrac{2.0281}{5.0562}\right] = [0.4011,\ 0.1978,\ 0.4011]$

Sum: $0.4011 + 0.1978 + 0.4011 = 1.0000$. ✅

Row 2 is the mirror image, $[0.1978, 0.4011, 0.4011]$. Row 3 scales to $[0.7071, 0.7071, 1.4142]$, and $e^{1.4142} = 4.1133$ over a sum of $8.1695$ gives

$A_{3,\cdot} = [0.2483,\ 0.2483,\ 0.5035], \qquad \textstyle\sum = 1.0001 \text{ (rounding)}$

Step 4 — apply to $V$.

$AV \approx \begin{bmatrix} 6.02 & 3.98 \\ 3.98 & 6.02 \\ 5.00 & 5.00 \end{bmatrix}$

Read the story off the numbers: token 1's query aligned with keys 1 and 3, so its output leans toward $v_1 = [10,0]$ but gets dragged toward the middle. Token 3, whose query is equally aligned with everything, lands at the balanced point $[5,5]$. And notice every entry is between 0 and 10 — the min and max of $V$'s columns. That's the convex-combination property, and it's a free correctness check on any attention implementation.

Q = K = np.array([[1., 0.], [0., 1.], [1., 1.]])
V = np.array([[10., 0.], [0., 10.], [5., 5.]])
out, A = attention(Q, K, V)

assert np.allclose(A.sum(axis=1), 1.0)          # rows are distributions
assert (A >= 0).all()
assert out.shape == (3, 2)
assert out.min() >= V.min() - 1e-9 and out.max() <= V.max() + 1e-9
print(np.round(A, 4))
print(np.round(out, 3))

4. Verifying the row-sum property properly

One hardcoded example proves little. Test the invariant on random inputs, including a masked case:

rng = np.random.default_rng(0)
n, d_model, d_k, d_v = 7, 16, 4, 8

X  = rng.normal(size=(n, d_model))
Wq = rng.normal(size=(d_model, d_k)) / np.sqrt(d_model)
Wk = rng.normal(size=(d_model, d_k)) / np.sqrt(d_model)
Wv = rng.normal(size=(d_model, d_v)) / np.sqrt(d_model)

out, A = self_attention(X, Wq, Wk, Wv)
assert A.shape == (n, n) and out.shape == (n, d_v)
assert np.allclose(A.sum(axis=1), np.ones(n))

M = causal_mask(n)
out_c, A_c = self_attention(X, Wq, Wk, Wv, mask=M)
assert np.allclose(A_c.sum(axis=1), np.ones(n))     # still distributions
assert np.allclose(A_c[~M], 0.0)                    # nothing leaks from the future
assert np.isclose(A_c[0, 0], 1.0)                   # token 0 can only attend to itself

Note what the causal mask does not break: rows still sum to 1, because softmax renormalizes over whatever survives. Row $i$ is a distribution over exactly $i+1$ positions. This is the detail most people get wrong when explaining masking — they imagine weights being zeroed out after softmax, which would leave rows summing to less than 1. The masking happens to the logits, before the normalization.

A second probe worth running, since it explains why positional encodings exist at all:

P = rng.permutation(n)
out_p, _ = self_attention(X[P], Wq, Wk, Wv)
assert np.allclose(out_p, out[P])   # shuffle inputs -> outputs shuffle identically

Unmasked self-attention is permutation-equivariant: it has no idea what order the tokens came in. Position information has to be injected into $X$ (or into the scores) from outside.

5. Try it yourself

Starting from the code above:

1. Sharpening. Replace Wq with 10 * Wq and recompute A. Report A.max(axis=1) and the row entropy $H_i = -\sum_j A_{ij}\log A_{ij}$ before and after. Then delete the / np.sqrt(d_k) and rerun with d_k = 4 vs d_k = 256, watching A.max(axis=1) — you should be able to state in one sentence why the $\sqrt{d_k}$ is there.
2. Break the mask on purpose. Build a mask that is all False in row 0 and confirm you get nan. Then explain to yourself where the nan comes from, given the max-subtraction trick.
3. The two-sentence explanation. Write down, without looking at this lesson, what the $(n,n)$ matrix contains and why its rows sum to 1. If your colleague's follow-up is "so what does $V$ do that $K$ doesn't?", have an answer ready — you'll formalize it next lesson when we look at what a single head actually learns to do.