The complete
Transformer block

1 What attention is missing

Attention solves the transport of information: it moves content from one position to another. But notice what it is, mathematically: a weighted average of V vectors. That is, a linear combination.

And there lies the problem. Stacking linear operations gains nothing: the composition of two linear transformations is another linear transformation. A hundred layers of pure attention would have exactly the same expressiveness as a single layer. You would gain nothing by stacking.

We need non-linearity. And we need a place where the model can process the information attention brought in, instead of merely transporting it. That place is the MLP.

2 Piece 1 — the MLP (where the knowledge lives)

The MLP is surprisingly simple: two linear layers with a non-linearity in between.

The shape is characteristic: expand, process, contract. Why expand to 3072 before coming back? Because the wide space gives the model room to separate patterns that are tangled in 768 dimensions — it is the same principle as a kernel in an SVM. The wide layer is where concepts become linearly separable.

The number that surprises

The MLP has twice as many parameters as attention in each block:

The non-linearity: GELU, not ReLU

GPT-2 uses GELU (Gaussian Error Linear Unit). The difference from the classic ReLU matters:

Read GELU as a probabilistic gate: instead of cutting sharply at zero, it multiplies the input by the probability that it "deserves to pass". Very negative values are nearly zeroed out, very positive values pass almost untouched, and in between there is a smooth transition.

Two concrete advantages over ReLU:

• It is differentiable everywhere. ReLU has a kink at x=0. Gradients like smooth curves.
• It does not kill neurons. In ReLU, a neuron that lands in the negative region receives a gradient of exactly 0 and may never recover (the "dying ReLU" problem). GELU always lets a sliver of gradient through.

In practice, GPT-2 uses a hyperbolic-tangent approximation of GELU — cheaper to compute and numerically identical for the model's purposes.

3 Piece 2 — the residual connection (the most important one)

If you can keep only one thing from this chapter, keep this one. The residual connection is what makes deep networks trainable, and it is ridiculously simple:

Why this saves training

In a deep network without residuals, the gradient has to cross every layer by successive multiplication. If each layer multiplies the gradient by, say, 0.9, then after 48 layers what remains is 0.9⁴⁸ ≈ 0.006 — the gradient has vanished. The first layers never learn. This is the classic vanishing gradient problem, which stalled deep learning for years.

With the residual, the derivative of the sum is:

The conceptual reading: the residual stream

There is a better way to see this, and it changed how people think about Transformers. Instead of "layers that transform", think of a shared blackboard of 768 dimensions that runs through the entire model from start to finish.

• The embedding writes the initial information on the board.
• Each block reads the board, thinks, and adds its contribution.
• No block erases what was there — the sum only adds.
• At the end, the board is read to produce the prediction.

This board is called the residual stream. It is the central concept of modern Transformer interpretability, and it explains why you can "read" what the model is thinking in the intermediate layers — the information is literally accumulated there, available.

4 Piece 3 — LayerNorm (and the change from GPT-1 to GPT-2)

As the vectors pass through dozens of blocks adding contributions, their magnitudes tend to grow out of control. LayerNorm renormalizes:

The change that defines GPT-2: Pre-LN

This is the most important architectural change between GPT-1 and GPT-2, and it is only a matter of where to put the LayerNorm.

This change alone is what allowed GPT-2 XL to have 48 layers. With Post-LN, training 48 layers would be a battle of numerical stability. After it, normalizing before attention became the standard choice for large Transformers, often in the form of RMSNorm.

GPT-2 also adds a final LayerNorm (called ln_f) after the last block, before projecting to the logits — needed precisely because, with Pre-LN, the output of the last block has never been through a normalization.

5 The assembled block

class Bloco(nn.Module):
    def __init__(self, d=768, h=12):
        super().__init__()
        self.ln_1 = nn.LayerNorm(d)
        self.attn = AtencaoMultiCabeca(d, h, causal=True)
        self.ln_2 = nn.LayerNorm(d)
        self.mlp  = nn.Sequential(
            nn.Linear(d, 4 * d),   # c_fc  : expands 768 -> 3072
            nn.GELU(),
            nn.Linear(4 * d, d),   # c_proj: contracts 3072 -> 768
        )

    def forward(self, x):
        # communication: the tokens talk
        x = x + self.attn(self.ln_1(x))
        # computation: each token thinks alone
        x = x + self.mlp(self.ln_2(x))
        return x

Two lines in the forward. Notice x = x + … in both: the original x is never destroyed, only added to. It is the blackboard being filled in.

The flow inside a block (5 tokens, GPT-2 small):

x            (5, 768) ---------------------------+  keeps the copy
  ln_1       (5, 768)                            |
  attn       (5, 768)                            |
                                                 v
x = x + attn(ln_1(x))                        (5, 768)  <- sum
  |
  +----------------------------------------------+
  ln_2       (5, 768)                            |
  c_fc       (5, 3072)  <- expands 4x            |
  GELU       (5, 3072)  <- non-linearity         |
  c_proj     (5, 768)   <- contracts             |
                                                 v
x = x + mlp(ln_2(x))                         (5, 768)  <- sum

INPUT SHAPE = OUTPUT SHAPE. That is why you can stack
12, 24, 36 or 48 of these without changing anything else.

6 In the lab

.venv\Scripts\python.exe lab\05_block.py

1. Builds a block from scratch and confirms that the shape goes in and comes out the same.
2. Compares Pre-LN and Post-LN in a 40-layer stack, measuring the norm of the gradient that reaches the first layer. Post-LN vanishes; Pre-LN survives. It is the demonstration of why the change was made.
3. Measures the growth of the residual stream layer by layer in the real GPT-2. You will see the norm grow almost monotonically — the blackboard being filled in.
4. Reads the blackboard halfway through — the technique called logit lens. See the box below: this is the most surprising experiment of the chapter.
5. Counts the parameters of a real block, piece by piece, confirming that the MLP takes two thirds.
6. Draws GELU against ReLU in ASCII, so you can see the difference in the critical region near zero.

You now know all the pieces. Chapter 06 assembles them into the complete GPT-2
and shows what really changes when you go from 124 million to 1.5 billion parameters.