Chapter 2: Layer Normalization in LLMs

Layer normalization plays a critical role in stabilizing training and improving convergence in transformer-based large language models (LLMs). This chapter introduces key normalization techniques and their practical implications in modern LLM architectures.

2.1 What Is Layer Normalization?

Layer Normalization (LayerNorm) standardizes the inputs across the hidden dimension for each sample independently. It computes the mean and variance over all hidden units and applies a learned scale (gamma) and shift (beta):

Where:

• gamma is a learnable scaling factor
• beta is a learnable shifting factor

2.2 What Is RMSNorm?

RMSNorm (Root Mean Square Normalization) is a simplified version of LayerNorm that removes the mean-centering step:

RMSNorm does not subtract the mean. It only rescales the vector using the RMS.This makes RMSNorm more efficient to compute. RMSNorm can yield comparable or even slightly better results in large-scale models.

2.3 What Is DeepNorm?

DeepNorm is a normalization variant introduced to stabilize deep transformer training. It modifies the residual path before applying normalization:

Where alpha > 1 upscales the residual connection before applying LayerNorm. Additionally, model weights are initialized using a downscaled factor beta <1 to further enhance stability.

DeepNorm helps prevent gradient explosion or vanishing in very deep transformer stacks.

def deepnorm(x): return LayerNorm(x + alpha * f(x))

Initialization method used in DeepNorm

def deepnorm_init(w): if w in [‘ffn’, ‘v_proj’, ‘out_proj’]: nn.init.xaviernormal(w, gain=α) elif w in [‘q_proj’, ‘k_proj’]: nn.init.xaviernormal(w, gain=β)

2.4 Where Does LayerNorm Appear in Transformers?

In transformer architectures, Layer Normalization can be applied in different positions relative to the residual connections. Each placement has its trade-offs in terms of training stability and model performance. In transformer blocks, LayerNorm can be placed in three main positions:

1. Post-LN:

• Position: LayerNorm is applied after the residual connection (x + f(x) → LayerNorm).
• Advantages: Straightforward to implement and conceptually simple.
• Drawbacks: Can lead to gradient instability in very deep models, especially as the depth increases.

2. Pre-LN:

• Position: LayerNorm is applied before the residual connection (LayerNorm(x) → f(x) → x + f(x)).
• Advantages: Improves training stability in deep transformers by normalizing inputs to each sublayer, which helps mitigate exploding or vanishing gradients.
• Drawbacks: May result in slightly reduced empirical performance compared to Post-LN in some benchmarks.

3. Sandwich-LN:

• Position: LayerNorm is applied both before and after the residual connection, effectively wrapping the residual block.
• Advantages: Offers enhanced stability by reducing gradient explosion; used in models like CogView
• Drawbacks: May introduce over-regularization, potentially slowing down convergence or degrading performance if not carefully tuned.

2.5 Which Layer Normalization Variants Do Popular LLMs Use?

Note: BLOOM adds LayerNorm after the embedding layer, improving training stability but sometimes at the cost of performance.