Out Wide and Back Again

Inside every transformer block, right after attention, sits the feed-forward sublayer. It takes the vector at one position, widens it, applies a non-linearity, and narrows it back to its original width:

If each position's vector holds $d$ numbers and the wide middle holds $h$:

• $W_1$ has $d$ rows and $h$ columns, and $b_1$ holds $h$ numbers.
• $W_2$ has $h$ rows and $d$ columns, and $b_2$ holds $d$ numbers.
• ReLU sets every negative number to 0 and leaves the rest alone.

Written out, hidden value $j$ is $\text{ReLU}\left(\sum_i x_i\, W_1[i][j] + b_1[j]\right)$ and output value $k$ is $\sum_j \text{hidden}_j\, W_2[j][k] + b_2[k]$.

The same weights are used at every position, and each position is handled on its own. No position's output depends on any other position.

Task: write feed_forward(x, w1, b1, w2, b2).

• x is a sequence: a list of $n$ position vectors, each a list of $d$ numbers.
• w1 and w2 are lists of rows; b1 and b2 are plain lists.
• Return a list of $n$ output vectors, each of $d$ numbers, with every value rounded to 4 decimal places.