Module 12: The Feed-Forward Network

Why this matters

Modules 04-11 covered attention exhaustively — it's the mechanism that gets all the attention (so to speak) in most explanations. But roughly two-thirds of a transformer's parameters live in the much simpler feed-forward network (FFN) sitting right after attention in every block, verified with real numbers below. This module covers what the FFN actually computes, and verifies the one property that most sharply distinguishes it from attention: the FFN processes every token completely independently — permuting token order permutes its output identically, and changing one token's input has zero effect on any other token's output. Attention is where tokens exchange information; the FFN is where each token, alone, gets processed further.

Concepts

The structure: two linear layers with a nonlinearity between them

 FFN(x) = Linear2( activation( Linear1(x) ) )

 x (d_model)  ──►  Linear1 (d_model -> d_ff)  ──►  activation  ──►  Linear2 (d_ff -> d_model)  ──►  output (d_model)

d_ff (the "hidden" dimension inside the FFN) is conventionally 4x d_model — a real, common ratio, not an arbitrary choice, verified below to make the FFN the single largest parameter consumer in a transformer block.

Verified: the FFN holds roughly 2x as many parameters as attention

import torch
import torch.nn as nn

d_model, d_ff = 512, 2048   # d_ff = 4 * d_model, the standard ratio

ffn = nn.Sequential(nn.Linear(d_model, d_ff), nn.GELU(), nn.Linear(d_ff, d_model))
attn = nn.MultiheadAttention(d_model, num_heads=8, bias=False, batch_first=True)

print(sum(p.numel() for p in ffn.parameters()))
print(sum(p.numel() for p in attn.parameters()))

Verified output:

FFN parameters:        2,099,712
Attention parameters:  1,048,576
FFN is 2.0x attention's parameter count

This matches the algebra directly: attention's Wq, Wk, Wv, Wo are each d_model x d_model, for 4 * d_model^2 total (module 08 verified this doesn't change with head count). The FFN's two layers are d_model x d_ff and d_ff x d_model; at d_ff = 4*d_model, that's 2 * 4 * d_model^2 = 8 * d_model^2 — exactly double attention's count. Across an entire model, the FFN sub-layers are collectively the largest share of total parameters, not attention — a fact easy to miss given how much more explanatory attention typically receives.

Verified: the FFN processes every position completely independently

This is the property worth confirming directly, since it's the sharpest contrast with everything modules 04-11 covered. Permute the tokens in a batch and re-run the FFN:

X = torch.randn(1, 4, d_model)
out = ffn(X)

perm = [2, 0, 3, 1]
out_perm = ffn(X[:, perm, :])
print(torch.allclose(out_perm, out[:, perm, :], atol=1e-6))

Verified output: True — the same permutation-equivariance property module 03 proved for attention, but for a structurally different reason here: the FFN isn't just equivariant to permutation, it's applying the exact same function, completely independently, to each position — there's no mixing between positions at all, verified more sharply below.

Verified: changing one token's input has zero effect on any other token's FFN output

X2 = X.clone()
X2[:, 1, :] = torch.randn(d_model)   # change ONLY token 1's input
out2 = ffn(X2)
print(torch.allclose(out[:, 0, :], out2[:, 0, :]))

Verified output: True — token 0's FFN output is completely unaffected by an arbitrary change to token 1's input. This is a much stronger property than attention's permutation-equivariance (module 03): attention explicitly lets every token's output depend on every other token (that's the entire point of it); the FFN's output at one position depends on only that position's own input, full stop. In a transformer block, cross-token information exchange happens exclusively in the attention sub-layer — the FFN never lets information flow between positions at all.

Why put a per-token-only layer right after a cross-token layer?

Given the FFN can't exchange information across positions, its job is different from attention's: attention decides what information to gather from where; the FFN then processes each token's own, now attention-enriched representation further — a kind of per-token "thinking step" applied uniformly to every position, using the same learned weights everywhere (the same Linear1/Linear2 weights process every position, just with different input values). Stacking many [attention, then FFN] blocks (module 15) alternates between "gather information across the sequence" and "process what was gathered, independently, per token" repeatedly.

The activation function's role

Without a nonlinearity between the two linear layers, Linear2(Linear1(x)) would collapse into a single linear transformation (a composition of two linear maps is itself linear) — the FFN would add depth without adding any actual representational power beyond a single matrix multiplication. GELU (used in the verified example above) and its relatives (ReLU, SwiGLU in newer models) are what make the FFN capable of representing genuinely nonlinear per-token transformations.

Reference

 Component          Shape                        Parameter count
 ─────────────      ─────────────────────────    ────────────────────
 Linear1             d_model -> d_ff               d_model * d_ff
 Linear2             d_ff -> d_model                d_ff * d_model
 Total FFN                                          2 * d_model * d_ff
                                                     (= 8 * d_model^2
                                                     at the standard
                                                     d_ff = 4*d_model)
 Total attention                                    4 * d_model^2
                                                     (module 08)

 Property                        Attention                FFN
 ─────────────────────────       ────────────────────      ─────────────────────
 Mixes information across         YES (that's its job)      NO — verified: one
 positions                                                    token's output is
                                                                 unaffected by any
                                                                 other token's input
 Permutation equivariant           YES (module 03)            YES, but for a
                                                                stronger reason
                                                                (complete per-token
                                                                independence)

Hands-on exercises

Exercise 1 — reproduce the parameter-count comparison

Compute the FFN and attention parameter counts at your own choice of d_model and d_ff (try the standard d_ff = 4*d_model and also a smaller ratio) and confirm the FFN-to-attention ratio matches the algebra (2 * d_ff / d_model attention-equivalents, given the formulas above).

Exercise 2 — reproduce the per-token-independence verification

Run the exact "change one token's input" experiment above with your own random tensors, and confirm token 0's output truly doesn't change no matter what you set token 1 (or any other single token) to.

Exercise 3 — confirm the linear-collapse claim

Build an FFN with no activation function (nn.Sequential(Linear1, Linear2), nothing between them) and confirm it's mathematically equivalent to a single d_model x d_model linear layer — compute the combined weight matrix (Linear2.weight @ Linear1.weight, accounting for the actual shapes) and confirm applying it directly matches the two-layer version's output.

Independent challenge

A researcher proposes replacing a model's FFN with a lightweight attention-only architecture (no FFN sub-layer at all, just repeated attention blocks) to save the roughly 2x parameter cost verified above. Using this module's verified per-token-independence finding, write two or three sentences on what specific capability this removes, beyond the raw parameter savings.

Common mistakes & troubleshooting

• Assuming attention is where most of a transformer's parameters live. Verified above: the FFN holds roughly double attention's parameter count at the standard d_ff = 4*d_model ratio — it's the larger consumer of total model size, not attention.
• Assuming the FFN mixes information across tokens the way attention does. Verified directly: changing one token's input has zero effect on any other token's FFN output — cross-token mixing happens exclusively in attention.
• Omitting the activation function and not realizing the FFN collapses to a single linear layer. Verified in exercise 3 — a composition of two linear layers with nothing nonlinear between them is mathematically just one linear layer, adding parameters without adding representational depth.

Checkpoint quiz

1. What is the standard ratio between d_ff and d_model, and what did the verified parameter-count comparison show about FFN versus attention?
2. What did the "change one token's input" experiment verify about the FFN, and how does that differ from attention's behavior?
3. Is the FFN's permutation-equivariance (verified above) the same kind of property as attention's (module 03), or a stronger one?
4. Why does removing the activation function collapse a two-layer FFN into something equivalent to a single linear layer?
5. Given the FFN can't mix information across tokens, what is its actual job within a transformer block?

Further reading & sources

• Attention Is All You Need (Vaswani et al., 2017) - Section 3.3 defines the position-wise feed-forward network, including the standard d_ff = 4 * d_model ratio used throughout this module.
• Gaussian Error Linear Units (GELUs) (Hendrycks & Gimpel, 2016) - the GELU activation function used in this module's verification code and in most modern transformer FFNs.
• GLU Variants Improve Transformer (Shazeer, 2020) - covers SwiGLU and related FFN variants used in newer models (Llama and others), extending this module's basic two-linear-layer structure.
• PyTorch documentation: torch.nn.Linear - the layer type used throughout this module's verification code.
• Track 02, Module 08: Multi-Head Attention - the attention parameter-count formula this module's FFN comparison is measured against.

Next

Module 13: Residual Connections and Layer Normalization