FFN with ReLU

§ 1 · Premise

The position-wise sublayer

The original transformer (Vaswani et al. 2017, “Attention Is All You Need”, §3.3) interleaves two kinds of sublayer in each block: multi-head self-attention, which mixes information across token positions, and a feed-forward network, which transforms each token’s representation in place. The feed-forward sublayer carries roughly two-thirds of the model’s parameter budget — at $d_{\text{model}} = 512$, $d_{\text{ff}} = 2048$, $H = 8$, the FFN holds $2 \cdot d_{\text{model}} \cdot d_{\text{ff}} \approx 2{,}097{,}152$ parameters per layer versus the attention block’s $4 \cdot d_{\text{model}}^2 \approx 1{,}048{,}576$. Whatever activation goes between the FFN’s two projections therefore has outsized influence on training dynamics, parameter-bound work, and the model’s representational capacity.

Three constraints shaped the 2017 choice.

Position-wise application. The FFN had to apply independently to each token’s hidden state. That ruled out anything with cross-token coupling (recurrence, convolutions of width > 1). Equivalently, the FFN can be described as two $1{\times}1$ convolutions over the sequence length (Vaswani et al. 2017, §3.3 footnote): same operation at every position, no parameter sharing across the depth direction.

Stable optimization with Adam. The transformer’s training recipe inherited Adam from the 2014 generation; the activation needed to interoperate with the Noam-Adam learning-rate schedule and the existing Xavier/Glorot initialization conventions without bespoke tuning. ReLU had a five-year track record of stable optimization (Glorot, Bordes & Bengio 2011, AISTATS; Krizhevsky, Sutskever & Hinton 2012, AlexNet) and was the default for almost every published deep architecture in 2016.

Throughput. Inference and training had to scale to a wall-clock budget that allowed weekly experiments on a 2017 8-GPU node. Tanh and sigmoid required exponentials and were known to saturate and slow training (Glorot & Bengio 2010, “Understanding the difficulty of training deep feedforward neural networks”, AISTATS); GELU (Hendrycks & Gimpel 2016) existed on arXiv but was untested at sequence-model scale; Swish would not be introduced until late 2017. ReLU was the conservative choice with the most prior evidence.

The 2017 FFN is what falls out: two linear projections with the simplest non-saturating nonlinearity between them, sized so the FFN’s parameter count roughly matches the attention block. Every subsequent FFN variant in this knowledge base — GELU, GeGLU, ReGLU, SwiGLU, Mixtral’s expert FFNs, DeepSeekMoE’s expert FFNs — keeps this two-projection scaffold and varies only what sits between.

§ 2 · Derivation

From a one-hidden-layer MLP to the transformer FFN

Start from a textbook one-hidden-layer MLP applied to a single token’s representation $\mathbf{x} \in \mathbb{R}^{d_{\text{model}}}$:

with $W_1 \in \mathbb{R}^{d_{\text{ff}} \times d_{\text{model}}}$, $W_2 \in \mathbb{R}^{d_{\text{model}} \times d_{\text{ff}}}$, biases of matching dimension, and elementwise nonlinearity $\phi$. This is the position-wise FFN with the row/column convention flipped. Vaswani et al. (2017, §3.3) wrote it in row-vector form,

which is the same operation; the choice of $\phi = \max(0, \cdot)$ is what makes this specifically the ReLU FFN.

Why ReLU specifically. Rewriting ReLU as a multiplicative gate exposes its structure:

ReLU multiplies its input by a hard $\{0, 1\}$ gate driven by the sign of the input. Three consequences follow:

• Sparsity. For symmetrically-initialized $W_1$ on a centered input, $\Pr[z > 0] \approx \tfrac{1}{2}$, so roughly half of the FFN’s $d_{\text{ff}}$ hidden units are zero at any step. Empirical work from the AlexNet era (Krizhevsky et al. 2012) and earlier (Glorot, Bordes & Bengio 2011, §3) showed this kind of activity-sparse code is both learnable and a decent prior for deep networks.
• Constant gradient on the active half. Where the gate is on, the gradient is exactly 1. Where it is off, the gradient is exactly 0. No gradient saturation, no vanishing-gradient problem from the activation itself.
• Dead-unit failure. If $W_1[i, :] \cdot \mathbf{x} + \mathbf{b}_1[i] < 0$ for every $\mathbf{x}$ ever seen, unit $i$ outputs zero and receives no gradient — it is permanently dead. The same Glorot et al. (2011) paper documents this failure mode at the same time it introduces the activation.

Sizing the hidden dimension. The original paper’s choice is $d_{\text{ff}} = 4 \cdot d_{\text{model}}$: at $d_{\text{model}} = 512$, $d_{\text{ff}} = 2048$. Two pieces of math motivate the ratio.

Parameter balance. The FFN holds $2 \cdot d_{\text{model}} \cdot d_{\text{ff}}$ parameters (biases are negligible). Multi-head attention, with $Q, K, V, O$ projections each of shape $d_{\text{model}} \times d_{\text{model}}$, holds $4 \cdot d_{\text{model}}^2$. Setting the two equal gives

The actual choice $d_{\text{ff}} = 4\,d_{\text{model}}$ is twice that — Vaswani et al. (2017, §3.3) do not explicitly defend the factor, but the doubling lets the FFN dominate the per-block parameter count by 2:1, which the rest of the field carried forward without questioning until GPT-2 (Radford et al. 2019, §2.1) and BERT (Devlin et al. 2019, §3.1) quietly inherited the same $4\times$ ratio.

Information capacity. A wider hidden layer gives the FFN more freedom to represent high-rank, low-correlation features per token. The Universal Approximation Theorem (Hornik 1991) shows one hidden layer of sufficient width can approximate any continuous function on a compact set; the $4\times$ ratio is empirically enough at transformer-base scale.

Parameter and FLOP count. Per layer, with batch $B$, sequence $T$, hidden $d = d_{\text{model}}$, expansion $f = d_{\text{ff}}/d_{\text{model}}$:

At $d = 512$, $f = 4$, this is $\sim 2.1$M parameters per FFN and $\sim 2.1$M FLOPs per token per layer. Across 6 transformer-base layers, the FFNs alone hold $\sim 12.6$M parameters — roughly two-thirds of the 65M total. This concentration of parameters is what makes the activation choice load-bearing.

Biases. The original FFN has additive bias terms $\mathbf{b}_1, \mathbf{b}_2$. Subsequent work (GPT-J, PaLM, Llama, Gemma) drops them: at modern scale they account for $<0.1\%$ of parameters but they slightly complicate fused-kernel layouts, and Chowdhery et al. (2022, PaLM, §2) report no measurable quality loss from removing them. The 2017 FFN keeps them; every modern descendant in this knowledge base does not.

§ 3 · Reference implementation

The 2017 sublayer in five lines

def ffn_relu(x, W1, b1, W2, b2):
    # x:  [B, T, d_model]
    # W1: [d_model, d_ff]    b1: [d_ff]
    # W2: [d_ff, d_model]    b2: [d_model]
    h = (x @ W1 + b1).clamp_min(0.0)   # ReLU = clamp_min(0)
    return h @ W2 + b2

Position-wise: every token at every position is mapped independently, with the same weights across positions. The clamp_min form makes the gating structure explicit; equivalently, F.relu(x @ W1 + b1) @ W2 + b2 or torch.maximum(x @ W1 + b1, 0) @ W2 + b2. Modern Llama-style FFNs replace this with a three-projection gated form; see the SwiGLU entry for the displacement.

§ 4 · Empirical evidence

What replaced it, and by how much

The 2017 transformer (Vaswani et al. 2017) shipped FFN-ReLU as part of the base-model specification (§5.1) and reported state-of-the-art BLEU on WMT14 En-De and En-Fr. Within five years, every dominant FFN benchmark moved to a different activation.

GELU’s margin. Shazeer (2020, “GLU Variants Improve Transformer”, arXiv 2002.05202, Table 1) ran the first systematic FFN-activation ablation on T5-base pretraining on C4. At matched FLOPs:

• $\mathrm{FFN}_{\mathrm{ReLU}}$: 1.997 log-perplexity
• $\mathrm{FFN}_{\mathrm{GELU}}$: 1.983
• $\mathrm{FFN}_{\mathrm{Swish}}$: 1.994

GELU beat ReLU by 0.014 in log-perplexity (~1.4% relative); Swish was nearly tied with ReLU. The empirical margin is small, but it is consistent: BERT (Devlin et al. 2019), GPT-2 (Radford et al. 2019), GPT-3 (Brown et al. 2020), and T5 (Raffel et al. 2020) all chose GELU and saw the same direction of improvement. By 2020, plain ReLU was no longer the default.

GLU variants’ larger margin. The same Shazeer (2020) ablation showed gated FFNs cutting another 0.04 in log-perplexity at matched parameters:

• $\mathrm{FFN}_{\mathrm{ReGLU}}$: 1.953
• $\mathrm{FFN}_{\mathrm{GEGLU}}$: 1.942
• $\mathrm{FFN}_{\mathrm{SwiGLU}}$: 1.944

Plain ReLU to gated SwiGLU is a $\sim 0.05$ log-perplexity gap — about 3× larger than the ReLU-to-GELU gap. From Llama 1 (Touvron et al. 2023, “LLaMA: Open and Efficient Foundation Language Models”, arXiv 2302.13971, §2.1) onward, the open-weights frontier ships SwiGLU or GeGLU; plain ReLU FFN does not appear in any production-decoder release in this knowledge base after 2022.

Where ReLU FFN still appears. Three settings preserve it. First, legacy inference: GPT-2 weights are still hosted and served by many providers, all running the original ReLU FFN. Second, ablation baselines: every FFN-activation paper since 2020 includes plain ReLU FFN as the bottom-of-table comparison (Shazeer 2020 Table 1; Su et al. 2024 RoPE-V2 Table 6; follow-ups). Third, retro-modern hybrids: the “SoLU” line (Elhage et al. 2022, Transformer Circuits Thread, §1) and some sparse-FFN work (Mirzadeh et al. 2023, “ReLU Strikes Back”, arXiv 2310.04564) revisit ReLU for interpretability and activation-sparsity reasons — not because it beats SwiGLU on perplexity, but because its hard-zero gate is easier to reason about. Mirzadeh et al. (2023, §3) report up to 90% FFN activation sparsity from ReLU with a small perplexity penalty, which can translate to inference-time skipping of dead units. This use case is research-scale; no frontier release ships it.

Public sensitivity studies. The $4\times$ hidden-dim ratio has been re-examined occasionally (Tay et al. 2022, “Scaling Laws vs Model Architectures”, arXiv 2207.10551, Table 4) but no systematic ablation of the ratio specifically for ReLU FFN exists at modern scale — the ratio question has migrated to the SwiGLU $8/3 \cdot d_{\text{model}}$ regime instead.

@article{arxiv1706_03762,
  title  = {Attention Is All You Need},
  author = {Ashish Vaswani and others (Google Brain)},
  year   = {2017},
  eprint = {1706.03762},
  archivePrefix = {arXiv},
  url    = {https://arxiv.org/abs/1706.03762}
}

@article{arxiv_1706_03762,
  title         = {Attention Is All You Need},
  author        = {Ashish Vaswani et al. (Google Brain)},
  year          = {2017},
  eprint        = {1706.03762},
  archivePrefix = {arXiv},
  url           = {https://arxiv.org/abs/1706.03762}
}

{
  "id": "arxiv_1706_03762",
  "type": "article-journal",
  "title": "Attention Is All You Need",
  "author": [
    {
      "literal": "Ashish Vaswani et al. (Google Brain)"
    }
  ],
  "issued": {
    "date-parts": [
      [
        2017
      ]
    ]
  },
  "URL": "https://arxiv.org/abs/1706.03762",
  "number": "1706.03762",
  "source": "arXiv"
}

TY  - JOUR
TI  - Attention Is All You Need
AU  - Ashish Vaswani et al. (Google Brain)
PY  - 2017
JO  - arXiv
AN  - arXiv:1706.03762
UR  - https://arxiv.org/abs/1706.03762
ER  -