Chapter 24 — Self-attention, step by step with real numbers

Step 1 — Score: how much does each word match each other word?

Take the dot product of a query with every key. Big dot product = pointing in the same direction = relevant.

Let's compute for "sat" (query q3 = [1,1]):

q3 . k1 = (1 x 1) + (1 x 0) = 1        how relevant is "the" to "sat"?
q3 . k2 = (1 x 0) + (1 x 1) = 1        how relevant is "cat" to "sat"?
q3 . k3 = (1 x 1) + (1 x 1) = 2        how relevant is "sat" to itself?

raw scores: [1, 1, 2]

Step 2 — Scale: divide by sqrt(d_k)

sqrt(d_k) = sqrt(2) = 1.414

scaled scores: [1/1.414, 1/1.414, 2/1.414] = [0.707, 0.707, 1.414]

Why divide? With larger d_k, dot products get large just from having more terms to add up. Large numbers going into softmax make it "spiky" — almost all weight on one word, and near-zero gradients (learning stalls). Dividing by sqrt(d_k) keeps the scores in a well-behaved range. It's a numerical-stability fix, nothing deeper.

Step 3 — Softmax: turn scores into percentages

exp(0.707) = 2.028
exp(0.707) = 2.028
exp(1.414) = 4.113
             ------
sum        = 8.169

weights = [2.028/8.169,  2.028/8.169,  4.113/8.169]
        = [0.248,        0.248,        0.503]        (sums to 1.0)

Read this: "sat" pays 25% attention to "the", 25% to "cat", and 50% to itself.

Step 4 — Weighted sum of the values

output_for_sat = 0.248 x v1 + 0.248 x v2 + 0.503 x v3

              = 0.248 x [1.0, 0.0]
              + 0.248 x [0.0, 1.0]
              + 0.503 x [0.5, 0.5]

              = [0.248, 0.000]
              + [0.000, 0.248]
              + [0.252, 0.252]
              ------------------
              = [0.500, 0.500]

That's it. The new representation of "sat" is [0.5, 0.5] — a blend of the whole sentence, weighted by relevance. This vector is what gets passed to the next part of the network.

Step 5 — The causal mask (what makes it a language model)

For an LLM predicting the next word, word 2 must not see word 3 — that would be looking at the answer. So before the softmax, we set future scores to negative infinity, which softmax turns into exactly zero:

Attention allowed (X = allowed, . = blocked):

              key: the   cat   sat
   query the       X      .     .        "the" sees only itself
   query cat       X      X     .        "cat" sees "the" and itself
   query sat       X      X     X        "sat" sees everything before it

Let's redo "cat" (query q2 = [0,1]) with the mask:

q2 . k1 = 0        q2 . k2 = 1        q2 . k3 = BLOCKED (-infinity)

scaled: [0/1.414, 1/1.414, -inf] = [0, 0.707, -inf]

exp(0)     = 1.000
exp(0.707) = 2.028
exp(-inf)  = 0.000
             -----
sum        = 3.028

weights = [0.330, 0.670, 0.000]

output_for_cat = 0.330 x [1,0] + 0.670 x [0,1] + 0 x [0.5,0.5]
               = [0.330, 0.670]

"cat" attends 33% to "the" and 67% to itself, and 0% to the future.

This single choice — masked or not — is most of the difference between GPT (masked, generates text) and BERT (unmasked, understands text).

Practice (30 min) — implement attention yourself

import numpy as np

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

def attention(Q, K, V, causal=True):
    d_k = Q.shape[-1]
    scores = Q @ K.T / np.sqrt(d_k)              # step 1 + 2
    if causal:                                    # step 5
        n = scores.shape[0]
        mask = np.triu(np.ones((n, n)), k=1).astype(bool)
        scores[mask] = -np.inf
    weights = softmax(scores)                     # step 3
    return weights @ V, weights                   # step 4

Q = np.array([[1,0],[0,1],[1,1]], dtype=float)
K = np.array([[1,0],[0,1],[1,1]], dtype=float)
V = np.array([[1,0],[0,1],[0.5,0.5]], dtype=float)

out, w = attention(Q, K, V, causal=True)
words = ["the", "cat", "sat"]

print("attention weights (row = who is looking, col = who they look at):")
print("         " + "  ".join(f"{x:>6}" for x in words))
for i, row in enumerate(w):
    print(f"  {words[i]:>5}  " + "  ".join(f"{x:6.3f}" for x in row))

print("\noutputs:")
for i, row in enumerate(out):
    print(f"  {words[i]:>5}  {row.round(3)}")

Verify the numbers match the hand calculation above. They should: sat gets [0.5, 0.5], cat gets [0.33, 0.67].

Then experiment:

1. Set causal=False. How do the weights change? What does "the" now see?
2. Change q3 to [5, 0] (a much stronger query toward k1). Recompute — how peaked does the attention become?
3. Remove the / np.sqrt(d_k) and set Q = Q * 10. Look at the weights: they become nearly one-hot (0.999, 0.000, 0.000). Now you've seen why the scaling matters.
4. Make V random and 4-dimensional. Confirm the shapes still work.

Write in notes.md: "Attention computes ___, which lets each word ___."