Build an LLMMiddleLesson 405 min read

Step 4 · The transformer block

Attention plus a small feed-forward network plus two normalisations plus two residual connections. Repeat N times. That is the entire model.

🧪 Lab · 20 min📝 Lesson test
Lesson in motion

In 60 seconds

Step 4 · The transformer block

Attention plus a small feed-forward network plus two normalisations plus two residual connections. Repeat N times. That is the entire model.

1/5
In simple words
One block is: "look around at everyone else" then "think about what you found". Stack forty of those and you have a language model.
ONE BLOCK · repeated N timesRMSNormSelf-attentionmix across tokensRMSNormFeed-forwardthink per tokenresidualresidualtokens talktokens think
Two operations, alternating. Attention moves information between positions. The feed-forward network processes each position independently. Neither can do the other's job, which is why both are there.

The feed-forward network

Per token, expand to roughly four times the width, apply a non-linearity, come back down. It holds the majority of the model's parameters and is where most factual knowledge appears to be stored.
Modern feed-forward: SwiGLUpython
class SwiGLU(nn.Module):
    def __init__(self, d_model, hidden=None):
        super().__init__()
        # 8/3 rather than 4x, because the gate adds a third matrix
        hidden = hidden or int(8 * d_model / 3 / 64) * 64
        self.up   = nn.Linear(d_model, hidden, bias=False)
        self.gate = nn.Linear(d_model, hidden, bias=False)
        self.down = nn.Linear(hidden, d_model, bias=False)

    def forward(self, x):
        return self.down(F.silu(self.gate(x)) * self.up(x))
The gate is the modern part. One branch decides how much of the other branch to let through, per dimension. It consistently outperforms a plain two-matrix network at the same parameter count, which is why Llama, Mistral and most current models use it.

Residuals and normalisation: the two things that make depth possible

  1. 1

    Residual connections

    x = x + block(x). Each block writes a correction onto a running stream rather than replacing it. Gradients flow straight down the addition path, which is what allows 80-layer networks to train at all.
  2. 2

    Pre-normalisation

    Normalise before each sub-layer, not after. Original transformers did it after and needed careful learning-rate warmup to avoid diverging. Pre-norm is dramatically more stable and is now universal.
  3. 3

    RMSNorm over LayerNorm

    Skip the mean-subtraction and the bias; just divide by the root-mean-square and scale. Slightly cheaper, works just as well, so modern models use it.
  4. 4

    No biases

    Most current models drop bias terms in linear layers entirely. Fewer parameters, no measured loss, slightly better stability.
The complete blockpython
class RMSNorm(nn.Module):
    def __init__(self, d, eps=1e-6):
        super().__init__()
        self.eps = eps
        self.weight = nn.Parameter(torch.ones(d))
    def forward(self, x):
        rms = x.pow(2).mean(-1, keepdim=True).add(self.eps).rsqrt()
        return x * rms * self.weight

class Block(nn.Module):
    def __init__(self, d_model, n_head):
        super().__init__()
        self.n1 = RMSNorm(d_model)
        self.attn = CausalSelfAttention(d_model, n_head)
        self.n2 = RMSNorm(d_model)
        self.ff = SwiGLU(d_model)
    def forward(self, x):
        x = x + self.attn(self.n1(x))    # mix across positions
        x = x + self.ff(self.n2(x))      # think within each position
        return x
Do this
Look at how short that is. A frontier model and this block differ in scale, data and training, not in idea. The architecture has been essentially stable since 2017, with a handful of refinements.

Watch and read more

But what is a GPT? Transformers explained3Blue1Brown · video
Let's build GPT: from scratchAndrej Karpathy · video

Lab

An ablation study of the modern transformer block.

~20 min

The problem

Train the same tiny model four ways: (a) full modern block, (b) LayerNorm instead of RMSNorm, (c) plain MLP instead of SwiGLU, (d) post-norm instead of pre-norm. Report the loss and the training stability of each.

You are done when

Hard questions

Try to answer before you reveal. If you can answer these, you understood the lesson.

Q1Post-norm diverges without careful warmup; pre-norm does not. Explain the mechanism.Reveal
With post-norm the residual stream is normalised after every block, so the identity path is repeatedly rescaled and gradients must traverse every normalisation on the way back — deep stacks see exploding or vanishing updates early, when weights are random. Pre-norm leaves the residual path clean: gradients flow straight down the additions, and each block contributes a normalised correction. That is why pre-norm made very deep transformers trainable without heroics.

Please sign in to continue.

Questions people ask

Where is knowledge stored?

Evidence points mostly at the feed-forward layers, which behave somewhat like key-value memories. Attention decides what to retrieve and route; the feed-forward network holds much of what gets retrieved. It is not a clean separation.

Why 4x expansion in the FFN?

Empirical. It has been near-optimal across many scales. Gated variants use about 8/3 because they have three matrices instead of two, keeping the parameter count comparable.

How deep should a model be?

There is a rough relationship between width and depth that works well: a 7B model is typically around 32 layers with d_model 4096. Very deep and thin, or very wide and shallow, both underperform at a fixed parameter budget.

Is anything actually new since 2017?

At the block level: RoPE, RMSNorm, gated FFNs, grouped-query attention, and mixture-of-experts routing. Real improvements, all incremental. The transformative changes have been in scale, data quality and post-training.

Lesson test

5 questions. Get 3 right (60%) to pass and complete this lesson.

Sign in with your phone number to take the test and save your progress