Layer Normalization › Layer Normalization with Different Normalized Shapes
Language AI HandbookAll chapters Layer Normalization › Layer Normalization with Different Normalized ShapesOne section · Sections in this chapter11 / 16 ↓
Reading tools Definitions & inline links
Choose your expertise level to adjust how many terms are explained. Beginners see more tooltips, experts see fewer to maintain reading flow. Hover over underlined terms for instant definitions.
Article links
Make inline references clickable
PyTorch's nn.LayerNorm accepts a normalized_shape parameter that controls which dimensions are normalized. For transformers, we typically normalize over the feature dimension only:
# Different normalized shapes
x = torch.randn(2, 4, 8) # (batch, seq_len, features)
# Normalize over features only (most common for transformers)
ln_features = nn.LayerNorm(8)
# Normalize over sequence and features
ln_seq_features = nn.LayerNorm([4, 8])
# Normalize over entire sample (batch, sequence, features)
ln_all = nn.LayerNorm([4, 8])
out_features = ln_features(x)
out_seq_features = ln_seq_features(x)LayerNorm with normalized_shape=(8,) - normalize over features: Output shape: torch.Size([2, 4, 8]) Each token normalized independently Token 0,0 mean: 0.000000 Token 0,1 mean: -0.000000 LayerNorm with normalized_shape=(4, 8) - normalize over seq+features: Output shape: torch.Size([2, 4, 8]) Each sample normalized as a whole Sample 0 mean: 0.000000
Notice the difference: with normalized_shape=(8,), each individual token has zero mean (Token 0,0 and Token 0,1 both have mean approximately 0). With normalized_shape=(4, 8), the entire sample is normalized together, so individual tokens may have non-zero means but the sample as a whole has zero mean.
Normalizing over features only is the standard choice for transformers because it treats each token independently, matching the autoregressive nature of language models and allowing the model to process variable-length sequences.
The choice of normalized_shape also has practical implications during inference. When you normalize over the feature dimension only, you can process a single token at a time in autoregressive generation, and layer normalization works exactly as during training. If you were to normalize over the sequence dimension, you would need to know the full sequence before normalizing any token, which breaks the causal autoregressive property required for efficient generation.