A LLaMA-Style Transformer Block

The problem

Time to assemble the pieces into the block that models in the style of LLaMA (Large Language Model Meta AI) stack dozens of times. A block has two sub-layers, each wrapped the same way:

• MHA (multi-head attention), where the tokens exchange information;
• an MLP (multilayer perceptron) of the SwiGLU (Swish-gated linear unit) kind, where each token is processed on its own;
• before each sub-layer, an RMSNorm (root mean square layer normalization).

The block

For x of shape $(B, T, D)$:

where $\mathrm{RMSNorm}_1$ and $\mathrm{RMSNorm}_2$ have their own gains. Read each line as: take a normalized copy of the stream, run the sub-layer on it, and add the result to the stream itself (the residual connection). The stream is never replaced, only corrected.

Example. A useful check. If the output matrices of both sub-layers are zero ($W_o = 0$ in attention, $W_2 = 0$ in the MLP), both sub-layers output zero, so $h = x + 0$ and $y = h + 0$: the block returns x unchanged. A block that normalizes the stream itself, instead of a copy, fails this check.

The SwiGLU MLP

where:

• $u$ is the normalized input, $W_1$ and $W_3$ have shape $(D, F)$ and widen each vector to $F$ numbers, and $W_2$, of shape $(F, D)$, narrows it back;
• $\mathrm{SiLU}(z) = z\,\sigma(z)$ is the SiLU (sigmoid linear unit): $z$ times its sigmoid $\sigma(z) = 1 / (1 + e^{-z})$;
• $\odot$ is the elementwise product: the branch $u W_3$ acts as a gate that scales each entry of the other branch.

Example. SiLU of a few numbers: $\mathrm{SiLU}(2) = 2 \cdot 0.881 \approx 1.762$, $\mathrm{SiLU}(0) = 0$, $\mathrm{SiLU}(-2) \approx -0.238$, $\mathrm{SiLU}(20) \approx 20$, $\mathrm{SiLU}(-20) \approx 0$. Large positive inputs pass almost unchanged, large negative ones are switched off, like ReLU (rectified linear unit) but smooth.

What to implement

• silu(z), elementwise.
• swiglu_mlp(u, W1, W3, W2).
• transformer_block(x, p, n_heads), where p is a dict with "norm1", "norm2" (RMSNorm gains, shape $(D,)$), "Wq", "Wk", "Wv", "Wo" (shape $(D, D)$) and "W1", "W3", "W2".

rmsnorm(x, gamma) and mha_forward(x, Wq, Wk, Wv, Wo, n_heads) (causal, from the earlier exercises) are given. To keep the block focused, it has no RoPE (rotary position embedding).

Requirements

• Pre-norm: normalize the input of each sub-layer, and add the sub-layer's output to the un-normalized stream.
• The second sub-layer reads $h$, the output of the first, not $x$.
• No loops.

Hint

For a sigmoid that never overflows, $\sigma(z) = \tfrac{1}{2}\left(1 + \tanh(z/2)\right)$ (section 2).