Chapter 25 — Multiple heads, position, and the full block

Three more pieces and you have a complete Transformer.

Piece 1 — Multi-head attention

One attention operation forces every kind of relationship through a single softmax. Real language has many simultaneous relationships: grammatical subject, what a pronoun refers to, topic, tone.

Solution: run several attentions in parallel, each with its own W_Q, W_K, W_V, then concatenate the results.

                        input word vectors
                                |
        +-------------+---------+---------+-------------+
        |             |                   |             |
     head 1        head 2              head 3   ...  head 12
   (maybe learns  (maybe learns      (maybe learns
    "next word")   "pronoun ->        "subject of
                    noun")             the verb")
        |             |                   |             |
        +-------------+---------+---------+-------------+
                                |
                     concatenate, then one more
                     matrix (W_O) to mix them
                                |
                             output

• Each head works in a smaller space (d_model / n_heads), so total compute is about the same as one big head.
• Typical: d_model = 4096 with 32 heads of 128 dims each (Llama-2-7B).
• Interpretability researchers have found heads that reliably do specific jobs: previous-token heads, induction heads (which spot "A B ... A" and predict "B" — believed central to in-context learning), and syntactic heads.

Modern variant you'll see: GQA (Grouped-Query Attention). Many query heads share a smaller number of key/value heads. It often preserves quality well while shrinking memory during generation by 4–8x. Many current production LLMs use it. Chapter 38 explains why that memory matters so much.

Piece 2 — Position information

Here's a subtle problem: attention has no built-in idea what order the words are in. Look back at the formula — it's a weighted sum, and sums don't care about order. Without position information, self-attention is permutation-equivariant: if you shuffle the input tokens, the outputs shuffle in the same way. The model has no separate signal that says which word came first, second, or third.

But "dog bites man" ≠ "man bites dog". So we must inject position.

The rough intuition for sinusoidal/RoPE: think of a clock with many hands moving at different speeds. The fast hand tells you fine position; slow hands tell you roughly where you are in the document. The combination of all hands uniquely identifies a position and — importantly — lets the model compute relative distance easily.

RoPE won because it encodes relative position naturally and can be stretched after training to handle longer contexts than the model was trained on (Chapter 45).

Piece 3 — The full Transformer block

    x  (the "residual stream" -- a running representation of every token)
    |
    +----------------------------+
    |                            |
    |                       normalise
    |                            |
    |                    MULTI-HEAD ATTENTION      <-- tokens TALK TO EACH OTHER
    |                            |
    +---------> ( + ) <----------+                 <-- residual: ADD, don't replace
                  |
                  x
                  |
    +-------------+--------------+
    |                            |
    |                       normalise
    |                            |
    |                     FEED-FORWARD NET         <-- each token THINKS ALONE
    |                (expand 4x, non-linearity,        (a 2-layer MLP)
    |                 shrink back)
    |                            |
    +---------> ( + ) <----------+
                  |
                  x   (passed to the next identical block)

Three things to understand:

1. Attention vs feed-forward — the division of labour.

• Attention moves information between positions. ("What in this sentence is relevant to me?")
• The feed-forward network processes information within each position. It's a plain 2-layer MLP applied to each token separately. It holds roughly two thirds of the model's parameters, and interpretability work suggests it's where much factual knowledge is stored.

2. Residual connections (x + something). Each sub-layer adds to x rather than replacing it. This is the gradient highway from Chapter 15 — without it, 32-layer Transformers would not train. Think of x as a shared notepad that every layer reads and adds annotations to.

3. Normalisation. Rescales activations to a sane range at each step, keeping training stable. Modern LLMs use RMSNorm (a cheaper variant) placed before each sub-layer.

The whole model

    token IDs:  [ 464, 3797, 3332 ]           ("The cat sat")
        |
        v
    EMBEDDING LOOKUP: each ID -> a vector    [3 x d_model]
        |
        v
    +---------------------+
    | Transformer block 1 |
    +---------------------+
        |
    +---------------------+
    | Transformer block 2 |
    +---------------------+
        |
       ...   (32, 80, or more identical blocks)
        |
    +---------------------+
    | Transformer block N |
    +---------------------+
        |
        v
    final normalise
        |
        v
    OUTPUT LAYER: multiply by a [d_model x vocab_size] matrix
        |
        v
    LOGITS: one score per vocabulary word     [3 x 50000]
        |
        v
    SOFTMAX -> probability of each word being next

That is a GPT. Everything else — RoPE, GQA, SwiGLU, FlashAttention — is refinement of this skeleton.

The catch: attention is quadratic

Look at the scores matrix: it's n x n where n is the number of tokens.

    100 tokens ->     10,000 scores    fine
  1,000 tokens ->  1,000,000 scores    fine
100,000 tokens -> 10,000,000,000 scores  per head, per layer!  problem

Doubling the context length quadruples the attention work. This single fact drives most of Part 7 and Chapter 45 (why context windows are limited and what people do about it).