Layer Normalization

Unlike Batch Normalization, which computes statistics across the batch dimension, Layer Normalization computes the mean and variance across the features of a single input vector.

In deep networks, covariate shift within hidden layers can slow down optimization. Layer Normalization (LayerNorm) mitigates this by computing the mean and variance across the feature dimension for each training sample individually. For a layer input vector \\(x \\in \\mathbb{R}^D\\), the mean \\(\\mu\\) is calculated as \\(\\mu = \\frac{1}{D} \\sum_{i=1}^D x_i\\), and the variance \\(\\sigma^2\\) is calculated as \\(\\sigma^2 = \\frac{1}{D} \\sum_{i=1}^D (x_i - \\mu)^2\\).

Once the mean and variance are computed, the inputs are normalized to have zero mean and unit variance. To avoid division by zero, a small numerical stability term \\(\\epsilon\\) is added to the variance in the denominator. The normalized value is computed as \\(\\hat{x}_i = \\frac{x_i - \\mu}{\\sqrt{\\sigma^2 + \\epsilon}}\\). To preserve the representational power of the network, LayerNorm introduces learnable parameters \\(\\gamma \\in \\mathbb{R}^D\\) and \\(\\beta \\in \\mathbb{R}^D\\), which scale and shift the normalized value. The final output is \\(y_i = \\gamma_i \\hat{x}_i + \\beta_i\\).

A major limitation of Batch Normalization is its dependence on batch size and batch statistics. This makes it difficult to apply to recurrent neural networks (RNNs) and Transformers where sequence lengths vary, and mini-batch sizes can be small. LayerNorm resolves this because its statistics are computed solely within a single sample's features, meaning the normalization behavior remains identical during training and inference.

In sequence models, the feature dimensions correspond to the embedding dimensions at each time step. LayerNorm normalizes each token's embedding vector independently, allowing the network to handle variable-length sequences without maintaining running averages of mean and variance across the entire training dataset. This batch independence stabilizes the training of massive Transformer architectures like GPT and BERT.

Understanding the geometric and domain differences between normalization methods is critical for selecting the right architecture.

The difference between Batch Normalization and Layer Normalization lies in the axes over which mean and variance are computed. In a tensor of shape \\((N, C, H, W)\\), Batch Normalization computes statistics across the batch dimension \\(N\\) and spatial dimensions \\(H, W\\) for each channel \\(C\\) independently. Thus, it normalizes activations across different samples.

Conversely, Layer Normalization computes statistics across the channel \\(C\\) and spatial dimensions \\(H, W\\) for each sample \\(N\\) individually. Geometrically, Batch Normalization performs vertical normalization across the batch, while Layer Normalization performs horizontal normalization across the features. This distinction determines their suitability for different network architectures and data modalities.

Batch Normalization introduces noise into the network during training because the normalization of a sample depends on the other samples in the mini-batch. This noise acts as a regularizer, often improving generalization in convolutional neural networks (CNNs) used for computer vision. Layer Normalization does not introduce this batch-dependent noise, making it a pure normalization technique without regularizing side-effects.

In terms of domain suitability, CNNs benefit more from Batch Normalization because spatial feature maps share statistical properties across channels. In contrast, natural language processing models and Transformers, which deal with highly variable embeddings and sequence lengths, rely on Layer Normalization to ensure stable gradient flow and sequence-level independence.

Let's look at how PyTorch provides built-in LayerNorm support, and how to verify its outputs manually.

PyTorch provides the nn.LayerNorm module, which takes the normalized shape as an argument. For instance, in NLP models, we often normalize over the last dimension (e.g., the embedding dimension). PyTorch handles the registration of learnable parameters weight (representing \\(\\gamma\\)) and bias (representing \\(\\beta\\)) and performs the forward pass efficiently.

The parameters of the LayerNorm layer, specifically the weight and bias vectors, are initialized to 1s and 0s respectively. During training, their gradients are calculated, allowing them to shift and scale the normalized features to restore the representational capacity of the layer.

To verify the mathematical mechanics of nn.LayerNorm, we can implement the operations manually using basic PyTorch tensor operations. By calculating the mean and variance along the last dimension (dim=-1), normalizing the tensor, and applying the weight and bias parameters, we can reproduce the exact outputs of PyTorch's built-in layer.

This manual implementation illustrates how the statistics are constrained within each individual sample's feature vector. It also highlights the role of the division by the standard deviation and the broadcast addition of the bias vector. Calculating variance with unbiased=False is necessary to match the population variance formula used by PyTorch's C++ backend.