Module 3: The real thing

The full head, in seven lines

Module 2 built every piece: queries and keys ask and advertise, values carry the goods, and $\sqrt{d_k}$ keeps softmax in its responsive range. Packaged together, with all three projection matrices now taken as inputs (in a real model, training supplies them), the mechanism is this function from the script:

def attention(X, W_q, W_k, W_v, causal=False):
    """One full head of self-attention: the real thing, in seven lines.

    X has one row per word. W_q, W_k, W_v are the head's three learned
    matrices (here: supplied by the caller). Returns the new row-per-word
    matrix and the attention weights.
    """
    # Project each word into its three roles.
    Q = X @ W_q   # what each word is looking for
    K = X @ W_k   # what each word advertises
    V = X @ W_v   # what each word hands over if attended to

    # Every query meets every key: scores[i, j] = Q[i] . K[j].
    scores = Q @ K.T / np.sqrt(W_k.shape[1])

    # A causal mask hides the future: word i may only look at words j <= i.
    # Setting a score to -infinity makes its softmax weight exactly 0.
    if causal:
        future = np.triu(np.ones(scores.shape, dtype=bool), k=1)
        scores = np.where(future, -np.inf, scores)

    # Each row of scores becomes weights; each output row is the
    # weights-blended average of the value rows.
    weights = softmax(scores)
    return weights @ V, weights

Ignore the causal branch for a moment (it is this module's second topic). Everything else is Module 2's pipeline verbatim, now with a real W_v instead of the wired head's "just hand over your embedding" choice. This function is self-attention as deployed: no simplifications left.

Does it match PyTorch?

PyTorch ships this same computation as torch.nn.functional.scaled_dot_product_attention, the workhorse behind its transformer layers. If our understanding is right, our NumPy and their implementation must compute the same numbers, not roughly, but exactly up to floating-point rounding. The script makes random projection matrices, runs both, and compares; this is the code that writes the capture below:

    # A general head: random learned-looking projections, d_model 4 -> 3.
    # Real models get these matrices from training; random ones are enough
    # to check that the two implementations compute the same function.
    W_q = rng.standard_normal((4, 3))
    W_k = rng.standard_normal((4, 3))
    W_v = rng.standard_normal((4, 3))

    ours, _ = attention(X, W_q, W_k, W_v)

    # PyTorch's fused kernel takes the already-projected Q, K, V.
    tQ = torch.from_numpy(X @ W_q)
    tK = torch.from_numpy(X @ W_k)
    tV = torch.from_numpy(X @ W_v)
    theirs = F.scaled_dot_product_attention(tQ, tK, tV).numpy()

The verdict (cross-check: results.log lines 84-88):

[S6] the general head, checked against PyTorch
  ours vs torch.nn.functional.scaled_dot_product_attention:
    max |difference| = 2.22e-16
  same comparison with the causal mask on both sides:
    max |difference| = 2.22e-16

Let's go deeper on that number, because it is not just "small". Double-precision numbers carry about 16 significant digits, and 2.22e-16 (which is $2^{-52}$) is machine epsilon: for numbers of ordinary size, like these outputs, it is the gap between one representable double and the very next one. The two implementations disagree by at most one final bit of rounding, which is what you get when two codebases compute the same function while ordering their arithmetic slightly differently. Twelve lines of NumPy, written from three ideas you can verify on paper, and the industrial implementation agree completely. (Note one API detail: PyTorch's function takes Q, K, V already projected, so the projections happen on our side in both runs; the comparison tests the attention core itself.)

The causal mask: no reading ahead

So far every word could attend to every word, including words that come after it. For some uses that is correct. But an LLM in the GPT family is trained to predict the next token (a token is the unit of text the model reads and writes, usually a word or a piece of one). During training, every position in a sentence predicts its successor simultaneously: position 1 of "the river bank" predicts "river" while position 2 predicts "bank". If attention let position 1 read position 2, the answer would be sitting in plain sight, the model would learn to copy instead of predict, and at generation time (when future words genuinely do not exist yet) the crutch would vanish. So generator-style LLMs run attention with a causal mask: word $i$ may look only at words 1 through $i$. (Models built to read text rather than generate it, like search-oriented encoders, skip the mask; that is the other major flavor of transformer.)

The mechanism needs one property from Module 0: softmax hands $e^0=1$ even to a zero score, so silencing a word takes more than a zero. The trick is $-\infty$: $e^{-\infty }=0$, an exactly-zero share. The causal branch above builds future, a matrix that is True strictly above the diagonal (that upper-right triangle is exactly the "word j comes after word i" region), and overwrites those scores with -np.inf. The script applies this to Module 2's legible wired head:

    # The causal mask on the legible wired head from [S4].
    Xw = embed(SENT_RIVER)
    Q = Xw @ W_Q_WIRED
    K = Xw @ W_K_WIRED
    scores = Q @ K.T / np.sqrt(2)
    future = np.triu(np.ones(scores.shape, dtype=bool), k=1)
    W_masked = softmax(np.where(future, -np.inf, scores))

Predict. The mask only removes words from view; softmax then renormalizes whatever is left. Before looking at the table: which of the three rows must come out unchanged from Module 2's unmasked table, and what must the first word's row become?

The masked weights (cross-check: results.log lines 89-94):

  wired head with a causal mask (upper right forced to 0):
                the   river    bank
    the       1.000   0.000   0.000
    river     0.406   0.594   0.000
    bank      0.159   0.500   0.341
  every row still sums to 1: 1.000000, 1.000000, 1.000000

Let's go deeper, row by row. "the", the first word, has no past except itself: its entire row collapses to weight 1.000 on itself, and this happens in every causal attention matrix regardless of the sentence or the weights (quiz question 2 asks you to say why). "river" may see only "the" and itself; its Module 2 weights get renormalized over that pair to 0.406 and 0.594. And "bank", the last word, sees the whole sentence, so its row [0.159, 0.500, 0.341] is identical to its unmasked Module 2 row: masking never changes the last word's view. The capture's last line confirms the rows still sum to 1: the mask reallocates shares, it does not leak them.

What attention alone still cannot do

One honest limitation, measured before we zoom out. The script's final check feeds the reversed sentence through Module 1's attention and compares bank's output in both:

    # Limitation worth measuring: attention alone ignores word ORDER.
    # Reversing the sentence permutes the output rows but changes nothing
    # about what each word computes.
    X_rev = embed(list(reversed(SENT_RIVER)))
    out_fwd, _ = simple_attention(embed(SENT_RIVER))
    out_rev, _ = simple_attention(X_rev)
    same = bool(np.allclose(out_fwd[2], out_rev[0]))

The result (cross-check: results.log lines 95-96):

  order-blindness: bank's output in "the river bank" equals its
  output in "bank river the" -> True (real LLMs add position information)

Nothing in any formula we built ever consulted a word's position: scores depend only on vector contents, so scrambling the sentence just scrambles the output rows correspondingly. Attention treats a sentence as a bag of words. Real models fix this before attention ever runs, by mixing position information into each embedding (that is the "+ position info" box in the diagram below); how they encode it is a workshop of its own, but you now know exactly why it is needed.

Where this sits inside GPT-2

Follow the stack downward. The text so far goes through the tokenizer box (text to token IDs: the lookup keys for the embedding table). The embeddings + position info box is Module 0's table plus the position fix we just measured the need for. Then comes the transformer block, repeated 12 times in GPT-2, and each block is two workers: the multi-head self-attention box, which is this workshop, and the MLP box, a per-word neural network that processes each (now context-aware) vector on its own, with no cross-word communication; attention is the only place words exchange information. The final box turns the last layer's vectors into probabilities for the next token.

"Multi-head" is the one term left to unpack. Our wired head could only ask one kind of question (topic-seeking, with a place-seeking side channel). Real sentences need many kinds of lookup at once: grammar, coreference, topic, position-adjacent context. So each layer runs several attention heads in parallel, each with its own learned $W_q,W_k,W_v$ and its own weight table, and their outputs are concatenated. GPT-2 uses 12 heads per layer. The script closes by pricing all of this at GPT-2's true sizes, where $d_{model}=768$ is the embedding width (our toy used 4):

    # Scale check: the same seven lines at GPT-2's size.
    d_model, d_k, heads, layers = 768, 64, 12, 12
    per_head = 3 * d_model * d_k
    per_layer = heads * per_head + d_model * d_model

The price tag (cross-check: results.log lines 97-98):

  one GPT-2 sized head: d_model = 768, d_k = 64 -> W_q + W_k + W_v = 147,456 numbers
  12 heads x 12 layers (+ output projections) -> 28,311,552 attention parameters of GPT-2's ~124M

Let's go deeper. One head's three matrices hold 3 x 768 x 64 = 147,456 learned numbers: our wired head made its point with 16. Across 12 heads and 12 layers, plus each layer's output projection (the 768 x 768 matrix that merges the concatenated heads back to embedding width), attention accounts for about 28.3 million of GPT-2's 124 million parameters, between a fifth and a quarter of the model. Every one of those parameters sits inside the seven-line function at the top of this page; scale changes the sizes of the matrices, not the mechanism.

What this workshop deliberately left out, so you know the road ahead: how training discovers useful projection matrices (they are learned by the same gradient-based process that learns the embedding table); what the MLP half of each block computes; how position information is encoded; and the engineering that makes generation fast (chiefly caching each step's keys and values instead of recomputing them). None of those change the mechanism you now own.

Reproducing this workshop

Every number and figure on these pages comes from one deterministic run of experiment.py (seed 0, fixed in the script):

• OS: Ubuntu 22.04 (Linux 6.8).
• Python: 3.10.12, with numpy 2.2.6, torch 2.12.1+cpu, matplotlib 3.10.6 (the same versions are stamped at the top of the captured log).
• Run: copy the script from the Source page into an experiments/ directory, then pip install numpy torch matplotlib and python3 experiment.py. It runs in seconds on a CPU and writes captures/results.log and the diagrams/*.png figures. Your log should match the captured one byte for byte.
• The two flow diagrams (the attention pipeline and the LLM stack) are drawn with D2 v0.7.1 and rasterized with rsvg-convert 2.52.5; their .d2 sources live in the workshop's diagrams/ directory.

You have the whole mechanism. Time for the final exam: 18 questions, 70 to pass. If you want to keep going afterward, re-run the experiment with your own sentences and axes: the Source page is yours to edit.