The Transformer Block

Course role: v1 is a conceptual bridge with compact scalar exercises. For the cumulative implementation sequence, numerical boundary cases, and learned multi-head projections, continue with deep learning v2. The v1 transformer exercise uses a single attention head even though its API accepts num_heads.

Goal

Assemble self-attention, layer normalization, feed-forward networks, and residuals into the core transformer block used by GPT-style models.

Prerequisites: Attention module.

---

1) Block structure (pre-norm)

For input x with shape (T, D):

y = x + Attn(LN(x))
z = y + FFN(LN(y))

This is the pre-norm layout used in GPT-2 and later for stability.

---

2) LayerNorm (per position)

LayerNorm normalizes across the feature dimension for each position independently:

mu  = mean(x)
var = mean((x - mu)^2)
ln(x) = (x - mu) / sqrt(var + eps)
output = gamma * ln(x) + beta

gamma and beta are learnable vectors of length D.

---

3) Feed-forward network (FFN)

Two linear layers with a nonlinearity:

FFN(x) = W2 * act(W1 * x + b1) + b2

Standard size is 4D for the hidden dimension. In this pack, ReLU is fine (GELU is a common modern variant).

---

4) Residual connections

Residuals preserve information and improve gradient flow:

x -> x + sublayer(x)

Without residuals, deep transformers are much harder to train.

---

5) Implementation notes (list-of-Value)

• Apply LayerNorm to each position: [ln(pos) for pos in x]
• SelfAttention returns (T, D)
• Residual sums are elementwise list additions

---

6) Single-head vs multi-head

The conceptual block often uses multi-head attention, but the pack code uses a single-head SelfAttention for clarity. The block API still passes num_heads for future extension.

---

Work a normalization and residual trace

For one position x=[1,3], mean is 2 and population variance is 1. With gamma=[1,1], beta=[0,0], and epsilon 1e-5, LayerNorm returns approximately [-0.999995,0.999995]. It normalizes the feature vector at this position; averaging across other sequence positions would create the wrong operation and can break causal dependencies.

To isolate residual routing, suppose the attention branch returns [0.2,-0.1] and the FFN branch after the second normalization returns [0.3,0.4]. First y=x+attention=[1.2,2.9], then z=y+FFN=[1.5,3.3]. The second residual uses y, not the original x. The skip path contributes an identity derivative in addition to derivatives through the branches.

Before seeing the check, set both branches' outputs identically to zero. What must the block return, even though LayerNorm still runs?

Check the residual invariant

It returns x exactly: y=x+0=x, then z=y+0=x. In a finite-difference check the whole block's input Jacobian is identity when both branch functions are identically zero. Nonzero normalization parameters do not change that direct residual result.

This module's quiz checks the design; the executable block is implemented inside the next mini-GPT exercise. If you want an independently graded block before the capstone, v2 transformer blocks provide that intermediate practice.

Key takeaways

1. TransformerBlock = pre-norm attention + pre-norm FFN with residuals.
2. LayerNorm is per-position, across features.
3. FFN adds per-position nonlinearity and capacity.
4. Residuals are essential for deep training stability.

Next: combine blocks into a full GPT-style language model.