LM4-T3 — Feed-Forward Layers and GELU

Learning objectives

After this lesson, you should be able to:

• Describe the feed-forward network of a Transformer block and give its shapes.
• Explain why a nonlinearity is required between the two linear layers.
• Evaluate GELU at a few inputs and compare it with ReLU.
• Compute the parameters and FLOPs of the feed-forward layers and their share of a block.

The idea: a small network applied to every token

After attention, each position’s vector contains information gathered from other tokens. The FFN takes that vector and transforms it: x → expand to 4d numbers → apply a nonlinearity → project back to d. The same two weight matrices are used for every position, which is why the layer is called position-wise: no information passes between positions inside it. Positions interact only in attention.

Why expand?

The wide hidden layer gives the network room to compute many features of the token’s vector at once and combine them. Expanding by a factor of four is the convention in GPT-2 and the original Transformer; recent models vary it, and some use gated variants (Module 10). A common way to think about the FFN is as a memory: attention decides what to read, and the FFN holds a large table of learned patterns that recognise some features of the vector and write something back. That view is a useful picture and an active research topic, not a proven fact about any particular model.

Why a nonlinearity?

Two linear layers in a row are one linear layer: W₂(W₁x) = (W₂W₁)x. Without something nonlinear in the middle, widening to 4d would add parameters and no expressive power. The activation function bends the space so that the composition can represent functions that a single matrix cannot.

GELU

The original Transformer used ReLU, which outputs the input if it is positive and 0 otherwise. GPT-2 uses the Gaussian Error Linear Unit. It multiplies the input by the probability that a standard normal variable is below it: large positive inputs pass almost unchanged, large negative ones are almost zeroed, and in between the transition is smooth. Compared with ReLU it has a non-zero gradient for small negative inputs, which avoids “dead” units and trains slightly better in practice. The exact form uses the error function; GPT-2 used a tanh approximation, which the book and this track follow.

The mathematics: expansion, GELU and the cost

For one token vector x of width d:

Applied to a tensor of shape (B, n, d) the same weights act on the last axis, so the hidden tensor has shape (B, n, 4d) and the output (B, n, d).

d

the model width; 4d is the hidden width (the expansion factor 4 is a convention).

W₁, b₁

the up-projection and its bias: 4d·d + 4d parameters.

W₂, b₂

the down-projection and its bias: d·4d + d parameters.

GELU

applied element-wise to the 4d hidden numbers.

GELU. The definition and the tanh approximation:

Φ(x)

the cumulative distribution function of the standard normal distribution: the probability that a N(0,1) variable is below x.

√(2/π)

a constant, about 0.7979.

0.044715

a fitted constant of the approximation.

Values (tanh form): GELU(−3) = −0.0036, GELU(−1) = −0.1588, GELU(−0.5) = −0.1543, GELU(0) = 0, GELU(0.5) = 0.3457, GELU(1) = 0.8412, GELU(2) = 1.9546, GELU(3) = 2.9964. The function dips to about −0.17 near x = −0.75 and then returns to 0.

Parameters. The FFN has

(4d² + 4d for the first layer and 4d² + d for the second). For d = 768: 8 · 589,824 + 5 · 768 = 4,722,432. Attention has 4d² + d = 2,360,064 (no bias on Q, K, V), so the FFN holds 4,722,432 / (4,722,432 + 2,360,064) = 66.7% of the weights of the two sublayers.

Compute. Each weight is used for one multiply-add per token, so the FFN costs about 2 · 8d² = 16d² FLOPs per token, twice the 8d² of the four attention projections. Together with the attention-matrix products (Lesson LM3-T6), the FFN is the main cost for sequences shorter than about 2d.

A worked example: a tiny feed-forward network

Take d = 2 with hidden width 4. Let W₁ = [[1, 0], [0, 1], [1, 1], [1, −1]], b₁ = 0, W₂ = [[1, 0, 0, 0], [0, 1, 0, 0]] (it reads out the first two hidden units), b₂ = 0, and input x = [1, −2].

Hidden pre-activation: W₁x = [1, −2, −1, 3]. After GELU: [0.8412, −0.0454, −0.1588, 2.9964]. (GELU(−2) = −0.0454.) Output: W₂ times that = [0.8412, −0.0454]. With ReLU the hidden vector would be [1, 0, 0, 3] and the output [1, 0].

The whole transformation is applied independently to every token of every sequence in the batch.

import math
import torch
import torch.nn as nn
import torch.nn.functional as F

def gelu(x):
    return 0.5 * x * (1 + torch.tanh(math.sqrt(2 / math.pi) * (x + 0.044715 * x ** 3)))

xs = torch.tensor([-3.0, -1.0, -0.5, 0.0, 0.5, 1.0, 2.0, 3.0])
print([round(v, 4) for v in gelu(xs).tolist()])      # [-0.0036, -0.1588, -0.1543, 0.0, 0.3457, 0.8412, 1.9546, 2.9964]
print(torch.allclose(gelu(xs), F.gelu(xs, approximate="tanh"), atol=1e-6))   # True

d = 768
ffn = nn.Sequential(nn.Linear(d, 4 * d), nn.GELU(approximate="tanh"), nn.Linear(4 * d, d))
nparams = sum(p.numel() for p in ffn.parameters())
print(nparams, nparams == 8 * d * d + 5 * d)       # 4722432 True

x = torch.randn(2, 5, d)                             # (B, n, d)
print(tuple(ffn[0](x).shape), tuple(ffn(x).shape))   # (2, 5, 3072) (2, 5, 768)
x2 = x.clone(); x2[:, 0] += 1.0                      # change token 0 only
print(torch.allclose(ffn(x)[:, 1:], ffn(x2)[:, 1:])) # True: no information moves between positions

The last line demonstrates what “position-wise” means: changing token 0 does not change the FFN output at any other position.

Engineering implications and limitations

Activation memory. The hidden tensor is four times wider than the residual stream: (B, n, 4d). For B = 8, n = 1,024, d = 768 it has 25.2 million entries (100 MB in float32) per layer, kept for the backward pass.

The FFN is the part to quantise and prune first. It has the most parameters. Methods in Module 8 and Module 10 (LoRA, quantisation, mixture of experts) often target these matrices.

Use the exact GELU you trained with. The tanh approximation and the exact erf form differ by up to about 10^-3. Loading pretrained weights into a model with the other variant gives small but real differences; use approximate="tanh" for GPT-2.

Bias terms. GPT-2 has them in the FFN. Many newer models do not. Count parameters with the setting of your configuration.

Gated variants. Modern models replace this FFN with a gated one (SwiGLU) with three matrices and a hidden width of about 8d/3, so that the parameter count stays near 8d². The structure “expand, nonlinearity, project” is unchanged.

Mixture of experts. Some large models replace the single FFN with many FFNs (“experts”) and route each token to a few of them, increasing parameters without increasing per-token compute. The routing layer is new; the experts are FFNs like this one (Module 10).

For robots. The feed-forward layers are where a language model’s knowledge appears to be stored, which is one reason why fine-tuning for a robot task (Module 8) usually adapts these matrices with low-rank updates instead of retraining everything.

Common misconceptions

Engineering takeaways

1. The FFN is FFN(x) = W₂ GELU(W₁x + b₁) + b₂ with a hidden width of 4d, applied independently at every position.
2. A nonlinearity is essential: two linear layers collapse into one.
3. GELU(x) = x·Φ(x), approximated by 0.5·x·(1 + tanh(√(2/π)(x + 0.044715x³))); GELU(1) = 0.8412.
4. The FFN has 8d² + 5d parameters (4,722,432 for d = 768), about two thirds of a block, and costs about 16d² FLOPs per token.
5. Its hidden tensor is four times wider than the residual stream, which matters for activation memory.

References and further reading

See Chapter 4 of the book (implementing a feed-forward network with GELU activations) and the PyTorch documentation for torch.nn.GELU.

• Hendrycks, D. and Gimpel, K. (2016). Gaussian Error Linear Units (GELUs). arXiv:1606.08415.
• Vaswani, A. et al. (2017). Attention Is All You Need. arXiv:1706.03762 (Section 3.3, position-wise feed-forward networks).
• Build a Large Language Model (From Scratch), Chapter 4: Implementing a GPT Model from Scratch (primary book reference)

Knowledge check

Three conceptual questions. Write an answer, then reveal the explanation. Your answers stay in your browser.

What’s next

Next: the Transformer block that combines attention and the feed-forward network, and the assembly of a whole GPT.

Related content

Nothing below is required to finish this lesson.

• Before
  LM4-T2: Layer Normalization and Residual Connections

  Normalise activations and add skip connections, and explain how both help deep networks train.

• Lab
  LM4-L2: Layer Normalization and Residual Numerics

  Implement layer norm and a residual update and see why activations stay stable.