Feed-Forward Networks: The Forgotten Half of Transformers
Published:
Intuition First: The FFN as a Pattern-Response Memory
Think of the FFN as a giant associative memory. The first matrix \(W_1\) acts as a bank of pattern detectors, each row is a template asking “does this token look like X?” The nonlinearity fires neurons that match. The second matrix \(W_2\) then says “when pattern X fires, add feature vector Y to the output.”
So for a token representing “Paris” in context “capital of France”, a neuron in the expanded layer might activate for the pattern “capital-of-Europe-city” and the corresponding \(W_2\) column adds a “France-related” feature vector to the output. This is factual retrieval, not via attention, but via the FFN’s stored patterns.
The FFN Is Half the Block
Every Transformer block follows this pattern:
x → MultiHeadAttention → residual + LN → FeedForward → residual + LN → output
The FeedForward (FFN) sub-layer is the second half of every block. In popular Transformer explanations, it is often described in one sentence and then forgotten in favour of attention. This is a mistake, the FFN is critical.
The Architecture of the FFN
The FFN is a simple two-layer MLP applied position-wise: each token is processed identically and independently.
- \(W_1 \in \mathbb{R}^{d_{\mathrm{model}} \times d_{\mathrm{ff}}}\): projects up from \(d_{\mathrm{model}}\) to \(d_{\mathrm{ff}}\)
- activation: nonlinearity (ReLU, GELU, or SwiGLU)
- \(W_2 \in \mathbb{R}^{d_{\mathrm{ff}} \times d_{\mathrm{model}}}\): projects back down
- \(d_{\mathrm{ff}} = 4 \times d_{\mathrm{model}}\) in most models (e.g., \(512 \to 2048\), or \(4096 \to 16384\))
The 4× expansion and contraction is standard but not derived from first principles, it was established empirically in the original paper and has remained the default.
Parameter Count: FFN Dominates
For a model with \(d_{\mathrm{model}} = 1024\) and \(d_{\mathrm{ff}} = 4096\), in each block:
| Sub-layer | Parameters |
|---|---|
| Multi-head attention (4 matrices) | \(4 \times 1024^2 = 4.2\mathrm{M}\) |
| FFN (2 matrices) | \(2 \times 1024 \times 4096 = 8.4\mathrm{M}\) |
The FFN holds twice as many parameters as the attention sub-layer. In a 96-layer model, FFNs collectively account for roughly 2/3 of all parameters.
What Does the FFN Actually Do?
Attention vs FFN: Division of Labour
Research into Transformer internals has revealed a rough division:
- Attention heads move information between positions, they determine which tokens influence each other and gather context
- FFN layers process information at a single position, they apply transformations and recall facts
This is why you can have a model that “knows” Paris is the capital of France even though that fact was not encoded in the positional attention pattern of the current context, the FFN retrieves it.
FFN as a Key-Value Memory
A 2020 paper (Geva et al., “Transformer Feed-Forward Layers Are Key-Value Memories”) showed that the FFN can be interpreted as:
- \(W_1\) rows (the “keys”): pattern detectors, each neuron in the expanded dimension activates for specific input patterns
- \(W_2\) columns (the “values”): for each activated key, the corresponding value vector is added to the output
When a token activates a key neuron (because it matches a learned pattern), the associated value is retrieved and added to the representation. This is analogous to a soft content-addressable memory, the FFN stores and retrieves (token, fact) associations.
The Nonlinearity: ReLU, GELU, SwiGLU
ReLU (original Transformer, 2017)
Simple and sparse, negative activations are exactly zero, which gives the FFN a sparse, efficient structure.
GELU (GPT-2, BERT, and successors)
Smooth approximation of ReLU with non-zero gradient for negative inputs. Empirically outperforms ReLU on most language tasks.
SwiGLU (LLaMA, PaLM, Mistral)
A gated variant: two parallel linear projections, one gating the other element-wise. SwiGLU-based FFNs use \(d_{\mathrm{ff}} = (8/3) \times d_{\mathrm{model}}\) (not 4×) to keep parameter count comparable. Consistently outperforms ReLU and GELU at large scale.
Position-Wise Independence: A Key Property
The FFN processes each token independently, it does not look at neighbouring tokens. There is no attention-like mechanism: the computation for position \(i\) uses only the vector at position \(i\).
This means:
- Parallelisable across positions (all tokens in a sequence processed simultaneously)
- No position-to-position information mixing, that is strictly the role of attention
- The FFN refines each token’s representation in place; it does not redistribute information
Worked Example: Parameter Count in GPT-3
GPT-3: \(d_{\mathrm{model}} = 12{,}288\) · \(d_{\mathrm{ff}} = 49{,}152\) (4×) · 96 layers
Per layer FFN parameters:
- \(W_1\): \(12{,}288 \times 49{,}152 = \mathbf{603.9M}\)
- \(W_2\): \(49{,}152 \times 12{,}288 = \mathbf{603.9M}\)
- Total FFN per layer: \(\approx \mathbf{1.21B}\)
Per layer MHA parameters (96 heads, \(d_k = d_v = 128\)):
- Q, K, V, O projections: \(4 \times 12{,}288^2 = \mathbf{603.9M}\)
Across 96 layers:
- All FFNs: \(96 \times 1.21\mathrm{B} \approx \mathbf{116B}\) parameters
- All MHA: \(96 \times 603.9\mathrm{M} \approx \mathbf{58B}\) parameters
- FFN share: \(\approx \mathbf{67\%}\) of the 175B total
This confirms the rule: in any standard Transformer, the FFN holds roughly two-thirds of all parameters. Scaling the model mostly means scaling the FFN.
Sparse FFNs: MoE
Mixture-of-Experts (MoE) Transformers replace the dense FFN with multiple expert FFNs, routing each token to only a subset (often 2 out of 64 or more experts):
token → router → expert_k → output
This allows vastly more total parameters (stored in expert FFNs) while keeping computation constant (only a fraction is used per token). Models like Mixtral 8×7B and GPT-4 (allegedly) use MoE in the FFN sub-layer.
Summary
| Property | Value |
|---|---|
| Architecture | Two-layer MLP with expansion |
| Expansion factor | 4× (ReLU/GELU) or 8/3× (SwiGLU) |
| Applied to | Each token independently |
| Parameter share | ~2/3 of total in standard models |
| Information role | Per-position processing and fact retrieval |
| Attention role comparison | Attention mixes positions; FFN refines each position |
The FFN is not attention’s sidekick. It is an equal partner, the knowledge storage and processing unit that sits beside attention’s information-routing mechanism.
References
- Vaswani, A., Shazeer, N., Parmar, N., Uszkoreit, J., Jones, L., Gomez, A. N., Kaiser, Ł., & Polosukhin, I. (2017). Attention Is All You Need. NeurIPS 2017 (Transformer FFN: two-layer MLP with ReLU, dimension 4d hidden, applied position-wise after each attention sublayer).
- Geva, M., Schuster, R., Berant, J., & Levy, O. (2021). Transformer Feed-Forward Layers Are Key-Value Memories. EMNLP 2021 (shows that FFN keys activate for human-interpretable input patterns and values store associated output information, FFN as learned key-value memory).
- Shazeer, N. (2020). GLU Variants Improve Transformer. arXiv 2020 (SwiGLU: gated linear units replacing ReLU in the FFN, now the dominant activation in LLaMA, Mistral, PaLM, and Gemini).