Feed-Forward Network

GPT-2

Easy

Position-wise FFN with GELU

Independent study note. Written from the public paper and official code. This is not TensorTonic Plus and does not reproduce their exercises, starter code, or tests. For the official version, subscribe on TensorTonic.

Overview

Each GPT-2 block has a position-wise feed-forward network after masked attention. Vaswani et al. (2017) introduced this as two linear maps with a ReLU in between and a hidden width of (4,d_{\mathrm{model}}). GPT and GPT-2 keep the 4× expansion and the position-wise contract, but swap ReLU for GELU (Hendrycks & Gimpel), the same nonlinearity GPT-1 used.

The 2019 paper does not write the FFN equation or name GELU in section 2.3. It says the model “largely follows” GPT, and the official mlp function is the concrete definition: expand with c_fc, GELU, project back with c_proj. The FFN is applied independently at each time step and never mixes tokens; all cross-position work is left to attention.

How it works

Let (x\in\mathbb{R}^{B\times T\times d}) with (d=) n_embd. The official width is n_state = nx * 4 where nx is the last dimension of (x), so (d_{\mathrm{ff}}=4d). For the 117M model, (d=768) and (d_{\mathrm{ff}}=3072). Then

\mathrm{FFN}(x)=\mathrm{GELU}(xW_1+b_1)\,W_2+b_2,

W_1\in\mathbb{R}^{d\times 4d},\; W_2\in\mathbb{R}^{4d\times d}.

conv1d in src/model.py is that linear map: a weight of shape [1, nx, nf] applied as a matmul plus bias. Scopes are c_fc (expand) and c_proj (contract). GELU is the tanh approximation also used in GPT:

\mathrm{GELU}(z)\approx 0.5\,z\left(1+\tanh\left[\sqrt{2/\pi}\,\bigl(z+0.044715\,z^3\bigr)\right]\right).

That is exactly gelu() in src/model.py. It is not the later erf form, and it is not ReLU.

Because the two matmuls are position-wise, one may reshape to ((BT)\times d), apply the MLP, and reshape back; the official code does this implicitly inside conv1d. There is no bias-free variant and no gated SwiGLU — those are later LMs.

In the surrounding block the FFN sees LN(x) (pre-norm) and its output is added to the residual. Those wrappers are not part of mlp itself. hparams is passed into mlp but unused; width is entirely nx*4.

Specified by the paper: Transformer LM following GPT, four model widths in Table 2. Specified by official code: GELU-tanh, 4× hidden, two biased conv1ds. Inferred: identity with Vaswani’s FFN up to the activation.

Official code

src/model.py: mlp (the FFN), gelu (activation), conv1d (linear). block calls mlp(norm(x, 'ln_2'), 'mlp', nx*4, hparams=hparams). Variable names c_fc / c_proj match released checkpoints. No other FFN implementation is in the public tree.

Watch-outs

• GELU ≠ ReLU ≠ GELU-erf. A ReLU FFN will not match GPT-2 activations; using scipy.special.erf instead of the tanh cubic also drifts.
• Hidden size is (4d_{\mathrm{model}}), not (4\times d_k). On 117M that is 3072, not (4\times 64=256).
• Both layers have biases. Dropping (b_1) or (b_2) (common in some “linear is 2D weight only” ports) disagrees with conv1d.
• c_proj here is the FFN output projection. Attention has its own c_proj. Reusing one tensor for both is a checkpoint-loading bug.

Sources

• Paper: Language Models are Unsupervised Multitask Learners (Radford et al., 2019), §2.3
• FFN pattern: Vaswani et al., “Attention Is All You Need,” 2017
• Code: openai/gpt-2 src/model.py (mlp, gelu, conv1d)