Feed-Forward Network

Attention (lessons 1.0–1.1) is the only place where tokens exchange information. The second half of every block is its opposite: the position-wise feed-forward network (FFN, also called the MLP), which processes each token on its own. Attention gathers; the FFN digests.

FFN(x) = W2 @ relu(W1 @ x + b1) + b2

Two linear layers, one nonlinearity, hourglass shape:

1. Expand: W1 maps d_model → d_ff, usually d_ff = 4 * d_model (Vaswani: 512 → 2048; GPT-2 small: 768 → 3072).
2. Gate: ReLU originally; GELU or SwiGLU in modern models.
3. Project back: W2 maps d_ff → d_model — the output must match the input's shape because the block adds it residually as x + FFN(x) (lesson 1.4).

"Position-wise" means the same W1/W2 are applied independently to every token — a 1×1 convolution along the sequence. No token mixing happens here (that's attention's job), so all T tokens go through in one parallel batched matmul.

Where the parameters live

GPT-2 small (d_model = 768, d_ff = 3072), per block:

• W1 + b1: 768×3072 + 3072 ≈ 2.36M
• W2 + b2: 3072×768 + 768 ≈ 2.36M
• FFN total ≈ 4.7M vs ≈ 2.36M for all four attention projections (Q, K, V, O at 768×768 each).

That's attention ≈ 4d² vs FFN ≈ 8d² — about two-thirds of a transformer's weights sit in its FFNs. When people say "the knowledge lives in the MLPs," this ratio is why.

What it computes: match a pattern, write back a value

A productive reading (Geva et al. 2021): W1's rows are keys, W2's columns are values. Row j of W1 is a pattern detector — z_j = row_j · x + b_j is large when the token matches pattern j — ReLU decides which fire, and each firing detector adds its W2 column to the output, scaled by how hard it fired.

Tiny trace: d_model = d_ff = 2, x = [1, -1], biases 0, W2 = identity:

W1 row [1, 1]:   z1 = 1·1 + 1·(-1) = 0   → relu → 0   (pattern "x1 ≈ x2": silent)
W1 row [1, -1]:  z2 = 1·1 - 1·(-1) = 2   → relu → 2   (pattern "x1 opposes x2": fires)
output = [0, 2]

The output is built only from patterns the token actually matched — at d_ff = 3072, thousands of micro-rules per token. The ReLU is the entire point: attention is linear in its values, so without the FFN's nonlinearity a stack of blocks would collapse toward one big linear map.

Pure-Python idiom — and the trap

Idiom: row-major — output[i] = row_i · input + b_i. In the W @ x convention W1 is a d_ff × d_model matrix; papers flip between W @ x and x @ W, so check orientation before you loop — a silent transpose only crashes when d_ff ≠ d_model, so square toy examples hide the bug.

Second trap: an activation after W2. The output feeds x + FFN(x); clamp it non-negative and the residual can never push a coordinate down. Nonlinearity goes between the linears — the output stays signed.

Variants

• GELU (BERT, GPT-2): a smooth ReLU, x · Φ(x).
• SwiGLU (Llama): gated — value path × swish(gate path), with d_ff ≈ 8/3 · d_model to keep params level.
• Mixture of Experts (Mixtral 8×7B): 8 FFNs per block, each token routed to its top 2 — ~8× FFN capacity at ~2 experts of compute.

Cost: the FFN is O(T · d_model · d_ff) — linear in sequence length — vs attention's O(T² · d_model): long contexts bottleneck on attention even though the FFN owns the parameters.

Next: causal masking (1.3), then residual connections (1.4) wrap both sublayers, and lesson 2.0 assembles the full block: x = x + Attn(LN(x)); x = x + FFN(LN(x)).

Your exercise

Implement one FFN forward pass by hand: y = W2 @ relu(W1 @ x + b1) + b2, with W1 and W2 arriving as flat row-major lists you index yourself. Two mistakes its grader catches:

• Orientation. W1 is H x D and W2 is D x H, both row-major, so row i of W1 starts at i*D and row i of W2 starts at i*H. Use one stride for both and the square toy cases still pass — it breaks the moment H != D.
• An activation after W2. ReLU belongs between the two linears only. Clamp the output too and the residual x + FFN(x) can never push a coordinate down.

Two lessons ahead, residual connections wrap this sublayer; that exercise substitutes the identity for FFN so you can test the residual plumbing alone.