Layer Normalization — Module 10.1

How This Lesson Fits the Module & Engineering Practice

Alongside residual connections, transformers stabilize training with layer normalization. Unlike BatchNorm (Vol 06), which normalizes across the batch for each feature, LayerNorm normalizes across features for each token—ideal for variable-length sequences and small or uneven batches.

Every encoder/decoder block you build next will combine attention or FFN with residual + LayerNorm. Getting the axes right prevents silent train/eval bugs that BatchNorm often causes on sequences.

Learning Objectives

By the end of this lesson, students should be able to:

Define LayerNorm and state which axes it reduces over.
Compare LayerNorm vs BatchNorm for sequence models.
Explain pre-norm vs post-norm block layouts at a high level.
Use nn.LayerNorm correctly in PyTorch.
Describe learnable \(\gamma\) and \(\beta\) affine parameters.
Choose LayerNorm over BatchNorm for typical NLP transformers.

Definition

Layer normalization re-centers and re-scales the features of each individual example (each token vector) so that, across the feature dimension, the mean is ~0 and variance is ~1, then applies learned gain \(\gamma\) and bias \(\beta\).

The LayerNorm Computation

For a token vector \(x \in \mathbb{R}^{d}\):

\[\mu = \frac{1}{d}\sum_{j=1}^{d} x_j,\quad \sigma^2 = \frac{1}{d}\sum_{j=1}^{d}(x_j-\mu)^2\]

\[\mathrm{LN}(x) = \gamma \odot \frac{x-\mu}{\sqrt{\sigma^2+\varepsilon}} + \beta\]

In batched form \((B, T, d)\), statistics are computed over the last dimension \(d\) independently for every \((b, t)\) position.

LayerNorm vs BatchNorm for Sequences

Property	LayerNorm	BatchNorm
Stats over	Features of one token	Batch (and often spatial) per channel
Depends on batch size?	No	Yes—noisy for tiny batches
Variable sequence length	Natural per-token	Awkward with padding / masks
Train vs eval	Same formula	Running averages at eval
Typical home	Transformers, RNNs	CNNs (see Vol 06)

Prefer LayerNorm

Token sequences, transformers.
Small or variable batches.
Need identical train/infer norm math.

Prefer BatchNorm

Large-batch vision CNNs.
Stable channel statistics.
Classic ResNet-style vision stacks.

Where It Sits in a Block

Post-norm (original)

\(x \leftarrow \mathrm{LN}(x + F(x))\)

Pre-norm (common now)

\(x \leftarrow x + F(\mathrm{LN}(x))\)

With attention

LN paired with MHA residual.

With FFN

LN paired with position-wise MLP.

Both patterns appear in production; pre-norm often trains more stably at large depth. Pair this lecture with Residual Connection and Feed-Forward Network.

PyTorch: LayerNorm on Token Streams

import torch
from torch import nn

d_model = 64
ln = nn.LayerNorm(d_model)  # normalize last dim

x = torch.randn(2, 10, d_model)  # (batch, seq, features)
y = ln(x)
print(y.shape)
print(y[0, 0].mean().abs().item() < 1e-5)  # ~True after LN (before affine nuance)

# Minimal pre-norm residual FFN slice
ff = nn.Sequential(
    nn.Linear(d_model, 4 * d_model),
    nn.GELU(),
    nn.Linear(4 * d_model, d_model),
)
x = x + ff(ln(x))
print(x.shape)

Why transformers use LN

Stable per-token feature scale.
Batch-size independent.
Plays well with padding masks.

Watch-outs

Wrong normalized_shape breaks axes.
Not a substitute for good init/LR.
Pre- vs post-norm changes dynamics.

Common Mistake

Dropping BatchNorm into a transformer “because norms help.” With padded sequences and small NLP batches, BatchNorm statistics are often misleading; LayerNorm is the default for token models.

Misconception

“LayerNorm normalizes across the sequence length.” Standard transformer LayerNorm normalizes across the feature dimension of each position, not across time. (There are variants, but nn.LayerNorm(d_model) is per-token features.)

Related module pages: Multi-Head Attention, Self Attention, Positional Encoding, Encoder, Decoder.

Knowledge Check

Short Answer: Over which dimension does standard transformer LayerNorm compute mean/variance? Answer: The feature dimension \(d_model\) for each token.
True/False: LayerNorm’s statistics depend on batch size the way BatchNorm’s do. Answer: False.
Multiple Choice: For NLP transformers, the usual norm is: (a) LayerNorm, (b) BatchNorm, (c) only max-pooling. Answer: (a).
Short Answer: Why is BatchNorm awkward for padded sequences? Answer: Batch/channel stats mix real tokens with padding and depend on batch composition.
True/False: LayerNorm includes learnable \(\gamma\) and \(\beta\). Answer: True.
Multiple Choice: Pre-norm applies LN: (a) before the sublayer \(F\), (b) only on the batch axis, (c) instead of residuals. Answer: (a).
Short Answer: Name the Vol 06 lecture that covers BatchNorm. Answer: Batch Normalization.
True/False: nn.LayerNorm(d_model) expects to normalize the last dimension of size \(d_model\). Answer: True.
Multiple Choice: Residuals and LayerNorm together help: (a) deep transformer training stability, (b) remove attention, (c) replace PE. Answer: (a).
Short Answer: What lecture completes the typical block after LN in this track? Answer: Feed-Forward Network.

Key Takeaways

LayerNorm normalizes each token’s features; BatchNorm normalizes across the batch.
Transformers prefer LayerNorm for sequences and variable batches.
Used with residuals in pre-norm or post-norm layouts.
Contrast with Vol 06 Batch Normalization.
Next: Feed-Forward Network.

Trainer’s Guide

Hands-on idea: Print mean/variance of a token vector before and after nn.LayerNorm and verify feature-axis normalization.

Discussion prompt: Why did vision ResNets standardize on BatchNorm while transformers standardized on LayerNorm?

Recap: LayerNorm stabilizes per-token features without batch statistics—the right default for attention stacks. Continue with Feed-Forward Network.