Layer Normalization

Unlike Batch Normalization, which normalizes across the batch dimension, Layer Normalization computes the mean and variance using all the hidden units in a single layer for an individual training case. This makes the computation independent of batch size, which is crucial for NLP tasks where batch sizes are often small due to memory constraints. The operation normalizes the summed inputs to a neuron across the feature dimension, re-centering and re-scaling the distribution.

A primary advantage is its independence from the batch dimension. In recurrent or Transformer models, sequence lengths vary, and distributed training often uses very small micro-batches. Layer Normalization performs identically during training and inference because its statistics are computed per time step, not per batch. This eliminates the train-test discrepancy and the reliance on running averages that plague Batch Normalization in sequence models.

By normalizing the pre-activation inputs, Layer Normalization drastically reduces internal covariate shift. It prevents the distribution of inputs to subsequent layers from drifting during training, which keeps the network in a regime where gradients are well-behaved. This allows for significantly higher learning rates, accelerating convergence and reducing sensitivity to weight initialization schemes in deep architectures like the Transformer.

After normalizing the input to have zero mean and unit variance, Layer Normalization introduces two learnable parameters: gamma (γ) and beta (β). These parameters scale and shift the normalized output, restoring the model's representational capacity. Without this affine transformation, the network's expressiveness would be severely limited, as the normalization would force every layer's output to have the same fixed distribution.

In the original Transformer architecture, Layer Normalization is applied in a post-layer normalization setup, placed after the residual addition. Modern architectures like GPT-2 and Llama often use pre-layer normalization, applying the normalization before the attention and feed-forward sub-layers. Pre-norm empirically provides a cleaner gradient highway, enabling the training of much deeper networks without the need for complex learning rate warm-up schedules.

An efficient alternative called RMSNorm has gained popularity in models like Llama. It hypothesizes that the re-centering operation (subtracting the mean) is less critical than the scaling operation. RMSNorm normalizes using only the root mean square statistic, discarding the mean calculation. This reduces computational overhead by approximately 7-15% with negligible performance degradation, making it a preferred choice for large-scale production models.

A pathway that adds the input of a sub-layer directly to its output before layer normalization is applied. This identity mapping prevents vanishing gradients and enables the training of very deep networks.

• Formula: Output = LayerNorm(x + Sublayer(x))
• Pre-LN vs. Post-LN: Modern Transformers typically apply normalization before the sublayer (Pre-LN) for training stability
• Gradient Highway: Allows error signals to propagate backward without degradation

The precursor to Layer Normalization that normalizes activations across the batch dimension rather than the feature dimension. It computes statistics over all samples in a mini-batch for each feature independently.

• Sequence Limitation: Performs poorly on variable-length sequences and small batches
• Recurrent Incompatibility: Cannot be applied effectively in RNNs due to unrolling complexities
• Training vs. Inference: Requires running averages of batch statistics for inference, adding complexity

The fundamental building block that integrates layer normalization with other core mechanisms. Each block typically consists of a multi-head self-attention sub-layer and a position-wise feed-forward network, each wrapped with residual connections and layer normalization.

• Standard Order: Attention → Add & Norm → FFN → Add & Norm
• Normalization Role: Stabilizes the distribution of inputs to each sub-layer
• Stacking: Deep models stack many identical blocks, relying on normalization to maintain signal quality

A position-wise, fully connected network that applies a non-linear transformation to each token's representation independently. Layer normalization ensures the FFN receives well-conditioned inputs regardless of the attention output's scale.

• Typical Structure: Two linear transformations with a GELU or ReLU activation
• Expansion Ratio: Often expands the hidden dimension by 4x before projecting back
• Normalization Dependency: Without LN, the FFN's gradients can explode due to large attention outputs

A simplification of Layer Normalization that uses only the root mean square statistic for re-scaling, omitting the mean-centering step. This reduces computational overhead while maintaining comparable training stability.

• Formula: Output = x / RMS(x) * γ
• Adoption: Used in LLaMA, Mistral, and other modern large language models
• Efficiency Gain: Removes the need to compute and subtract the mean, saving approximately 5-10% of normalization compute

The phenomenon where the distribution of network activations changes during training as parameters are updated. Layer Normalization directly addresses this by re-normalizing the statistics of each training case's features.

• Original Motivation: Reducing this shift was the primary justification for batch normalization
• LayerNorm's Advantage: Normalizes each sample independently, decoupling the statistics from the batch
• Modern View: Normalization may primarily smooth the loss landscape rather than reduce covariate shift