Layer Normalization & RMSNorm

Table of Contents

1. Learning Objectives
2. Notation
3. Core Intuition
4. Batch Normalization Revisited
5. Layer Normalization
6. RMSNorm
7. Pre-Norm vs Post-Norm
8. Gradient Analysis
9. Common Pitfalls
10. Summary
11. Exercises

Learning Objectives

1. Contrast batch normalization (statistics over batch) with layer normalization (statistics over features).
2. Derive the LayerNorm forward pass and its Jacobian.
3. Derive RMSNorm as a simplified LayerNorm without mean centering.
4. Explain why transformers prefer LayerNorm over BatchNorm.
5. Analyze Pre-Norm vs Post-Norm placement in residual blocks.

Notation

• $\mathbf{h} \in \mathbb{R}^d$ — hidden representation (single token/sample)
• $\mu = \frac{1}{d}\sum_{i=1}^d h_i$ — mean over features
• $\sigma^2 = \frac{1}{d}\sum_{i=1}^d(h_i - \mu)^2$ — variance over features
• $\gamma, \beta \in \mathbb{R}^d$ — learnable scale and shift parameters
• $\text{RMS}(\mathbf{h}) = \sqrt{\frac{1}{d}\sum_i h_i^2}$ — root mean square

Core Intuition

Normalization stabilizes training by preventing activations from growing or shrinking across layers. Batch normalization normalizes over the batch dimension (requiring batch statistics), which is problematic for variable-length sequences and small batches. Layer normalization normalizes over the feature dimension of a single sample, making it independent of batch size and applicable to recurrent/transformer architectures.

Batch Normalization Revisited

For a mini-batch of $B$ samples, BatchNorm computes statistics per feature:

Limitations: Depends on batch statistics → requires running averages at inference; fails for batch size 1; not applicable to autoregressive models where future tokens shouldn't influence current normalization.

Layer Normalization

LayerNorm computes statistics over the feature dimension of a single sample:

Properties:

• Independent of batch size and other samples.
• Same computation at training and inference (no running statistics).
• Learnable affine parameters $\gamma, \beta$ restore representational power.

Jacobian. Let $\hat{\mathbf{h}} = (\mathbf{h} - \mu\mathbf{1})/\sqrt{\sigma^2 + \epsilon}$. Then:

RMSNorm

RMSNorm (Zhang & Sennrich, 2019) removes the mean centering step:

Motivation: Empirically, the re-centering in LayerNorm contributes little to performance but adds computation. RMSNorm is ~10–15% faster.

Used in: LLaMA, Gemma, and most modern LLMs.

Pre-Norm vs Post-Norm

In a transformer block with residual connection:

Post-Norm (original transformer): $\mathbf{h} + \text{Norm}(\text{Attn}(\mathbf{h}))$ — normalizes after the sublayer.

Pre-Norm (GPT-2 onwards): $\mathbf{h} + \text{Attn}(\text{Norm}(\mathbf{h}))$ — normalizes before the sublayer.

Analysis. Pre-Norm produces better-behaved gradients:

• The gradient through the skip connection is unmodified ($\partial h^{(L)}/\partial h^{(l)}$ includes a direct path of magnitude 1).
• Post-Norm can cause gradient explosion in very deep transformers without careful warmup.
• Pre-Norm enables stable training without learning rate warmup.

Gradient Analysis

For Pre-Norm with $L$ blocks:

where $\mathbf{J}_k$ is the Jacobian of block $k$. The leading term $\partial L/\partial \mathbf{h}^{(L)}$ provides a direct gradient path, preventing vanishing.

Common Pitfalls

Pitfall 1. Using BatchNorm in transformers for NLP. Variable sequence lengths and autoregressive masking make batch statistics unreliable.

Pitfall 2. Forgetting $\epsilon$ in the denominator. Without it, zero-variance features cause division by zero.

Pitfall 3. Applying LayerNorm to the wrong dimension. For a tensor of shape $(B, T, d)$, normalize over the last dimension $d$ (features), not $T$ (sequence) or $B$ (batch).

Summary

• BatchNorm normalizes over the batch → dependent on other samples, problematic for inference.
• LayerNorm normalizes over features of a single sample → independent, stable, used in transformers.
• RMSNorm drops mean centering → faster, empirically equivalent, used in modern LLMs.
• Pre-Norm placement enables stable deep training by preserving gradient flow through skip connections.

Exercises

Exercise 1. Derive the Jacobian of LayerNorm (equation 4).

Exercise 2. Show that LayerNorm is invariant to uniform scaling and shifting of its input: $\text{LN}(a\mathbf{h} + b\mathbf{1}) = \text{LN}(\mathbf{h})$ (ignoring $\gamma, \beta$).

Exercise 3. Prove that RMSNorm projects inputs onto the unit sphere of radius $\sqrt{d}$, scaled by $\gamma$.

Exercise 4. For a 2-layer Post-Norm transformer, derive the gradient $\partial L/\partial \mathbf{h}^{(0)}$ and identify conditions for instability.

Exercise 5. Compare the computational cost (FLOPs) of LayerNorm vs RMSNorm for a vector of dimension $d$.