12. The position-wise feed-forward network

Attention has allowed positions to exchange information. Now each position needs time to process what it retrieved. The feed-forward sublayer performs that different job by transforming every position independently with the same learned function.

The paper uses two linear transformations with a ReLU between them:

FFN(x) = max(0, xW₁ + b₁)W₂ + b₂

In the course model, a 128-feature vector expands to 512 features and contracts back to 128:

self.input_projection = nn.Linear(128, 512)
self.activation = nn.ReLU()
self.output_projection = nn.Linear(512, 128)

PyTorch applies a linear layer over the final dimension, so a tensor shaped batch × time × 128 remains batched and sequential while each position passes through the same network.

attention:    information moves between positions
feed-forward: each position -> expand -> ReLU -> contract -> same position

Proving that positions do not mix

Change one input position and run the feed-forward layer twice. The changed position should produce a different output; every other position should remain bit-for-bit identical.

.venv/bin/python lessons/09-feed-forward-and-normalisation/example.py

The retained example reports exactly zero change at unaffected positions. That property contrasts with attention, where changing one visible position may intentionally influence later positions.

Why the hidden layer is wider

The expansion provides more intermediate features in which ReLU can create nonlinear combinations. Without a nonlinear activation, two consecutive linear layers could collapse into one linear transformation and add little expressive power.

The feed-forward network accounts for a substantial share of Transformer parameters. It is easy to overlook because the attention diagram is more visually distinctive, but attention alone would only remix value vectors. The nonlinear per-position computation gives each block additional capacity to transform what it retrieved.

Follow the code

PositionWiseFeedForward is only three modules, but the testable contract matters: the shape is preserved and changing one position cannot alter another.

In your fork, remove the ReLU and rerun the focused tests. The shapes will still work, which is a useful reminder that shape tests alone cannot tell us whether two linear layers have collapsed into an effectively linear operation.

Next: residuals, LayerNorm, and dropout