Four Parts, One Pass

Self-attention, add and normalize, feed-forward, add and normalize. Four parts, in that order, and that is the entire block — the thing GPT stacks ninety-six times without changing anything about it.

Task: write transformer_block(x, Wq, Wk, Wv, W1, b1, W2, b2) for a single masked attention head, returning the block's output, rounded to 4 decimal places. Both normalizations have their learned scale and shift at the neutral values $1$ and $0$, so there is nothing extra to apply.

1 — Masked self-attention. With $Q = x W_q$, $K = x W_k$ and $V = x W_v$,

$\text{score}_{ij} = \frac{Q_i \cdot K_j}{\sqrt{d_k}} \qquad \text{for } j \le i$

Row $i$ may use positions $0$ through $i$ only; everything to its right is blocked. Softmax each row over the positions it is allowed to see, then blend the values: $a_i = \sum_{j \le i} w_{ij} V_j$.

2 — Add and normalize. $h = \operatorname{LN}(x + a)$.

3 — Feed-forward, each position on its own. $f = \operatorname{ReLU}(h W_1 + b_1)\, W_2 + b_2$, where $\operatorname{ReLU}$ keeps positive entries and sets negative ones to $0$.

4 — Add and normalize again. The block returns $\operatorname{LN}(h + f)$.

$\operatorname{LN}$ works on each position's vector separately, using that vector's own $d$ slots:

$\operatorname{LN}(v)_j = \frac{v_j - m}{\sqrt{\dfrac{1}{d}\sum_{l=1}^{d}(v_l - m)^2 + 10^{-5}}}, \qquad m = \frac{1}{d}\sum_{l=1}^{d} v_l$

Round only the numbers you return.

Two things are worth noticing once it runs. The output has exactly the shape of the input, which is the property that lets you stack this ninety-six times. And the second residual adds $h$, not $x$: every sublayer adds to whatever it was handed, so one running workspace travels all the way up the stack.