Lesson 82: LayerNorm, FFN with GELU, and the Full Encoder Block

1. LayerNorm, FFN with GELU & the Full Encoder Block

Title

USAAIO · Lesson 82 · Phase 3 (Transformers)

The two sublayers that wrap multi-head attention: LayerNorm (normalize over features, not batch), GELU feed-forward network, residual connections, and the complete pre-norm EncoderBlock as a reusable nn.Module.

2. By the end of this lesson you can

Objectives

1. Implement LayerNorm from scratch and verify it against nn.LayerNorm
2. Explain why layer norm (over features) beats batch norm for variable-length sequences
3. Contrast pre-norm (Norm → sublayer) and post-norm and state which is more stable
4. Build the FFN sublayer (two Linear layers + GELU, d_ff = 4·d_model)
5. Assemble the full EncoderBlock nn.Module with residual connections and both sublayers

3. What survived from Multi-Head Attention?

Warm-up

4. LayerNorm — normalize over features

Section

Part 1 of 3

5. What LayerNorm does

Concept

LayerNorm normalizes each token's feature vector independently: compute the mean and variance across the d_model features for that one token, then rescale.

\[ \hat{x}_i = \frac{x_i - \mu}{\sqrt{\sigma^2 + \varepsilon}}, \quad \mu = \frac{1}{d}\sum_{i=1}^d x_i, \quad \sigma^2 = \frac{1}{d}\sum_{i=1}^d (x_i - \mu)^2 \]

After normalization, learnable parameters γ (scale) and β (shift) — both of shape d_model — let the network undo the normalization if needed. nn.LayerNorm trains γ and β automatically.

6. Break it if you can: What LayerNorm does

Counterexample

7. LayerNorm vs BatchNorm: the axis difference

Concept

BatchNorm normalizes across the batch (per feature channel). LayerNorm normalizes across features (per token). For sequences, batch sizes vary and sequence lengths vary — BatchNorm statistics become unreliable.

For a token embedding x ∈ ℝ^d, LN's statistics are entirely self-contained — independent of what other tokens or other sequences are in the batch (Lesson 17 CE context).

8. Fill in: mean/var computed over for LayerNorm vs BatchNorm: the axis difference

Comparison

9. Guess the shape of the answer: LayerNorm from scratch

Estimation

10. LayerNorm from scratch

Worked example

Normalize x = [1, 3, 5, 7] manually, then verify against nn.LayerNorm.

import torch, torch.nn as nn

x = torch.tensor([[1.0, 3.0, 5.0, 7.0]])
mean = x.mean(dim=-1, keepdim=True)          # 4.0
var  = x.var(dim=-1, keepdim=True, unbiased=False)  # 5.0
xhat = (x - mean) / (var + 1e-5).sqrt()     # no gamma/beta yet

ln = nn.LayerNorm(4, elementwise_affine=False, eps=1e-5)
print('scratch:', xhat.tolist())
print('nn.LN:  ', ln(x).tolist())
print('match:  ', torch.allclose(xhat, ln(x), atol=1e-4))

11. What each one costs: LayerNorm from scratch

Trade off

12. Something is wrong here: normalizing over the wrong dimension

Anomaly

13. Trap: normalizing over the wrong dimension

Trap

14. Break it on purpose: normalizing over the wrong dimension

Break the constraint

15. Pre-norm vs Post-norm & residual connections

Section

Part 2 of 3

16. Pre-norm: the modern default

Concept

Two orderings exist for inserting LayerNorm relative to a sublayer (MHA or FFN):

Pre-norm keeps the residual stream unnormalized, so gradients flow back through the + x path without touching the norm. This is why GPT-2, GPT-3, LLaMA, and most modern transformers use pre-norm.

17. By analogy: Pre-norm: the modern default

Analogy

18. Residual connections: the gradient highway

Concept

The residual formula output = x + Sublayer(Norm(x)) means the network learns a correction (the residual) on top of the identity path.

\[ \frac{\partial \mathcal{L}}{\partial x} = \frac{\partial \mathcal{L}}{\partial(x + F(x))} \cdot \left(1 + \frac{\partial F}{\partial x}\right) \]

The 1 in the gradient ensures that, even if ∂F/∂x ≈ 0 in early training, a full gradient signal still flows to every earlier layer (Lesson 9 backprop context). Without residuals, deep networks suffer vanishing gradients.

19. Teach it back: Residual connections: the gradient highway

Explain it

20. What has to be given first: Residual gradient flow: with vs without

Missing information

21. Residual gradient flow: with vs without

Worked example

Train an 8-layer network with and without residual connections. Compare the gradient norm reaching layer 1 (deepest backprop path).

import torch, torch.nn as nn
torch.manual_seed(7)
depth, d = 8, 16

class NoRes(nn.Module):
    def __init__(self):
        super().__init__()
        self.layers = nn.ModuleList([nn.Linear(d, d) for _ in range(depth)])
        for l in self.layers: nn.init.normal_(l.weight, std=0.3)
    def forward(self, x):
        for l in self.layers: x = torch.tanh(l(x))
        return x

class WithRes(nn.Module):
    def __init__(self):
        super().__init__()
        self.layers = nn.ModuleList([nn.Linear(d, d) for _ in range(depth)])
    def forward(self, x):
        for l in self.layers: x = x + torch.tanh(l(x))
        return x

x = torch.randn(4, d)
for Cls, name in [(NoRes,'no-res'), (WithRes,'residual')]:
    model = Cls(); model(x).sum().backward()
    gnorms = [l.weight.grad.norm().item() for l in model.layers]
    print(f'{name}: L1={gnorms[0]:.3f}, L8={gnorms[-1]:.3f}')

22. Watch it run: Residual gradient flow: with vs without

Pattern

23. FFN sublayer — GELU and the 4× rule

Section

Part 3 of 3

24. The FFN sublayer

Concept

Each encoder position independently passes through a two-layer MLP: expand to d_ff = 4·d_model, apply GELU, project back. The expansion and contraction happen at every token in parallel.

\[ \text{FFN}(x) = W_2\, \text{GELU}(W_1 x + b_1) + b_2, \quad W_1 \in \mathbb{R}^{d_{ff} \times d}, \; W_2 \in \mathbb{R}^{d \times d_{ff}} \]

Why 4×? Empirical finding from the original transformer paper. The extra width gives the network capacity to store factual associations (Lesson 82 onwards). Why GELU over ReLU? Smoother gradient near zero; small negative values not completely zeroed out.

25. What rests on this: The FFN sublayer

Socratic

26. GELU vs ReLU: the smooth gate

Concept

GELU (Gaussian Error Linear Unit) weights the input by the probability it is positive: GELU(x) ≈ x · Φ(x) where Φ is the standard normal CDF. It smoothly suppresses small or negative inputs.

The smooth gradient at x < 0 avoids the dead neuron problem of ReLU (Lesson 44 activation context). BERT, GPT-2, and GPT-3 all use GELU.

27. Teach it back: GELU vs ReLU: the smooth gate

Explain it

28. Guess the shape of the answer: FFN sublayer in PyTorch

Estimation

29. FFN sublayer in PyTorch

Worked example

Build the FFN as an nn.Module. Use d_model=4, d_ff=16 (4×). Verify parameter count and output shape.

import torch, torch.nn as nn, torch.nn.functional as F
torch.manual_seed(0)

class FFN(nn.Module):
    def __init__(self, d_model, d_ff):
        super().__init__()
        self.w1 = nn.Linear(d_model, d_ff)
        self.w2 = nn.Linear(d_ff, d_model)
    def forward(self, x):
        return self.w2(F.gelu(self.w1(x)))

ffn = FFN(d_model=4, d_ff=16)
x   = torch.tensor([[1.0, 0.0, -1.0, 0.5]])
hid = F.gelu(ffn.w1(x))
out = ffn(x)
params = sum(p.numel() for p in ffn.parameters())
print(f'hidden shape: {hid.shape}')   # (1, 16)
print(f'output shape: {out.shape}')   # (1,  4)
print(f'output:       {out.squeeze().tolist()}')  # [-0.3974, 0.2011, 0.0934, -0.1582]
print(f'params:       {params}')      # 148

30. Which is which, by tensor shape

Discrimination

31. Something is wrong here: applying the FFN across the sequence instead of…

Anomaly

32. Trap: applying the FFN across the sequence instead of per-token

Trap

33. Which of these survive contact with Lesson 82: LayerNorm, FFN with GELU, and the…?

Two truths and a lie

34. Assembling the full EncoderBlock

Concept

The complete pre-norm encoder block chains two sublayers, each guarded by a LayerNorm and a residual connection:

1. Sublayer 1: x = x + MHA(LayerNorm(x)) — multi-head self-attention (Lesson 81)
2. Sublayer 2: x = x + FFN(LayerNorm(x)) — position-wise feed-forward network
3. Output x has the same shape as input — it goes into the next encoder block

The two LayerNorms (norm1, norm2) are distinct: norm1 guards the attention sublayer, norm2 guards the FFN sublayer. Each has its own learned γ and β.

35. By analogy: Assembling the full EncoderBlock

Analogy

36. Guess the shape of the answer: EncoderBlock as nn.Module

Estimation

37. EncoderBlock as nn.Module

Worked example

Implement the full pre-norm EncoderBlock. Use d_model=8, num_heads=2, d_ff=32. Verify shape and parameter count.

import torch, torch.nn as nn, torch.nn.functional as F
torch.manual_seed(42)

class MultiHeadSelfAttn(nn.Module):
    def __init__(self, d, h):
        super().__init__()
        self.h = h; self.dk = d // h
        self.Wq = nn.Linear(d, d, bias=False)
        self.Wk = nn.Linear(d, d, bias=False)
        self.Wv = nn.Linear(d, d, bias=False)
        self.Wo = nn.Linear(d, d, bias=False)
    def forward(self, x):
        B, T, D = x.shape
        def split(W): return W(x).view(B,T,self.h,self.dk).transpose(1,2)
        Q, K, V = split(self.Wq), split(self.Wk), split(self.Wv)
        attn = (Q @ K.transpose(-2,-1) / self.dk**0.5).softmax(-1)
        return self.Wo((attn @ V).transpose(1,2).reshape(B,T,D))

class EncoderBlock(nn.Module):
    def __init__(self, d, h, d_ff):
        super().__init__()
        self.norm1 = nn.LayerNorm(d)
        self.attn  = MultiHeadSelfAttn(d, h)
        self.norm2 = nn.LayerNorm(d)
        self.ffn   = nn.Sequential(
            nn.Linear(d, d_ff), nn.GELU(), nn.Linear(d_ff, d))
    def forward(self, x):
        x = x + self.attn(self.norm1(x))  # sublayer 1
        x = x + self.ffn(self.norm2(x))   # sublayer 2
        return x

block = EncoderBlock(d=8, h=2, d_ff=32)
x = torch.randn(2, 5, 8)
out = block(x)
params = sum(p.numel() for p in block.parameters())
print(f'input:  {x.shape}')    # (2, 5, 8)
print(f'output: {out.shape}')  # (2, 5, 8)
print(f'params: {params}')     # 840

38. Fill in: params (d=8) for EncoderBlock as nn.Module

Comparison

39. Without one step: The Encoder Block recipe

Constraint

40. The Encoder Block recipe

Pattern

1. LayerNorm: mean = x.mean(-1); var = x.var(-1); xhat = (x-mean)/(var+ε).sqrt() — normalize each token's feature vector independently
2. Pre-norm vs post-norm: pre-norm wraps the sublayer (x + Sub(Norm(x))); post-norm wraps the residual sum (Norm(x + Sub(x))). Use pre-norm for stable deep training
3. Residual connection: x = x + Sublayer(...) — the +x term keeps gradient magnitude healthy through depth (the 1 + in the chain rule)
4. FFN sublayer: Linear(d → 4d) → GELU → Linear(4d → d); applied position-wise (same weights every token, no cross-token mixing)
5. Full EncoderBlock: norm1 → MHA → residual then norm2 → FFN → residual; output shape = input shape (B, T, d_model)

41. Where does it stop working: The Encoder Block recipe

Edge cases

42. Rule out three: Check yourself — LayerNorm axis

Elimination

43. Check yourself — LayerNorm axis

Check

Compute before clicking.

44. Answer it before you see the options: Check yourself — GELU vs ReLU

Prediction

45. Check yourself — GELU vs ReLU

Check

What distinguishes them?

46. Rule out three: Check yourself — pre-norm formula

Elimination

47. Check yourself — pre-norm formula

Check

Identify the correct data flow.

48. Your turn: build EncoderBlock

Section

Project

49. Project: LayerNorm + FFN + EncoderBlock

Concept

Build all three components from scratch: LayerNorm (verify vs nn.LayerNorm), FFN sublayer (Linear+GELU+Linear), full EncoderBlock nn.Module with pre-norm residual wiring.

Build rules: type every line, implement LayerNorm manually before reaching for nn.LayerNorm, and verify shapes after each sublayer with print(x.shape).

50. Break it if you can: Project: LayerNorm + FFN + EncoderBlock

Counterexample

51. Milestone 1 — LayerNorm from scratch

Worked example

Your turn: implement LayerNorm manually on x = [1, 3, 5, 7]. Predict xhat before running.

Hint: mean = x.mean(dim=-1, keepdim=True); var = x.var(dim=-1, keepdim=True, unbiased=False); xhat = (x - mean) / (var + 1e-5).sqrt().

import torch, torch.nn as nn
x = torch.tensor([[1.0, 3.0, 5.0, 7.0]])
mean = x.mean(dim=-1, keepdim=True)
var  = x.var(dim=-1, keepdim=True, unbiased=False)
xhat = (x - mean) / (var + 1e-5).sqrt()
ln   = nn.LayerNorm(4, elementwise_affine=False, eps=1e-5)
print('scratch:', [round(v,4) for v in xhat.squeeze().tolist()])
print('nn.LN:  ', [round(v,4) for v in ln(x).squeeze().tolist()])
print('match:  ', torch.allclose(xhat, ln(x), atol=1e-4))

52. Milestone 2 — FFN sublayer

Worked example

Your turn: implement FFN(d_model=4, d_ff=16), pass x = [[1, 0, −1, 0.5]], and count the parameters. Predict the param count before running.

Hint: nn.Linear(d_model, d_ff) → F.gelu(...) → nn.Linear(d_ff, d_model). Params = d_model*d_ff + d_ff + d_ff*d_model + d_model.

import torch, torch.nn as nn, torch.nn.functional as F
torch.manual_seed(0)
class FFN(nn.Module):
    def __init__(self, d_model, d_ff):
        super().__init__()
        self.w1 = nn.Linear(d_model, d_ff)
        self.w2 = nn.Linear(d_ff, d_model)
    def forward(self, x): return self.w2(F.gelu(self.w1(x)))

ffn = FFN(4, 16)
x   = torch.tensor([[1.0, 0.0, -1.0, 0.5]])
out = ffn(x)
params = sum(p.numel() for p in ffn.parameters())
print('output:', [round(v,4) for v in out.squeeze().tolist()])
print('params:', params)

53. What each one costs: Milestone 2 — FFN sublayer

Trade off

54. Milestone 3 — full EncoderBlock

Worked example

Your turn: assemble EncoderBlock(d=8, h=2, d_ff=32) with pre-norm residuals. Verify that output.shape == input.shape and params == 840.

Hint: forward: x = x + attn(norm1(x)); x = x + ffn(norm2(x)). Use nn.Sequential(nn.Linear(d,d_ff), nn.GELU(), nn.Linear(d_ff,d)) for the FFN.

import torch, torch.nn as nn
torch.manual_seed(42)
# (reuse MultiHeadSelfAttn from Milestone 2 shell)
class EncoderBlock(nn.Module):
    def __init__(self, d, h, d_ff):
        super().__init__()
        self.norm1 = nn.LayerNorm(d)
        self.attn  = MultiHeadSelfAttn(d, h)   # from prev. milestone
        self.norm2 = nn.LayerNorm(d)
        self.ffn   = nn.Sequential(
            nn.Linear(d, d_ff), nn.GELU(), nn.Linear(d_ff, d))
    def forward(self, x):
        x = x + self.attn(self.norm1(x))
        x = x + self.ffn(self.norm2(x))
        return x

block = EncoderBlock(d=8, h=2, d_ff=32)
x = torch.randn(2, 5, 8)
out = block(x)
params = sum(p.numel() for p in block.parameters())
print(f'input:  {x.shape}')    # torch.Size([2, 5, 8])
print(f'output: {out.shape}')  # torch.Size([2, 5, 8])
print(f'params: {params}')     # 840

55. Which is which, by note

Discrimination

56. The full program

Concept

import torch, torch.nn as nn, torch.nn.functional as F
torch.manual_seed(42)

# LayerNorm from scratch
def layernorm_scratch(x, eps=1e-5):
    mu  = x.mean(-1, keepdim=True)
    var = x.var(-1, keepdim=True, unbiased=False)
    return (x - mu) / (var + eps).sqrt()

# Multi-head self-attention
class MHA(nn.Module):
    def __init__(self, d, h):
        super().__init__()
        self.h = h; self.dk = d // h
        self.Wq = nn.Linear(d,d,bias=False); self.Wk = nn.Linear(d,d,bias=False)
        self.Wv = nn.Linear(d,d,bias=False); self.Wo = nn.Linear(d,d,bias=False)
    def forward(self, x):
        B,T,D = x.shape
        def sp(W): return W(x).view(B,T,self.h,self.dk).transpose(1,2)
        Q,K,V = sp(self.Wq), sp(self.Wk), sp(self.Wv)
        return self.Wo(((Q@K.transpose(-2,-1)/self.dk**0.5).softmax(-1)@V)
                       .transpose(1,2).reshape(B,T,D))

# EncoderBlock (pre-norm)
class EncoderBlock(nn.Module):
    def __init__(self, d, h, d_ff):
        super().__init__()
        self.norm1=nn.LayerNorm(d); self.attn=MHA(d,h)
        self.norm2=nn.LayerNorm(d)
        self.ffn=nn.Sequential(nn.Linear(d,d_ff),nn.GELU(),nn.Linear(d_ff,d))
    def forward(self, x):
        x = x + self.attn(self.norm1(x))
        x = x + self.ffn(self.norm2(x))
        return x

# verify
xv = torch.tensor([[1.0,3.0,5.0,7.0]])
print('LN match:', torch.allclose(layernorm_scratch(xv), nn.LayerNorm(4,elementwise_affine=False)(xv),atol=1e-4))
enc = EncoderBlock(8,2,32)
xe  = torch.randn(2,5,8)
print('shape:', enc(xe).shape)    # torch.Size([2, 5, 8])
print('params:', sum(p.numel() for p in enc.parameters()))  # 840

Stack N of these EncoderBlock modules — that's the BERT/GPT encoder. Every block: normalize → attend → normalize → feed-forward, with clean residual paths throughout.

57. Fill in: what it confirms for The full program

Comparison

58. Show it off

Concept

Out loud, slides closed: explain (1) why LayerNorm normalizes over the feature dimension rather than the batch dimension; (2) the pre-norm formula and why it stabilizes deep training; (3) what the FFN's 4× expansion achieves and why GELU is preferred over ReLU.

Stretch (homework from L82): implement LayerNorm as a proper nn.Module with trainable γ and β; run the pre-norm vs post-norm training experiment on a 6-layer model and log the loss curves; experiment with d_ff = 2·d_model vs 4·d_model and observe quality difference. Next: Lesson 83 — positional encodings and full encoder stack.

59. Connect it up: Lesson 82: LayerNorm, FFN with GELU, and the Full Encoder Block

Connect it up

60. What you can do now

Recap

• Implement LayerNorm from scratch (mean/var over last dim) and verify against nn.LayerNorm
• Explain LayerNorm vs BatchNorm: LN normalizes per-token (feature axis), BN per-feature (batch axis)
• State the pre-norm formula (x = x + Sub(Norm(x))) and why it stabilizes deep transformers
• Build the FFN sublayer (Linear → GELU → Linear, d_ff = 4d), position-wise
• Assemble the full EncoderBlock nn.Module with two pre-norm sublayers and residual connections