12.7.4 — The Transformer Block

12.7.3 ended with three things single attention cannot do. It computes only one kind of relationship. It has no idea of word order, so "the courier delivered the parcel" and "the parcel delivered the courier" look identical to it. And it only averages information from other tokens without transforming it. A real model needs all three fixed, and the fixes, together with the wiring that makes deep stacks trainable, form the transformer block: the single unit that an LLM repeats 32, 80 or 120 times.

This page assembles that block piece by piece, then shows how three ways of connecting blocks give the three families of transformer, and why one of them, decoder-only, became the shape of nearly every LLM.

1. Multi-head attention: several kinds of relevance at once

Consider what the token "it" in "The kettle arrived but it was broken" needs from its context. It needs to know which noun it refers to ("kettle"). It needs to know the verb that follows ("was"). It may need to know that the sentence contains a contrast ("but"). One attention operation, with one set of query, key and value matrices, learns one notion of "relevant", and a weighted average cannot point strongly at three different places at once without blurring them together.

Multi-head attention runs several attention operations side by side, each with its own W_Q, W_K and W_V, called heads. Each head can learn a different kind of relationship. Their outputs are joined end to end and passed through one more learned matrix, W_O, which mixes the heads' findings back into a single vector per token.

Multi-head attention. Several heads attend in parallel, each learning its own kind of relationship; their outputs are joined and mixed by one output matrix. The roles shown are typical of what researchers find inside trained models; nobody assigns them.

Two facts make this cheap. First, the heads split the work rather than multiplying it: with a model width of 4,096 and 32 heads, each head works in only 4{,}096 / 32 = 128 dimensions, so the 32 small heads together cost about the same as one full-width head. Second, all heads run at the same moment, which suits the GPU.

The roles in the figure are the kind of thing researchers find when they inspect trained models: heads that attend to the previous token, heads that link pronouns to nouns, heads that match opening and closing brackets, and, in models that process code and lists, heads that copy patterns from earlier in the context. Nobody assigns these roles. The heads start with different random weights, and training pushes each towards whatever kind of relationship best reduces the loss. Many heads do things nobody has yet described in plain words.

Sharing keys and values: grouped-query attention

One variation matters for cost. In plain multi-head attention, every head has its own keys and values, and during generation all of them must be kept in the KV cache for every token so far, which 12.7.6 shows is often the biggest memory cost in serving. Grouped-query attention, or GQA, keeps many query heads but lets groups of them share one key head and one value head. With 32 query heads sharing 8 key-value heads, the cache shrinks by a factor of four, for a small loss in quality. Most current open models use it. The extreme version, one shared key-value head for all queries, is called multi-query attention.

2. Position: putting word order back

Attention compares contents, never places, so a transformer must be told where each token is. Three methods have been used, and the third is the one you will meet in nearly every current model.

Sinusoidal position vectors. The original 2017 transformer added to each token's embedding a fixed vector made of sine and cosine waves of different frequencies, a unique pattern for each position. It works, and it is a neat idea, but it is no longer common.

Learned position vectors. GPT-2 and BERT simply learned one vector per position, up to a maximum length, and added it to the token embedding. Simple, with one serious flaw: a model trained with positions up to 1,024 has no vector at all for position 1,025, so it cannot handle longer text.

Rotary position embedding. RoPE, introduced by Jianlin Su and colleagues in 2021, does not add anything to the embeddings. Instead, inside every attention head, it rotates each query and each key by an angle proportional to its position, before the dot product. Think of the numbers in a query as pairs, each pair a little arrow in a plane. A token at position 1 has its arrows turned by one step, a token at position 10 by ten steps.

Why rotation? Because of what it does to the dot product. From 12.2.1, the dot product depends on the angle between two arrows. If a query at position m is turned by m\theta and a key at position n by n\theta, the angle between them changes by (m - n)\theta, so the score depends only on how far apart the two tokens are, not on where they sit in the text. A worked check, with a step of \theta = 0.5 radians and query and key both starting as the arrow (1, 0):

Query position	Key position	Gap	Dot product after rotating
3	1	2	0.540
10	8	2	0.540
3	2	1	0.878

Positions 3 and 1, and positions 10 and 8, give exactly the same score, \cos(1.0) \approx 0.540, because the gap is 2 in both cases. A gap of 1 gives a different score. Relative distance is what matters for language ("the word just before me", "the noun three words back"), and RoPE builds exactly that into every head. Real models use many pairs per head, each turning at a different speed, so that both near and far distances can be told apart.

RoPE also made a practical trick possible: extending the context window. A model trained on sequences of 8,000 tokens can be adapted to much longer ones by slowing down the rotation speeds so that long distances look like the distances it already knows, followed by a short extra training run on long documents. Variants of this rescaling are how many models' context windows grew from thousands to hundreds of thousands of tokens. A simpler alternative you may see named is ALiBi, which skips rotations and directly subtracts a penalty from each attention score in proportion to the distance between the tokens.

3. The feed-forward layer: processing what was gathered

After attention has gathered information between tokens, each token's vector goes, on its own, through a small two-layer network, the multilayer perceptron of 12.5.1. The same network, with the same weights, is applied separately at every position:

\text{FFN}(\mathbf{x}) = W_2\, \text{GELU}(W_1 \mathbf{x} + \mathbf{b}_1) + \mathbf{b}_2

Read it as "a weight matrix, then the GELU bend from 12.5.1, then another weight matrix". The first matrix widens the vector, traditionally to 4 times the model width, so 4,096 numbers become about 16,000, and the second narrows it back. Many current models use a slightly richer version called SwiGLU, with three matrices instead of two, in which one widened copy of the vector acts as a set of learned on-off gates for the other.

The split of work is worth remembering: attention moves information between tokens; the feed-forward layer transforms the information within each token. And the feed-forward layers are big. Here is the count for one block of Llama 2 7B, whose model width is 4,096 and whose SwiGLU layer widens to 11,008:

• Attention: four 4,096 × 4,096 matrices (W_Q, W_K, W_V, W_O) ≈ 67 million weights.
• Feed-forward: three 4,096 × 11,008 matrices ≈ 135 million weights.

So about two-thirds of every block is feed-forward. Thirty-two such blocks give about 6.48 billion weights, and the input embedding table and output layer, each 32,000 tokens × 4,096, add about 0.26 billion more, for 6.74 billion in total: exactly the size of Llama 2 7B. There is good evidence that a lot of what a model "knows", facts such as "the capital of France is Paris", is stored in these feed-forward weights, with particular neurons acting somewhat like keys that recognise a pattern and values that add the associated information to the token's vector.

4. The wiring: residual stream and normalisation

The two sublayers are connected with the residual connections and layer normalisation of 12.5.4, in this order:

def transformer_block(x):
    x = x + multi_head_attention(layer_norm(x))   # (1) gather from other tokens, add it on
    x = x + feed_forward(layer_norm(x))           # (2) process each token, add it on
    return x

Line (1) normalises each token's vector, runs attention, and adds the result to the vector rather than replacing it. Line (2) does the same with the feed-forward layer. Normalising before each sublayer, as here, is called pre-norm. The original 2017 design normalised after the addition, post-norm, which needed delicate learning-rate warmup to train at all; pre-norm is far more stable and is what large models use. Many current models also replace layer normalisation with the slightly cheaper RMSNorm from 12.5.4.

The residual stream. Each token's vector travels along a single path from embedding to output. Every attention and feed-forward sublayer reads a normalised copy and adds its result back, so the vector is refined step by step rather than replaced.

This picture gives a useful way to think about the whole model, often called the residual stream. Each token has one vector that travels from the embedding at the start to the unembedding at the end. Nothing ever overwrites it. Every sublayer, in every block, reads a normalised copy, computes something, and adds its contribution back. Early blocks tend to add simple information (what kind of word this is, what is next to it); later blocks add more abstract information (what the sentence is about, what should come next). The gradient also flows back along this stream with nothing to shrink it, which is why the residual connections of 12.5.4 let these models be a hundred blocks deep.

5. Encoders, decoders, and why decoder-only won

The same block can be wired into three kinds of model, and the difference is entirely in what each token is allowed to attend to.

Encoder-only. No causal mask: every token attends to every other, before and after it. The model therefore sees the whole text at once in both directions, which is ideal for understanding tasks. BERT, from Google in 2018, is the famous example. It was trained by hiding about 15 percent of the words in a sentence and predicting them from both sides, which is called masked language modelling. Encoders cannot generate text left to right, but they are excellent at classification and at producing the sentence embeddings of 12.6.2: most embedding models are encoders.

Decoder-only. The causal mask of 12.7.3 on every attention layer: each token sees only itself and what came before. Trained to predict the next token. GPT, Claude, Llama, Gemini and nearly every current LLM are decoder-only.

Encoder-decoder. An encoder reads the input in both directions; a decoder generates the output left to right, and in each of its blocks has an extra attention layer, cross-attention, whose queries come from the decoder and whose keys and values come from the encoder's output. This is the original 2017 design, built for translation, and it lives on in models such as Google's T5 and in many speech-recognition models, where the input is audio and the output is text.

Why did decoder-only win for LLMs? Because "predict the next token" can express almost every task. Translation is "here is a sentence in English; its Hindi version is …". Summarising, answering, classifying and writing code are all continuations of some text. One architecture with one training objective, applied to all the text available, turned out to scale more simply and further than designs with separate parts for input and output. When the strategy is "make it bigger and train it on more", the simplest design that can do everything wins.

6. The whole model, in one picture

Put it together and a decoder-only LLM is:

1. A token embedding table (one learned vector per token).
2. A stack of L identical blocks, each: pre-norm, multi-head causal attention with RoPE, add to the stream; pre-norm, feed-forward, add to the stream.
3. A final normalisation, and an output matrix that turns the last position's vector into one logit per vocabulary entry.

Nothing else. The differences between models of the same family are mostly in the numbers: width (4,096 or 8,192), number of blocks (32, 80, 126), number of heads and key-value heads, vocabulary size, feed-forward width, and how many tokens of what text it was trained on. A minimal working block in PyTorch fits in a few lines:

import torch, torch.nn as nn

class Block(nn.Module):
    def __init__(self, d=512, heads=8):
        super().__init__()
        self.n1, self.n2 = nn.LayerNorm(d), nn.LayerNorm(d)
        self.attn = nn.MultiheadAttention(d, heads, batch_first=True)           # (1)
        self.ffn = nn.Sequential(nn.Linear(d, 4 * d), nn.GELU(), nn.Linear(4 * d, d))  # (2)

    def forward(self, x):
        n = x.shape[1]
        causal = torch.triu(torch.ones(n, n, dtype=torch.bool), diagonal=1)    # (3)
        h = self.n1(x)
        x = x + self.attn(h, h, h, attn_mask=causal, need_weights=False)[0]    # (4)
        x = x + self.ffn(self.n2(x))                                            # (5)
        return x

x = torch.randn(2, 10, 512)        # 2 sequences, 10 tokens, 512 numbers each
print(Block()(x).shape)            # → torch.Size([2, 10, 512])

Line (1) creates multi-head attention with 8 heads of 512 / 8 = 64 dimensions each. Line (2) is the feed-forward layer, widening to 4 times and back. Line (3) builds the causal mask: True above the diagonal means "not allowed to attend". Lines (4) and (5) are the two residual updates of section 4, with pre-norm. The output has the same shape as the input, which is what lets blocks be stacked any number of times. (This sketch leaves out RoPE and other refinements that production models add.)

Recall

• Multi-head attention runs several heads in parallel, each with its own projections and its own learned kind of relationship, then concatenates and mixes them with W_O; splitting the width keeps the cost about that of one head. Grouped-query attention lets query heads share key-value heads to shrink the KV cache.
• Attention ignores order, so position is added. RoPE rotates queries and keys by an angle proportional to position, so scores depend only on the gap between tokens (positions 3 and 1 score the same as 10 and 8); slowing the rotations is how context windows are extended.
• The feed-forward layer widens each token's vector, bends it, and narrows it back, separately per token; it holds about two-thirds of the weights (Llama 2 7B: 67 million attention + 135 million feed-forward per block, 6.74 billion total).
• Pre-norm residual wiring makes the residual stream: each token's vector is refined by every sublayer, never overwritten.
• Encoder-only (BERT) sees both directions and suits understanding and embeddings; encoder-decoder suits translation and speech; decoder-only, with the causal mask and next-token prediction, won for LLMs because one objective expresses almost every task.

Self-test: With a width of 8,192 and 64 heads, how many dimensions does each head use? Using the RoPE table, why must positions 10 and 9 give 0.878? Roughly how many weights does one block have with width 4,096 and a plain 4× feed-forward layer?

Next: 12.7.5 Pretraining and Scaling — the block is simple; what turns a stack of them into something that can answer questions is training on trillions of tokens, and the surprisingly regular laws that say how much better models get with more.