Feed-Forward Networks

A transformer block is two things: attention, and this. Attention is the part everyone studies and the part the papers are named after. The feed-forward block is where the parameters are.

Project up to a wider space, apply a nonlinearity, project back. The conventional width is $d_{ff} = 4d$, which gives $8d^2$ parameters against attention’s $4d^2$ — hence two thirds of the model, in the half of the block that gets a third of the attention.

Position-wise

The same $W_1$ and $W_2$ are applied to every token, independently. Nothing in the block sees a neighbour. This is the division of labour that makes the transformer work: attention is the only operation that mixes positions, and the feed-forward block is the only one with the capacity to transform what it finds there.

Why 4×

The ratio is inherited from the original paper and has never been strongly justified. What evidence exists says the block is not very sensitive to it — between roughly $2d$ and $8d$, quality tracks parameter count rather than the ratio itself, so the choice is a question of how you would rather spend a parameter budget.

Two things have moved it. Gated variants use three matrices instead of two, so the width is scaled by $\tfrac{2}{3}$ to hold the parameter count fixed — $d_{ff} \approx \tfrac{8}{3}d$, usually rounded to a multiple of 256 for the hardware. And mixture-of-experts replaces the single block with many, at which point the per-expert width becomes a granularity question rather than a capacity one.

from torch import Tensor, nn


class FeedForward(nn.Module):
    def __init__(self, dim: int, hidden: int | None = None, bias: bool = False):
        super().__init__()
        hidden = hidden or 4 * dim
        self.up = nn.Linear(dim, hidden, bias=bias)
        self.down = nn.Linear(hidden, dim, bias=bias)
        self.act = nn.GELU()

    def forward(self, x: Tensor) -> Tensor:                   # [B, T, D]
        return self.down(self.act(self.up(x)))

Biases are gone from most modern implementations. They are the smallest tensors in the model and removing them measurably improves training stability at scale — one of several places where a parameter that costs nothing turns out to be worth less than nothing.

The mechanistic reading is that $W_1$‘s rows are pattern detectors and $W_2$‘s columns are what gets written back when a pattern fires — a key-value memory, with $d_{ff}$ entries per layer. That view is well supported and it explains a great deal: why factual knowledge localises to feed-forward blocks, why editing a fact means editing $W_2$, and why the block’s activations are extremely sparse in practice even though nothing asked them to be.