Transformer Block Forward Pass

Implement a single transformer block forward pass.

A transformer block consists of:

Multi-head self-attention with residual connection and layer norm: x = LayerNorm(x + Attention(x))
Feed-forward network with residual connection and layer norm: x = LayerNorm(x + FFN(x))

For simplicity:

Input:

Output: Output matrix after one transformer block, rounded to 4 decimal places.

Input:

2 2
1 0
0 1

Output:

[ 0.9998 -0.9998]
[-0.9998  0.9998]

Reasoning:

First, we calculate the self-attention: since it's single-head attention with no projection weights, Attention(x) simplifies to $x \cdot x^T$ normalized by the square root of the sequence length ( $\sqrt{2}$ ). For the input matrix $X = \begin{bmatrix} 1 & 0 \\ 0 & 1 \end{bmatrix}$ , Attention(x) becomes $\begin{bmatrix} 1 & 0 \\ 0 & 1 \end{bmatrix}$ .
Then, we apply the layer norm to the sum of x and Attention(x): $x = \frac{x + Attention(x) - mean}{\sqrt{var + 1e-5}}$ . Since $x + Attention(x) = \begin{bmatrix} 2 & 0 \\ 0 & 2 \end{bmatrix}$ , $mean = 1$ , and $var = 0$ , $x$ becomes $\begin{bmatrix} 2 - 1 \\ 2 - 1 \end{bmatrix} / \sqrt{0 + 1e-5} \approx \begin{bmatrix} 1 & 0 \\ 0 & 1 \end{bmatrix}$ , but due to the specifics of the layer norm calculation, it actually results in $\begin{bmatrix} 1 & 0 \\ 0 & 1 \end{bmatrix}$ .
Next, the feed-forward network (FFN) is applied: $FFN(x) = ReLU(x \cdot W1) \cdot W2$ . Given $d_{ff} = 4 \cdot d = 8$ , $W1$ is a $2 \times 8$ matrix and $W2$ is an $8 \times 2$ matrix, both initialized with $\np.random.seed(42)$ and scale $0.1$ . After computation, $FFN(x)$ is calculated and then added to $x$ .
Finally, another layer norm is applied to the sum of $x$ and $FFN(x)$ , resulting in the output matrix after one transformer block, which after computation and rounding to 4 decimal places gives the output: $\begin{bmatrix} 0.9998 & -0.9998 \ -0.9998 & 0.9998 \end{b