The transformer block

Inside the MLP

The MLP does three things to each token, one token at a time:

1. Expand. A linear layer turns the token's $d$ numbers into $4d$ numbers. With a width of 64, that's 256 numbers.
2. Bend. GELU, the smooth activation from Phase 4, bends those numbers. Without a bend, the two layers could only learn straight-line patterns.
3. Project back. A second linear layer turns the $4d$ numbers back into $d$.

The result has to be $d$ numbers long, because it gets added onto the token's vector. You can't add a list of 256 numbers to a list of 64.

A complete block

The code below builds a full block in PyTorch. Here's a map of it before you read it.

• self.ln1 and self.ln2 are the two layer norms.
• self.qkv makes the queries, keys and values for all heads in one go. self.proj is the output projection, $W_O$. The attention method is the multi-head attention from earlier in this phase, causal mask included.
• self.mlp is the expand, GELU, project-back network.
• forward is the two lines from the top of this lesson.

The code then runs 2 sequences of 10 tokens each through a block of width 64 with 4 heads. Finally, it counts the parameters in each part.

import math
import torch
import torch.nn as nn

class Block(nn.Module):
    def __init__(self, d, n_heads):
        super().__init__()
        self.h = n_heads
        self.ln1, self.ln2 = nn.LayerNorm(d), nn.LayerNorm(d)
        self.qkv = nn.Linear(d, 3 * d)                  # queries, keys, values for all heads
        self.proj = nn.Linear(d, d)                     # attention output projection (W_O)
        self.mlp = nn.Sequential(nn.Linear(d, 4 * d), nn.GELU(), nn.Linear(4 * d, d))

    def attention(self, x):
        B, T, d = x.shape
        q, k, v = self.qkv(x).split(d, dim=-1)
        q, k, v = (t.view(B, T, self.h, d // self.h).transpose(1, 2) for t in (q, k, v))
        s = q @ k.transpose(-2, -1) / math.sqrt(d // self.h)
        s = s.masked_fill(~torch.tril(torch.ones(T, T)).bool(), float("-inf"))
        y = (torch.softmax(s, dim=-1) @ v).transpose(1, 2).reshape(B, T, d)
        return self.proj(y)

    def forward(self, x):
        x = x + self.attention(self.ln1(x))             # communicate
        x = x + self.mlp(self.ln2(x))                   # compute
        return x

torch.manual_seed(0)
d, heads = 64, 4
block = Block(d, heads)
x = torch.randn(2, 10, d)
print("in:", tuple(x.shape), " out:", tuple(block(x).shape))

counts = {name: sum(p.numel() for p in mod.parameters())
          for name, mod in [("attention", nn.ModuleList([block.qkv, block.proj])), ("MLP", block.mlp),
                            ("layer norms", nn.ModuleList([block.ln1, block.ln2]))]}
total = sum(counts.values())
for name, n in counts.items():
    print(f"{name:<12} {n:>6,} parameters ({n / total:.0%})")
print(f"total        {total:>6,}   vs 12·d² = {12 * d * d:,}")

Look at the first line it prints. The output shape, (2, 10, 64), matches the input. That's what lets you stack blocks: each block's output is a valid input for the next one.

Now the counts. The MLP holds about two-thirds of the parameters (66%). Attention holds about a third (33%). The layer norms hold only 256 numbers, well under 1%. Python rounds that share up, so it prints 1%.

The block also stays causal: a token's output never depends on later tokens. The mask in attention makes sure of that. The layer norms and the MLP can't break it, because they never look at other tokens.

Here torch is the course's in-browser teaching version of PyTorch. It covers the parts of PyTorch this course uses.

Why about 12·d²

The code printed a total close to "12·d²". Here's where that rule of thumb comes from.

A grid of weights (a matrix) that turns 3 numbers into 5 numbers holds 3 × 5 = 15 weights. So a $d \times d$ matrix holds $d^2$ weights, $d$ times $d$.

• Attention has four $d \times d$ matrices: queries, keys, values and the output projection. That's $4d^2$.
• The MLP has a $d \times 4d$ matrix to expand and a $4d \times d$ matrix to come back. Each holds $4d^2$, so together they hold $8d^2$.

In plain words: square the width, then multiply by 12. A third of that is attention, and two-thirds is the MLP. The $\approx$ sign means "about".

Try it with a tiny width of 4. Then $d^2 = 16$.

• Attention: 4 × 16 = 64.
• MLP: 8 × 16 = 128.
• Total: 64 + 128 = 192, which is 12 × 16.

The same count in Python, plus the estimate for a real model:

d = 4
attention = 4 * d * d                 # queries, keys, values and output: four d x d grids
mlp = d * (4 * d) + (4 * d) * d       # expand d to 4d, then come back from 4d to d
print(attention, mlp, attention + mlp, 12 * d * d)

blocks, width = 12, 768               # GPT-2 small
print(f"{blocks * 12 * width * width:,}")

The first line prints 64, 128, 192 and 192. The second line is GPT-2 small: width 768, 12 blocks. It prints 84,934,656, about 85 million parameters in its blocks. The rest of GPT-2's parameters sit in its token and position embedding tables.

The block you built earlier has width 64, so the estimate is 12 × 64² = 49,152. The code counted 49,984. The extra 832 are biases and layer norms. They grow with $d$, not $d^2$, so in a big model they barely register.

d = 100
print(12 * d * d)

Variations you'll see in model cards

A model card is the fact sheet published alongside a model. Model cards often list small changes to the block:

• Post-norm vs pre-norm. The original 2017 transformer normalized after adding the residual. Almost every model since puts the norm first. That trains more stably when many blocks are stacked.
• RMSNorm instead of LayerNorm. It's a simpler normalization: it rescales the numbers but skips subtracting their average. Llama-style models use it.
• SwiGLU instead of GELU in the MLP, also in Llama-style models. It's a gated activation: one part of the MLP decides how much of another part gets through. The hidden size is adjusted to keep the parameter count similar.
• Grouped-query attention. Several query heads share the same keys and values. There are fewer key and value heads than query heads, which saves memory while the model generates text.
• Mixture of experts. The single MLP is replaced by several MLPs and a router. The router picks which MLPs each token uses. A later lesson covers it.

The skeleton is the same in all of them: attention, then the MLP, each wrapped in a norm and a residual connection.

x = layer_norm(x + attention(x))      # post-norm (2017): add first, then normalize
x = x + attention(layer_norm(x))      # pre-norm (today): normalize only the step's input

import numpy as np

x = np.array([[1.0, 0.0],      # token 0
              [0.0, 2.0],      # token 1
              [2.0, 1.0]])     # token 2
T = len(x)

# Communicate: each token averages itself and the tokens before it (causal, equal weights).
weights = np.tril(np.ones((T, T)))
weights = weights / weights.sum(axis=1, keepdims=True)
x = x + weights @ x
print("after attention:", x.tolist())

# Compute: the same small function runs on each token separately.
def mlp(token):
    return np.maximum(0, token - 1)
x = x + np.array([mlp(t) for t in x])
print("after MLP:      ", x.tolist())