The Feed-Forward Network

Attention finished by moving information between positions: every token’s vector now carries some of what it read from the tokens before it. The feed-forward network does the opposite kind of work. It touches one position at a time, with no sideways reach, and reshapes that single vector on its own. It is the fourth and last mechanism a block runs, and it is where most of a block’s parameters — and most of its arithmetic — actually live.

The reason it exists is that attention only gathers. It produces, for each position, a weighted blend of other positions’ values, but a blend is a linear combination, and stacking linear combinations only ever yields another linear combination. For the model to compute something genuinely nonlinear about what it gathered — to recognize a pattern, not just average one — it needs a step that bends. The feed-forward network is that step: it takes the emb_dim vector attention handed back, expands it into a wider space where a nonlinearity has room to work, and projects it back down.

6.1 What it computes

Here is the whole of it, from forge’s model.py.

class FeedForward(nn.Module):
    def __init__(self, config: ModelConfig) -> None:
        super().__init__()
        d, h, dt = config.emb_dim, config.hidden_dim, config.dtype
        self.fc1 = nn.Linear(d, h, bias=False, dtype=dt)
        self.fc2 = nn.Linear(d, h, bias=False, dtype=dt)
        self.fc3 = nn.Linear(h, d, bias=False, dtype=dt)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.fc3(nn.functional.silu(self.fc1(x)) * self.fc2(x))

Listing 6.1 · forge’s feed-forward network. Three linear projections and one gated product.

Three linear layers, no biases. fc1 and fc2 both widen x from emb_dim — 1024 in the 0.6B model — to hidden_dim, which is 3072; fc3 brings the result back down from 3072 to 1024. The dtype is threaded through so every weight matches the model it loads into. The work is the single expression in forward, and it reads from the inside out.

self.fc1(x) and self.fc2(x) are two independent wide views of the same input: two different learned projections, each producing a 3072-wide vector. nn.functional.silu wraps the first of them. SiLU is z * sigmoid(z) — a smooth activation that passes large positive values almost unchanged and squashes negative ones toward zero. That activated view then multiplies the other view, self.fc2(x), element by element. The product is the hidden signal, and fc3 sums it back down to the original width. This particular gated form — two projections, an activation on one, an elementwise product, a projection down — is called SwiGLU.

6.2 Why two projections and a gate

A plainer feed-forward network needs only one wide projection, an activation, and a projection back: fc3(silu(fc1(x))). SwiGLU adds the second projection fc2 and the elementwise multiply, and that addition changes what the activation does. On its own, silu(fc1(x)) would just transform the signal. Multiplied into fc2(x), it becomes a gate: for each of the 3072 hidden units, silu(fc1(x)) decides how much of fc2(x) passes through. One projection learns what to look for; the other learns what to pass; the product is content weighted by relevance, computed per unit.

The width is deliberate too. A block applies only one nonlinearity here, so it applies it in a space three times wider than the vector it started with. The extra room lets the activation separate features that overlap when packed into 1024 dimensions, and fc3 then compresses whatever survives back down to the width the residual stream expects.

6.3 Correct, to the bit

The feed-forward network gets the same treatment as the mechanisms before it: an independent implementation of the same three projections and the same gate, the same weights loaded into both, one seeded batch fed to each, and a required match. The companion driver builds forge’s feed-forward network and the reference’s, hands both a random (2, 9, 32) batch, and prints the largest disagreement:

Exactly zero, like RMSNorm and for the same reason: forge’s forward and the reference run the identical operations in the identical order — the same two projections, the same SiLU, the same elementwise product, the same projection down — so no rounding noise can open up between them. There is nowhere for a difference to hide, which makes any nonzero number here a real divergence rather than float dust. The checkpoint makes one appear.

With the feed-forward network in place, all four of a block’s mechanisms are built: RMSNorm sets the scale, rotary embeddings stamp position, grouped-query attention mixes across tokens, and the feed-forward network transforms each position on its own. The block from Chapter 2 is now whole — every line inside it is one you have read and checked against an independent implementation of the same architecture. What remains in Part I is not a new mechanism but a matter of memory. Attention re-reads every earlier token on every step, recomputing keys and values that never change once a token is in the past. The next chapter stores them once and reuses them. That is the KV cache.