Transformer Block

Your program will assemble everything so far into one block of a GPT: normalise, attend with a causal mask, add back, normalise, feed forward, add back. All the weights come from a seed, so a short input fixes a whole model and your output can be checked number for number.

The block

One block takes a matrix x of T rows (tokens) by D columns and returns a matrix of the same shape. Pre-norm, as in GPT-2:

x1  = x  + Attention( LN1(x) )          first sublayer:  mix information between tokens
out = x1 + FFN( LN2(x1) )               second sublayer: process each token on its own

Attention(h):  q = h @ Wq    k = h @ Wk    v = h @ Wv          (each T x D)
               split q, k, v into H heads of width D/H          (lesson 6)
               per head and per query i: softmax over keys j <= i of q_i . k_j / sqrt(D/H)   (lessons 5, 8)
               mix the v rows, join the heads, then  @ Wo       (lesson 6)
FFN(h):        W2 @ gelu( W1 @ h ), per row, biases 0           (lesson 7)
LN:            gamma 1, beta 0, eps 1e-5, per row               (lessons 3, 9)

Nothing in the block is new. What is new is the order and where the additions happen. The two + signs are the residual connections of lesson 9. The skip carries the raw x and x1, and only the sublayer inputs are normalised.

Parameters

Per block: Wq, Wk, Wv, Wo are D x D (4 D^2 numbers), W1 is F x D and W2 is D x F (2 D F, which is 8 D^2 at the usual F = 4D). That is 12 D^2 per block, so GPT-2 small (D = 768) has about 7.1 million per block and 85 million across its 12. LayerNorm adds only 2D numbers each.

A model from a seed

To make the test input short, the weights are generated, not listed. A small random generator fills them in a fixed order:

state  <- seed                                  (32-bit unsigned)
next:  state <- (state * 1664525 + 1013904223) mod 2^32
       u     <- (state >> 8) / 16777216          a number in [0, 1), exact in a double
weight <- u - 0.5

For seed = 1 the first draw is state = 1015568748, u = 0.2364..., weight -0.2636. The matrices are drawn in the order Wq, Wk, Wv, Wo, W1, W2, each row by row. The starter already contains the generator, the matrices, and helpers for LayerNorm, matrix products, softmax and GELU.

Orientation

For the attention weights the row vector comes first: q = h @ Wq, so q[c] = sum over j of h[j] * Wq[j][c]. For the FFN the matrix comes first, as in lesson 7: W1 @ h gives F numbers with W1 being F x D.

Debugging a block

• One token. With T = 1 the softmax has a single entry equal to 1, so attention returns LN(x) @ Wv @ Wo. If your one-token output is wrong, the problem is not in the heads or the mask.
• No peeking. Change the last row of the input: the earlier output rows must stay identical, because of the causal mask.
• Shape. Every intermediate is T x D except the FFN hidden layer, which is T x F.

Common mistakes

• Feeding the normalised vector into the skip connection (x1 = LN(x) + attn). The skip adds the raw input.
• Normalising x a second time. The second LayerNorm is applied to x1, the output of the first residual add.
• Drawing the weights in the wrong order, or column by column. The first number drawn is Wq[0][0], the second Wq[0][1].
• Scaling the scores by sqrt(D) instead of sqrt(D/H).

Your exercise

Write causal_mha (the heads, scores, mask and mixing) and block (the wiring above). The starter reads the input, seeds the generator and prints your result.