A Minimal Transformer Block

A transformer is a stack of identical blocks. Each block lets tokens exchange information with attention, lets every token think on its own with a small feed-forward network, and keeps both steps stable with residual connections and layer normalization. After this lesson your program assembles those pieces, which you already built, into a working block.

The two sublayers

a  = LayerNorm(x)
x  = x + SelfAttention(a)        mix information between tokens
b  = LayerNorm(x)
x  = x + MLP(b)                  transform each token on its own

Each sublayer does the same three things: normalize, transform, then add the result back onto the input. That addition is the residual connection. Think of x as a shared stream that every block reads from and writes small corrections into.

Why the residual path matters

Because the output is x + something, the derivative of the output with respect to x always contains an identity term. Gradients flow straight back through a stack of 96 blocks without passing through 96 multiplications that could shrink them, the same trick that made 100-layer image networks (ResNets) trainable. It also means a freshly initialised block with near-zero sublayers acts as the identity, so deep stacks start out well behaved.

Pre-norm and post-norm

The original 2017 transformer normalized after each addition (post-norm):

x = LayerNorm(x + SelfAttention(x))
x = LayerNorm(x + MLP(x))

Modern models normalize at the start of each sublayer (pre-norm, shown above). Pre-norm leaves the residual stream completely untouched by normalization, which keeps gradients steady in very deep stacks; post-norm tends to need careful learning-rate warm-up. The exercise supports both so you can see the numbers differ.

The MLP sublayer

Two linear layers with a ReLU between them: widen from D to a hidden size H (often 4*D), apply max(0, .), then project back down. Attention moves information between positions; the MLP processes each position independently. Because H is wide, roughly two thirds of a block's parameters live in the MLP.

Self-attention here

The attention is the one from the previous lesson with Q, K and V all equal to the normalized input (identity projections, one head). A decoder-only language model also applies the causal mask. For a sequence of one token the single attention weight is 1, so attention just returns that token's normalized vector.

A worked trace

D = 2, one token x = [2, 6], H = 2, W1 the identity, b1 = 0, W2 = [[2, 0], [0, 3]], b2 = [0.5, 0], pre-norm.

LayerNorm(x)           = [-1.0000, 1.0000]       (mean 4, std 2)
attention (1 token)    = [-1.0000, 1.0000]
x = x + attention      = [ 1.0000, 7.0000]
LayerNorm(x)           = [-1.0000, 1.0000]       (mean 4, std 3)
MLP: relu(W1 * that)   = [0, 1]    then W2*[0,1] + b2 = [0.5, 3.0]
x = x + MLP            = [ 1.5000, 10.0000]

What a real model adds

Learned Q, K, V and output projections per head, many heads in parallel, token embeddings plus positional information at the bottom, a final layer norm and a projection to the vocabulary at the top, and dropout. GPT-3 stacks 96 such blocks; the toy you build here has the same skeleton.

Common mistakes

• Putting the normalization after the residual add when pre-norm was asked for, or the reverse.
• Adding the sublayer output to the normalized vector instead of to the original stream.
• Forgetting that the second sublayer starts from the stream after the first residual.
• Using one set of norm parameters for both norms.
• Applying the ReLU after the second linear layer.

Your exercise

The starter holds the parameters, a working layer_norm and self_attention (the earlier lessons' code), and the command loop. Implement mlp and block_forward, which runs one block in pre-norm or post-norm form.