A Minimal Transformer Block

A transformer block wraps self-attention and a feed-forward network (MLP) in residual connections with layer normalization. Stack a dozen or more of these blocks and you have GPT, BERT, or any modern transformer.

a = LayerNorm(x)
x = x + SelfAttn(a, a, a)
b = LayerNorm(x)
x = x + MLP(b)

Two sublayers, each wrapped the same way:

1. Normalize.
2. Transform (attention, or MLP).
3. Add the result back to the input (residual connection).

Pre-Norm vs Post-Norm

The original "Attention Is All You Need" paper (2017) normalized after the residual add (post-norm):

x = LayerNorm(x + SelfAttn(x, x, x))
x = LayerNorm(x + MLP(x))

Modern transformers (GPT-2 onward) normalize before each sublayer instead (pre-norm), as shown at the top of this lesson. Pre-norm keeps a clean residual path from input to output with no normalization in the way, which stabilizes gradients in very deep stacks. Post-norm can train unstable at depth without careful learning-rate warmup. Pre-norm is now the default choice for training stability.

The Residual Connection

x = x + sublayer(norm(x)) — the + x is what lets gradients flow straight through every block, unimpeded, all the way back to the input embedding. Without it, stacking 12+ blocks would make backpropagation vanish long before reaching early layers (the same problem ResNet solved for CNNs in 2015).

The MLP Sublayer

The feed-forward network in each block is intentionally simple: two linear layers with a nonlinearity between them, expanding to a wider hidden dimension and projecting back down.

This is where a transformer does most of its "thinking" per token — attention moves information between tokens, the MLP processes information within a token. Roughly two-thirds of a transformer's parameters live in these MLP layers.

Self-Attention Recap

Each block reuses the scaled dot-product attention from the previous lesson. With identity Q/K/V projections (as in this exercise), attention reduces to:

scores[i][j] = (x_i . x_j) / sqrt(D)
weights = softmax(scores)
output_i = sum_j weights[i][j] * x_j

For a sequence of length 1, softmax over a single score is always 1.0, so attention becomes a no-op: the output equals the (normalized) input unchanged. This is a useful sanity check when testing your implementation.

Putting a Full Model Together

A real transformer decoder is: embedding + positional encoding → N stacked blocks (like the one you're building here) → final LayerNorm → output projection. GPT-3 uses 96 of these blocks; a small toy model might use 2-4. Everything you need for a working (if tiny) transformer is now covered across this chapter's two lessons.