The layer holding two thirds of the parameters

Everything so far has been attention: building the weights, adding order, splitting the head. But a transformer block has one more thing beside it, and its name is just feed-forward.

A short description is not a small component. Counting says the opposite.

Two thirds of the parameters

Count one block, with d_model = d and an expansion factor of 4.

attention      Wq, Wk, Wv, Wo   4 x d^2   = 4d^2
feed-forward   W1 (d->4d), W2 (4d->d)     = 8d^2
total                                       12d^2

8/12, exactly 2/3. At d=512 that is 1,048,576 parameters for attention against 2,097,152 for the feed-forward. Ten parts have been spent on a third of the parameters.

This layer does not mix tokens

What the feed-forward does is one line.

FFN(x) = W2 @ relu(W1 @ x + b1) + b2

x is one token’s vector. Feed a whole sentence and every row passes through on its own. Measured, that is exactly what happens.

np.allclose(F(X), np.vstack([F(X[i:i+1]) for i in range(n)]))   # True
np.allclose(F(X[perm]), F(X)[perm])                             # True

Part eight called permutation equivariance a problem for attention. For the feed-forward it is correct. Shuffle the places and the results should shuffle with them. Connecting tokens is attention’s entire job; this layer computes inside each token and nowhere else.

Reading a block that way is clean. Attention moves, feed-forward processes.

Without the nonlinearity, the expansion buys exactly nothing

The layer widens d to 4d and narrows back to d. Why widen? Without a nonlinearity the answer is no reason at all.

Drop the relu and W2(W1 x) = (W2 W1) x, where W2 W1 is a single d x d matrix whose rank cannot exceed min(d, 4d) = d.

W2 @ W1 is 64x64, rank 64  (ceiling min(64, 256) = 64)
replaced exactly by one unexpanded 64x64: residual 3.3e-16

32,768 parameters buying the same family of functions as 4,096. The widened space is pure waste. It is part four’s point about stacked linear layers collapsing into one, except here the waste shows up as an exact parameter count.

So what this layer really is is not the expansion but the nonlinearity sitting inside it.

What width buys: pieces

Then what exactly grows when you widen? A relu network is piecewise linear, and one neuron makes one kink. With a one-dimensional input, neuron i kinks at the single point x = -b_i / w_i. Width m gives at most m kinks and m+1 pieces.

Here is sin(3x) actually fitted at several widths.

   m     RMSE    kinks inside the interval
   2   0.6221         2
   4   0.2251         4
   8   0.0989         7
  16   0.0388        16
  32   0.0250        31

Width 2 has two kinks and cannot imitate sin at all. Sixteen times the width takes the error from 0.6221 to 0.0250. What grew is not expressive power in the abstract but the number of straight pieces available to trace a curve.

The choice of 4d reads out of this too. Wider means finer pieces, and parameters grow in proportion to width. Four is a conventional compromise between those, and plenty of models use 3 or 8/3 instead.

Reading it as keys and values

One more framing. Call a row of W1 a key k_i and a column of W2 a value v_i, and the layer rewrites as

FFN(x) = Σ relu(k_i · x + b_i) · v_i

k_i · x measures how well this token matches pattern i, and where it matches, that much of v_i is added. Look up, then write. As with part ten’s heads, it is a sum.

np.abs(F(X) - sum(np.outer(relu(X @ W1[i] + b1[i]), W2[:, i]) for i in range(m))).max()
# 3.3e-16

Each of the 4d neurons is one rule: “if this pattern, add this”. That is where the reading of a large model’s feed-forward as a memory store comes from.

Measured at random initialisation, though, 0.5030 of the neurons fire for a given token. Half of them responding at once is hard to call one rule apiece. Clean separation is a result of training, not something the structure hands over.

So

• 2/3 of a block’s parameters are feed-forward. 4d^2 against 8d^2
• The layer does not mix tokens. Attention moves, this layer processes
• Without a nonlinearity the 4x expansion is exactly pointless: 32,768 parameters buying the same family as 4,096, residual 3.3e-16
• What width buys is pieces. One neuron is one kink, width m gives m+1 pieces, and going from width 2 to 32 takes the error 0.6221 -> 0.0250
• The layer is exactly Σ relu(k_i · x) v_i. Look up, then write

Next time the blocks get stacked for real. The pieces have been examined one at a time; what remains is the order they have to be wired in for twenty layers to train at all.