Skip to content
Feed-Forward Network
step 1/5

Reading — step 1 of 5

Read

~4 min readAttention

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

python

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

Nothing to submit here — the quiz below is the checkpoint. The FFN's slot gets graded two lessons ahead: the Residual Connections exercise implements y = x + alpha * sub(LayerNorm(x)), sub being the identity stand-in for this FFN. Two mistakes its grader catches:

  • Sample variance in LayerNorm. Divide by D, not D−1. For input X 1,2 with ALPHA 1 the expected stdout is exactly -1.0000,4.0000 (mean 1.5, population variance 0.25 → LN gives [-1, 1] → y1 = [0, 3] → y2 = [-1, 4]). With the D−1 variance you print -0.4142,3.4142 and fail every non-trivial test.
  • Formatting. Four decimals, comma-joined: ",".join(f"{v:.4f}" for v in y2). print(round(v, 4)) emits -1.0,4.0 — wrong string, zero marks, correct math.

Discussion

Ask a question, share an insight, or help someone who’s stuck.

Sign in to post a comment or reply.

Loading…