Chapter 5 · Section 5.3 · 45 min

The feed-forward network

Here is a question that sounds like a trick. Attention lets every token gather information from every other token, with learned, input-dependent weights. Why does a Transformer block need anything else?

Look at what attention outputs. For each token, it returns a weighted average of value vectors, and each value vector is a linear projection of a token's stream. Attention is superb at moving information between positions, but once the weights are fixed, what it computes about the moved information is linear. You met the precise version of this worry at the end of §4.7. A stack of pure self-attention layers, with no skip connections and no MLPs, collapses toward a rank-1 output (every token ending up with the same vector) doubly exponentially fast in depth (Dong, Cordonnier & Loukas, 2021). The skip connections of §5.1 are half of the cure, and the MLP is the other half.

So every block pairs attention with a second sublayer, a small multilayer perceptron applied to each token on its own. It is the plainest component in the Transformer and holds about two thirds of its parameters. By the end of this section you will be able to write it in three variants, count it exactly, explain why nearly every modern LLM uses a gated version with a strange hidden size like 11,008, and read it as a key-value memory that stores much of what the model knows.

The MLP sublayer, applied to each token

For one token's stream vector $\mathbf{x} \in \R^{d}$ , written as a row, the feed-forward network (FFN) is

\operatorname{FFN}(\mathbf{x}) = \phi\big(\mathbf{x}\mathbf{W}_1 + \mathbf{b}_1\big)\,\mathbf{W}_2 + \mathbf{b}_2, \qquad \mathbf{W}_1 \in \R^{d \times d_{\text{ff}}},\ \mathbf{W}_2 \in \R^{d_{\text{ff}} \times d},

where $d_{\text{ff}}$ is the hidden width, $\mathbf{b}_1 \in \R^{d_{\text{ff}}}$ and $\mathbf{b}_2 \in \R^{d}$ are biases, and $\phi$ is an elementwise nonlinearity. For a sequence $\mathbf{X} \in \R^{T \times d}$ the same weights are applied to every row, $\operatorname{FFN}(\mathbf{X}) = \phi(\mathbf{X}\mathbf{W}_1 + \mathbf{b}_1)\mathbf{W}_2 + \mathbf{b}_2$ , which is why it is often called position-wise. No information crosses between tokens here. That is attention's job.

In plain words, project each token up into a wider space, apply a nonlinearity there, and project back down to the stream.

Counting it

The FFN has $2\,d\,d_{\text{ff}} + d_{\text{ff}} + d$ parameters. The original Transformer used $d = 512$ and $d_{\text{ff}} = 2048$ , an expansion ratio of 4, and almost every dense model since kept that ratio or something matched to it. With $d_{\text{ff}} = 4d$ ,

\text{FFN parameters} = 8d^2 + 5d, \qquad \text{attention parameters} = 4d^2 + 4d,

since attention has four $d \times d$ projections ( $\mathbf{W}_Q, \mathbf{W}_K, \mathbf{W}_V, \mathbf{W}_O$ , see §4.5). So the FFN is two thirds of every block's weights. Its compute is in the same proportion, because a matrix multiply costs 2 FLOPs per weight per token (one multiply, one add). Forward, the FFN costs $2 \cdot 2\,d\,d_{\text{ff}} = 16d^2$ FLOPs per token against $8d^2$ for attention's projections.

A worked example with GPT-2 small, $d = 768$ and $d_{\text{ff}} = 3072$ . The FFN has $2 \cdot 768 \cdot 3072 + 3072 + 768 = 4{,}722{,}432$ parameters per layer, and the attention sublayer has $4 \cdot 768^2 + 4 \cdot 768 = 2{,}362{,}368$ . Across 12 layers the FFNs hold about 56.7 million of the model's 124 million parameters. Section 5.6 counts the rest of the block, and §6.2 the rest of the model.

Why a ratio of 4? There is no theorem. Vaswani et al. chose it, and it has survived because wider hidden layers buy capacity cheaply. The FFN is a dense matrix multiply, the operation GPUs do best, so adding hidden units costs little wall-clock time per parameter. Mixture-of-experts models (Chapter 13) push the same idea further by keeping many FFNs and activating only a few per token.

Activation functions, from ReLU to GELU and SiLU

The original Transformer used ReLU, $\phi(z) = \max(0, z)$ . BERT and GPT-2 switched to GELU, and LLaMA-style models use SiLU inside a gate. Here is where GELU comes from, because its definition explains its shape.

Hendrycks & Gimpel (2016) started from dropout. Dropout multiplies a unit's input by a random 0 or 1. Suppose instead the mask depends on the input itself, $m \sim \operatorname{Bernoulli}(\Phi(z))$ , where $\Phi$ is the standard normal CDF, so larger inputs are more likely to be kept. The expected output is

\operatorname{GELU}(z) = \mathbb{E}[m\,z] = z\,\Phi(z) = \frac{z}{2}\Big(1 + \operatorname{erf}\big(z/\sqrt 2\big)\Big).

In plain words, GELU passes an input through with a probability that rises smoothly from 0 to 1 as the input grows, and returns the average.

For large positive $z$ , $\Phi(z) \to 1$ and GELU is the identity. For large negative $z$ , it is 0. In between it is smooth and dips slightly below zero, reaching its minimum of about $-0.170$ at $z \approx -0.752$ . Its derivative is $\Phi(z) + z\,e^{-z^2/2}/\sqrt{2\pi}$ , which is positive for every input above about $-0.75$ , zero only at the minimum, and slightly negative below it, where it fades smoothly toward zero (it never drops below about $-0.13$ ). Compare that with ReLU, whose gradient is exactly zero for every negative input, so a unit that drifts negative on all inputs can stop learning forever (a "dead" unit).

Computing $\operatorname{erf}$ used to be slow, so the paper gives an approximation that GPT-2 and many codebases use,

\operatorname{GELU}(z) \approx \frac{z}{2}\Big(1 + \tanh\Big(\sqrt{2/\pi}\,\big(z + 0.044715\,z^3\big)\Big)\Big).

Its maximum error is about $4.7 \times 10^{-4}$ . The paper also gives a cruder sigmoid form, $z\,\sigma(1.702z)$ , which is off by about $2 \times 10^{-2}$ near $|z| \approx 2$ (the last line of the cell below). Both errors are small, but not zero, and a model trained with one variant should be run with the same one.

The SiLU, $\operatorname{SiLU}(z) = z\,\sigma(z)$ with $\sigma$ the logistic sigmoid, is the same idea with the logistic CDF in place of the Gaussian one. It was named in the GELU paper, used in reinforcement learning by Elfwing et al. (2018), and rediscovered by an automated search as "Swish" (Ramachandran et al., 2017, who proposed $z\,\sigma(\beta z)$ , of which SiLU is the case $\beta = 1$ ). Its minimum is about $-0.278$ at $z \approx -1.278$ .

One more activation is worth knowing, because it keeps turning up in efficient recipes. So et al. (2021) searched over Transformer components for language modeling, and one of the two changes their search found most useful was squared ReLU, $\phi(z) = \max(0, z)^2$ . It keeps ReLU's exact zeros, so most hidden units are exactly off for any given token, while growing faster than linearly for large inputs. NVIDIA's Nemotron-4 models use it. Like the choice between ReLU and GELU, this is an empirical finding about what trains well, not a derived optimum.

Python · cell 1

Warm-up

GELU, exact and approximate

Write gelu(x) (the exact erf form) and gelu_tanh(x) (the tanh approximation) for NumPy arrays of any shape. NumPy has no vectorized erf, so use np.vectorize(math.erf) or scipy.special.erf.

import numpy as np
import math

def gelu(x):
    raise NotImplementedError

def gelu_tanh(x):
    raise NotImplementedError

Hidden tests, hints, and AI help, from $1.99 a month.

Core

The FFN forward pass and its parameter count

Write ffn(X, W1, b1, W2, b2) using the tanh GELU (as GPT-2 does) for X of shape (..., d), and ffn_param_count(d, d_ff, bias=True, gated=False). A gated FFN (next subsection) has three $d \times d_{\text{ff}}$ -sized matrices, and its biases (when present) are $2d_{\text{ff}} + d$ . An ungated one has two matrices and biases $d_{\text{ff}} + d$ .

import numpy as np

def ffn(X, W1, b1, W2, b2):
    raise NotImplementedError

def ffn_param_count(d, d_ff, bias=True, gated=False):
    raise NotImplementedError

Hidden tests, hints, and AI help, from $1.99 a month.

Gated linear units and SwiGLU

In 2020 Noam Shazeer ran a simple experiment. Replace the FFN's first layer with a gated linear unit (Dauphin et al., 2017), which computes two projections and multiplies them elementwise, one passed through a nonlinearity,

\operatorname{FFN}_{\text{GLU}}(\mathbf{x}) = \big(\sigma(\mathbf{x}\mathbf{W}_g) \odot \mathbf{x}\mathbf{W}_u\big)\mathbf{W}_d, \qquad \mathbf{W}_g, \mathbf{W}_u \in \R^{d\times d_{\text{ff}}},\ \mathbf{W}_d \in \R^{d_{\text{ff}}\times d}.

We call $\mathbf{W}_g$ the gate projection, $\mathbf{W}_u$ the up projection, and $\mathbf{W}_d$ the down projection, the names used in the Hugging Face LLaMA code (gate_proj, up_proj, down_proj). Meta's original code calls the same three matrices w1, w3, and w2. Shazeer tried other functions in place of the sigmoid and dropped the biases. The variant that took over the field uses SiLU (he called it Swish),

\operatorname{FFN}_{\text{SwiGLU}}(\mathbf{x}) = \big(\operatorname{SiLU}(\mathbf{x}\mathbf{W}_g) \odot \mathbf{x}\mathbf{W}_u\big)\mathbf{W}_d.

In plain words, compute two different linear features of the token, squash one of them into a soft on/off switch, use it to scale the other, and project back down.

Matching the parameter count, the 8/3 rule

A gated FFN has three matrices instead of two. To compare fairly, Shazeer shrank the hidden width so the parameter count (and the FLOPs) stayed the same. Setting $3\,d\,d_{\text{ff}}' = 2\,d\,(4d)$ gives

d_{\text{ff}}' = \tfrac{2}{3}\cdot 4d = \tfrac{8}{3}\,d .

For his $d = 768$ baseline that meant $d_{\text{ff}}' = 2048$ instead of 3072. LLaMA uses the same rule, then rounds up to a multiple of 256 so the matrices tile nicely on GPUs. For $d = 4096$ , $\tfrac{8}{3}\cdot 4096 = 10{,}922.67$ , the integer part is $10{,}922$ , and rounding up to a multiple of 256 gives $43 \times 256 = 11{,}008$ . That is where LLaMA-7B's odd hidden size comes from. Later models multiply by an extra factor before rounding and round to coarser multiples. LLaMA 2 70B and LLaMA 3 use a factor of 1.3 with multiples of 4,096 or 1,024, giving 28,672 for $d = 8192$ and 14,336 for $d = 4096$ .

What gating bought

Shazeer trained T5-base-sized models on the span-filling task, with every variant matched in parameters and compute. The held-out log-perplexities after 524,288 steps (lower is better, from his Table 1) are below.

FFN variant	hidden layer	log-perplexity
ReLU (baseline)	$\max(0, \mathbf{x}\mathbf{W}_1)$	1.677
GELU	$\operatorname{GELU}(\mathbf{x}\mathbf{W}_1)$	1.679
Swish	$\operatorname{SiLU}(\mathbf{x}\mathbf{W}_1)$	1.683
GLU	$\sigma(\mathbf{x}\mathbf{W}_g)\odot\mathbf{x}\mathbf{W}_u$	1.663
Bilinear	$\mathbf{x}\mathbf{W}_g\odot\mathbf{x}\mathbf{W}_u$	1.648
ReGLU	$\max(0, \mathbf{x}\mathbf{W}_g)\odot\mathbf{x}\mathbf{W}_u$	1.645
GEGLU	$\operatorname{GELU}(\mathbf{x}\mathbf{W}_g)\odot\mathbf{x}\mathbf{W}_u$	1.633
SwiGLU	$\operatorname{SiLU}(\mathbf{x}\mathbf{W}_g)\odot\mathbf{x}\mathbf{W}_u$	1.636

Two things jump out. Changing the activation alone (ReLU, GELU, Swish) barely matters, while every gated variant wins. And the bilinear layer, which has no nonlinearity at all apart from the product, beats all three ungated ones. The paper's conclusion famously declines to explain the result and attributes it to "divine benevolence". This is an empirical finding, and the best explanations we have are hypotheses.

Multiplicative interactions. A gated unit computes a product of two linear functions of the input, a second-order feature that an ungated unit would need several units to approximate. A linear path for gradients. The up branch is linear in the input, so gradients flow through it without passing through a saturating nonlinearity, much as the gates of an LSTM (§0.8) protect the cell state. Input-dependent feature selection. The gate decides, token by token, which hidden features are active, which is a soft version of the sparsity that mixture-of-experts makes explicit.

PaLM, LLaMA, Mistral, Qwen, and DeepSeek use SwiGLU. Gemma and T5 v1.1 use GEGLU. Because $d_{\text{ff}}$ is shrunk to compensate, the gate costs nothing extra, and the FLOPs per token stay at $16d^2$ .

Interview

A LLaMA-style SwiGLU FFN

Implement the two pieces you would write for a LLaMA-style model.

swiglu_ffn(X, W_gate, W_up, W_down) computes $(\operatorname{SiLU}(\mathbf{X}\mathbf{W}_g)\odot\mathbf{X}\mathbf{W}_u)\mathbf{W}_d$ for X of shape (..., d), with no biases. SiLU is $z\,\sigma(z)$ .
llama_hidden_dim(d, multiple_of=256, ffn_dim_multiplier=None) returns the hidden width. Start from int(2 * (4 * d) / 3), multiply by ffn_dim_multiplier and take int again if one is given, then round up to a multiple of multiple_of.

import numpy as np

def swiglu_ffn(X, W_gate, W_up, W_down):
    raise NotImplementedError

def llama_hidden_dim(d, multiple_of=256, ffn_dim_multiplier=None):
    raise NotImplementedError

Hidden tests, hints, and AI help, from $1.99 a month.

The backward pass of a gated FFN

If you are asked to backpropagate through SwiGLU, name every intermediate first. With $\mathbf{G} = \mathbf{X}\mathbf{W}_g$ , $\mathbf{U} = \mathbf{X}\mathbf{W}_u$ , $\mathbf{A} = \operatorname{SiLU}(\mathbf{G})$ , $\mathbf{H} = \mathbf{A}\odot\mathbf{U}$ , and $\mathbf{Y} = \mathbf{H}\mathbf{W}_d$ , and with $\partial\mathcal{L}/\partial\mathbf{Y}$ given, the chain rule (§0.2) gives, in order,

\frac{\partial\mathcal{L}}{\partial\mathbf{W}_d} = \mathbf{H}^\top\frac{\partial\mathcal{L}}{\partial\mathbf{Y}},\quad \frac{\partial\mathcal{L}}{\partial\mathbf{H}} = \frac{\partial\mathcal{L}}{\partial\mathbf{Y}}\mathbf{W}_d^\top,\quad \frac{\partial\mathcal{L}}{\partial\mathbf{U}} = \frac{\partial\mathcal{L}}{\partial\mathbf{H}}\odot\mathbf{A},\quad \frac{\partial\mathcal{L}}{\partial\mathbf{G}} = \frac{\partial\mathcal{L}}{\partial\mathbf{H}}\odot\mathbf{U}\odot\operatorname{SiLU}'(\mathbf{G}),

then $\partial\mathcal{L}/\partial\mathbf{W}_u = \mathbf{X}^\top\,\partial\mathcal{L}/\partial\mathbf{U}$ , $\partial\mathcal{L}/\partial\mathbf{W}_g = \mathbf{X}^\top\,\partial\mathcal{L}/\partial\mathbf{G}$ , and $\partial\mathcal{L}/\partial\mathbf{X} = (\partial\mathcal{L}/\partial\mathbf{G})\mathbf{W}_g^\top + (\partial\mathcal{L}/\partial\mathbf{U})\mathbf{W}_u^\top$ , the sum because $\mathbf{X}$ feeds both branches. The SiLU derivative follows from $\sigma' = \sigma(1-\sigma)$ ,

\operatorname{SiLU}'(z) = \sigma(z) + z\,\sigma(z)\big(1 - \sigma(z)\big) = \sigma(z)\big(1 + z(1 - \sigma(z))\big).

In plain words, the gate and the up branch each receive the hidden gradient scaled by the other branch's value (the product rule), the gate's share is further scaled by SiLU's slope, and the input collects gradient from both branches.

Interview

Backpropagate through SwiGLU

Write swiglu_backward(dY, X, W_gate, W_up, W_down) returning (dX, dW_gate, dW_up, dW_down) for X of shape (N, d) and dY of shape (N, d). Recompute the forward intermediates inside the function.

import numpy as np

def swiglu_backward(dY, X, W_gate, W_up, W_down):
    """Returns dX, dW_gate, dW_up, dW_down."""
    raise NotImplementedError

Hidden tests, hints, and AI help, from $1.99 a month.

The FFN is a key-value memory

Now rewrite the FFN in a way that changes how you see it. Let $\mathbf{k}_i \in \R^d$ be the $i$ -th column of $\mathbf{W}_1$ and $\mathbf{v}_i \in \R^d$ the $i$ -th row of $\mathbf{W}_2$ . Then $\mathbf{x}\mathbf{W}_1$ is the vector of dot products $(\mathbf{x}\cdot\mathbf{k}_1, \dots, \mathbf{x}\cdot\mathbf{k}_{d_{\text{ff}}})$ , and multiplying by $\mathbf{W}_2$ sums the rows of $\mathbf{W}_2$ weighted by the hidden activations,

\operatorname{FFN}(\mathbf{x}) = \sum_{i=1}^{d_{\text{ff}}} \phi\big(\mathbf{x}\cdot\mathbf{k}_i + b_{1,i}\big)\,\mathbf{v}_i \;+\; \mathbf{b}_2 .

In plain words, the FFN holds $d_{\text{ff}}$ memories. Each memory has a key that is compared with the current token, and a value that is added to the stream in proportion to how strongly the key matched.

Put this next to attention, $\sum_n a_n\mathbf{v}_n$ with $a_n = \operatorname{softmax}(\mathbf{q}\cdot\mathbf{k}_n/\sqrt{d_k})$ . The shape is the same. The differences are what make them complementary. Attention's keys and values come from the tokens in the context and change with every input. The FFN's keys and values are parameters, the same for every input, so they store knowledge learned in training. And attention's weights are normalized by a softmax to sum to one, while the FFN's are unnormalized and mostly near zero, so many memories can fire at once or none at all. Sukhbaatar et al. (2019) took the analogy literally and replaced the FFN with attention over learned "persistent" key-value vectors, without losing quality.

Geva et al. (2021) then checked whether real language models use their FFNs this way. They found that individual keys in a trained model fire on human-interpretable patterns in the input prefixes (shallow layers on surface patterns, like a particular word ending a phrase, upper layers on semantic ones, like a topic). Projecting a value vector through the output embedding gives a distribution over the vocabulary, and in upper layers it tends to concentrate on the tokens likely to follow the key's pattern. The layer's output is a composition of many memories, which later layers refine through the residual stream. Meng et al. (2022) located factual associations ("the Eiffel Tower is in Paris") in middle-layer FFNs and edited them with a single rank-one update to $\mathbf{W}_2$ . These are empirical findings about trained models, not theorems, and real memories are messier than a lookup table. Many units respond to several unrelated patterns, a symptom of superposition, the packing of more features than neurons into one layer.

Here is a hand-built FFN that stores three facts, so you can see the mechanism with nothing learned.

Python · cell 2

The query "Tokyo" activates only its own key, and the write points at "Japan". The token "the" matches nothing above the threshold and writes nothing. A trained FFN does this with thousands of overlapping, superposed memories and smooth activations, but the bookkeeping is the same. The widget lets you move the query and watch memories fire.

Core

Decompose an FFN into its memories

Write ffn_memories(x, W1, b1, W2, k) for a ReLU FFN without output bias. For one token x of shape (d,), return (coef, top, partial) where coef (shape (d_ff,)) holds every memory's activation $\max(0, \mathbf{x}\cdot\mathbf{k}_i + b_{1,i})$ , top is an int array of the k memories with the largest activations (largest first), and partial is the sum of just those k memories' contributions $\sum_{i \in \text{top}} \text{coef}_i\,\mathbf{v}_i$ .

import numpy as np

def ffn_memories(x, W1, b1, W2, k):
    raise NotImplementedError

Hidden tests, hints, and AI help, from $1.99 a month.

In PyTorch

This is the FFN as LLaMA writes it, with the three projections named as in its code. Note the absence of biases and the hidden size computed by the rule you implemented.

🔥 In PyTorchfor a machine with PyTorch

FFN	matrices	params at matched size	activations	used by
ReLU MLP	2	$8d^2$ ( $d_{\text{ff}} = 4d$ )	ReLU	original Transformer, T5 v1.0
GELU MLP	2	$8d^2$	GELU (tanh or exact)	BERT, GPT-2, GPT-3
squared-ReLU MLP	2	$8d^2$	$\max(0, z)^2$	Primer, Nemotron-4
GEGLU	3	$8d^2$ ( $d_{\text{ff}} = 8d/3$ )	GELU gate times linear	T5 v1.1, Gemma
SwiGLU	3	$8d^2$ ( $d_{\text{ff}} = 8d/3$ )	SiLU gate times linear	PaLM, LLaMA, Mistral, Qwen, DeepSeek
mixture of experts	3 per expert	many times $8d^2$ , few active	SwiGLU per expert	Mixtral, DeepSeek-V3

🎯 Interview questions on this section

If attention already mixes information across tokens, why does a Transformer need the FFN?

Attention outputs convex combinations of linear projections of the tokens, so it moves information well but computes little about it per token. Pure attention stacks without MLPs and skips collapse toward rank-1 outputs doubly exponentially in depth (Dong et al., 2021). The FFN adds a per-token nonlinear transformation with a large hidden width, and it is where much of the model's stored knowledge lives (it holds two thirds of the block's parameters at the standard 4x ratio).

Write the SwiGLU FFN and explain the 8/3 factor.

$\operatorname{FFN}(\mathbf{x}) = (\operatorname{SiLU}(\mathbf{x}\mathbf{W}_g)\odot\mathbf{x}\mathbf{W}_u)\mathbf{W}_d$ with $\operatorname{SiLU}(z) = z\sigma(z)$ and no biases. It has three $d\times d_{\text{ff}}$ matrices instead of two, so to keep the parameters and FLOPs of a standard $d_{\text{ff}} = 4d$ MLP, set $3d\,d_{\text{ff}} = 8d^2$ , which gives $d_{\text{ff}} = 8d/3$ . LLaMA computes $\lfloor 8d/3\rfloor$ , optionally multiplies by a factor, and rounds up to a multiple of 256, so $d = 4096$ gives 11,008.

Derive GELU and compare it with ReLU.

Multiply the input $z$ by a random mask $m\sim\operatorname{Bernoulli}(\Phi(z))$ , where $\Phi$ is the standard normal CDF, and take the expectation, $\mathbb{E}[mz] = z\Phi(z)$ . It equals the identity for large $z$ and zero for very negative $z$ , is smooth everywhere, and dips to about $-0.17$ near $z = -0.75$ . Unlike ReLU its gradient is nonzero for moderately negative inputs, so units do not die, and its smoothness helps optimization. In Shazeer's matched experiments, though, swapping ReLU for GELU barely changed perplexity, while adding a gate helped a lot.

Explain the view of the FFN as a key-value memory. How does it differ from attention?

Writing the columns of $\mathbf{W}_1$ as keys $\mathbf{k}_i$ and rows of $\mathbf{W}_2$ as values $\mathbf{v}_i$ , the FFN is $\sum_i\phi(\mathbf{x}\cdot\mathbf{k}_i + b_i)\mathbf{v}_i$ , a lookup that adds each value in proportion to how well its key matches the token. Attention has the same form, but its keys and values are computed from context tokens and its weights are softmax-normalized. The FFN's keys and values are learned parameters that persist across inputs and its weights are unnormalized, so it behaves like a long-term memory. Geva et al. (2021) found keys that detect interpretable input patterns and values that promote likely next tokens, and Meng et al. (2022) edited facts stored in mid-layer FFNs.

How many FLOPs does an FFN cost per token, and how does it compare with attention?

Each weight costs one multiply and one add per token, so the forward pass costs $2\times$ the parameter count, $2\cdot 2d\,d_{\text{ff}} = 16d^2$ for $d_{\text{ff}} = 4d$ (and the same for SwiGLU at $8d/3$ ). Attention's four projections cost $8d^2$ , plus the score and weighted-sum terms that grow with context length, $4Td$ per token for a context of $T$ . So for $T$ well below $2d$ the FFN dominates. Section 5.6 turns this into the full count.

Why might a bilinear layer, with no nonlinearity, beat a ReLU MLP of the same size?

The elementwise product of two linear projections is itself nonlinear in the input (it is quadratic), so the layer can represent multiplicative feature interactions directly. Both branches are linear in the input, which gives clean gradient paths. In Shazeer's experiments the bilinear FFN reached 1.648 log-perplexity against 1.677 for ReLU at matched parameters. The honest answer is that this is an empirical result whose full explanation is still open.

🧭 What to remember

The FFN, $\phi(\mathbf{x}\mathbf{W}_1 + \mathbf{b}_1)\mathbf{W}_2 + \mathbf{b}_2$ , transforms each token separately. At $d_{\text{ff}} = 4d$ it holds $8d^2$ parameters, two thirds of a block, and costs $16d^2$ FLOPs per token.
GELU is $z\Phi(z)$ , the expected output of an input-dependent dropout mask. Its tanh approximation is close but not identical.
Gated FFNs multiply a nonlinear gate by a linear branch. SwiGLU, $(\operatorname{SiLU}(\mathbf{x}\mathbf{W}_g)\odot\mathbf{x}\mathbf{W}_u)\mathbf{W}_d$ , is the modern default, with $d_{\text{ff}} = 8d/3$ to match parameters, which is how 11,008 arises.
Gating improved perplexity in matched experiments while the choice of activation alone barely mattered. Why gating helps is not settled.
An FFN is a key-value memory, $\sum_i\phi(\mathbf{x}\cdot\mathbf{k}_i)\mathbf{v}_i$ , with learned, persistent keys and values. Trained models appear to store patterns and facts this way.
Next, §5.4 tackles what attention and the FFN both lack, any notion of where a token sits in the sequence.

Shazeer (2020). GLU Variants Improve Transformer. The matched comparison of gated FFNs, in four pages.
Hendrycks & Gimpel (2016). Gaussian Error Linear Units (GELUs). The derivation of GELU, its approximations, and the naming of SiLU.
Dauphin, Fan, Auli & Grangier (2017). Language Modeling with Gated Convolutional Networks. Where the gated linear unit comes from.
Geva, Schuster, Berant & Levy (2021). Transformer Feed-Forward Layers Are Key-Value Memories. The memory view, tested on a trained language model.
Meng, Bau, Andonian & Belinkov (2022). Locating and Editing Factual Associations in GPT. Facts in mid-layer FFNs, and editing them.
Dong, Cordonnier & Loukas (2021). Attention is Not All You Need. Why pure attention collapses without MLPs and skips.
Sukhbaatar et al. (2019). Augmenting Self-attention with Persistent Memory. Replacing the FFN with attention over learned keys and values.
So et al. (2021). Primer, Searching for Efficient Transformers for Language Modeling. The architecture search that found squared ReLU.

Zero to Frontier is made for larger screens

The feed-forward network

The MLP sublayer, applied to each token

Counting it

Activation functions, from ReLU to GELU and SiLU

GELU, exact and approximate

The FFN forward pass and its parameter count

Gated linear units and SwiGLU

Matching the parameter count, the 8/3 rule

What gating bought

A LLaMA-style SwiGLU FFN

The backward pass of a gated FFN

Backpropagate through SwiGLU

The FFN is a key-value memory

Decompose an FFN into its memories

In PyTorch

Read more

Watch next