Feed-Forward Network

Your program will compute the second half of a transformer block: a small two-layer network applied to every token on its own. Attention gathers information from other tokens; the feed-forward network (FFN, also called the MLP) then processes each token's gathered result by itself.

The shape: expand, bend, contract

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

x          D numbers      (one token)
W1 @ x     H numbers      "expand":   W1 is H rows x D columns
act(...)   H numbers      "bend":     a nonlinearity, element by element
W2 @ ...   D numbers      "contract": W2 is D rows x H columns

H is the hidden width, usually 4 * D (the original transformer: 512 to 2048; GPT-2 small: 768 to 3072). The output has the input's width because the block adds it back to the input in lesson 9.

Position-wise means the same W1 and W2 are applied to every token separately. No token reads another here, so a sequence of T tokens is just T independent calls.

Why the bend matters

Without act the two layers collapse: W2 (W1 x) is one matrix times x. The nonlinearity is what lets the network compute anything that is not a linear map, and attention itself is almost linear in its values, so the FFN is where the block gets its expressive power.

• ReLU: max(0, z). Hard zero for every negative input.
• GELU (GPT-2, BERT): a smooth version that lets small negative values leak through. The tanh form GPT-2 uses is 0.5 * z * (1 + tanh(0.79788456 * (z + 0.044715 * z^3))). Some values: gelu(-3) = -0.0036, gelu(-1) = -0.1588, gelu(0.5) = 0.3457, gelu(1) = 0.8412, gelu(3) = 2.9964.

Worked trace

D = 2, H = 3. x = [1, -2]. W1 rows [1,1], [1,-1], [0,2]; b1 = [0, 0, 1]; W2 rows [1,0,1], [0,2,0]; b2 = [0.5, 0].

W1 @ x + b1 :  1*1 + 1*-2 = -1      1*1 + -1*-2 = 3      0*1 + 2*-2 + 1 = -3
ReLU        :  [0, 3, 0]
W2 @ h + b2 :  row 0: 1*0 + 0*3 + 1*0 + 0.5 = 0.5
               row 1: 0*0 + 2*3 + 0*0 + 0   = 6
y = [0.5, 6]
with GELU in place of ReLU the hidden vector is [-0.1588, 2.9964, -0.0036]  ->  y = [0.3376, 5.9927]

Only the middle hidden unit survives ReLU, so only column 1 of W2 contributes to the output.

What it computes, and where the weights are

One productive reading: each row of W1 is a pattern detector, the matching column of W2 is what that detector writes to the output when it fires. A hidden width of 3072 means thousands of such rules per token. The parameter count shows why people say the knowledge lives in the MLPs:

• Attention at D = 768: W_Q, W_K, W_V, W_O are 4 matrices of D x D, about 2.4 million numbers.
• The FFN at H = 4D: W1 and W2 are 2 matrices of D x 4D, about 4.7 million numbers.

Roughly two thirds of a block's weights are in its FFN. Llama replaces GELU with SwiGLU, a gated variant with a third matrix, and Mixtral replaces one FFN by eight experts of which each token uses two.

Storage and orientation

In this exercise the matrices arrive flat and row-major. W1 is H x D: row i starts at index i*D. W2 is D x H: row i starts at index i*H. Papers write both W @ x and x @ W, so decide the orientation before you loop, and use the right stride for each matrix.

Common mistakes

• Using the same stride for both matrices. Square toy cases (H = D) still pass; the first H != D case reads the wrong entries.
• Applying the activation to the output as well. The result is added to the residual stream, which must be able to move a coordinate down as well as up; a final ReLU removes that.
• Forgetting a bias, or adding b2 before the second multiplication instead of after.
• Treating a sequence as one big vector. The FFN is applied row by row.

Your exercise

Write the activation and the FFN for one token; the starter calls it once per X line.