LayerNorm

Layer Normalization normalizes each data point across its feature dimension, then rescales with learnable parameters. It is the default normalizer inside Transformers because it is independent of batch size and works the same at training and inference.

What is Layer Normalization?

Info

Layer Normalization (LayerNorm) is a technique used to stabilize and improve the training of deep neural networks by normalizing the outputs of each layer.

LayerNorm ⇒ Normalize by rows (data items)

1. The Problem Addressed by LayerNorm

• When training deep networks, the distribution of layer inputs can change as weights are updated, which may lead to a problem known as internal covariate shift. This can hinder the training process and slow down convergence.

2. How Layer Normalization Works

• Layer normalization addresses this issue by normalizing the inputs of each layer across all features (rather than across the batch like Batch Normalization).
• For each input $x$ in a layer, LayerNorm computes the mean $\mu$ and variance ${\sigma }^2$ over the $H$ features:

$$\mu =\frac{1}{H}\sum _{i=1}^H{x}_i{\sigma }^2=\frac{1}{H}\sum _{i=1}^H(x_i-\mu {)}^2$$

• It then normalizes the input and applies learnable scale $\gamma$ and shift $\beta$:

$$\hat{x}=\frac{x-\mu }{\sqrt{{\sigma }^2+ϵ}}y=\gamma \odot \hat{x}+\beta$$

where $ϵ$ is a small constant added for numerical stability.

3. Why LayerNorm is Useful

• Consistent Performance: By normalizing across features for each data point, LayerNorm stabilizes the learning process and leads to more consistent performance across different inputs.
• Better Gradient Flow: LayerNorm helps prevent the vanishing/exploding gradient problem, ensuring that gradients remain well-behaved during backpropagation, especially in deep networks.

Caution

Key Differences from Batch Normalization: Unlike Batch Normalization, which normalizes across the batch dimension, LayerNorm normalizes across the feature dimension. This makes LayerNorm especially useful for tasks with varying batch sizes or for recurrent neural networks (RNNs) where the sequence length can differ.

LayerNorm vs BatchNorm

Aspect	LayerNorm	BatchNorm
Normalize over	feature dim (per token)	batch dim (per feature)
Batch size dependence	none	strong
Train vs eval	identical	needs running statistics
Best for	Transformers, RNNs	CNNs with large batches

Because BatchNorm mixes statistics across the batch, it behaves differently at inference and struggles with small or variable batches; LayerNorm sidesteps both issues.

PyTorch Implementation

import torch
import torch.nn as nn
 
class LayerNorm(nn.Module):
    def __init__(self, dim, eps=1e-5):
        super().__init__()
        self.gamma = nn.Parameter(torch.ones(dim))
        self.beta = nn.Parameter(torch.zeros(dim))
        self.eps = eps
 
    def forward(self, x):
        # normalize over the last (feature) dim
        mu = x.mean(dim=-1, keepdim=True)
        var = x.var(dim=-1, unbiased=False, keepdim=True)
        x_hat = (x - mu) / torch.sqrt(var + self.eps)
        return self.gamma * x_hat + self.beta

Pre-Norm vs Post-Norm

Where the norm sits in a residual block matters for training stability.

• Post-norm (original Transformer): $x+\text{Sublayer}(x)$ then normalize. Higher final quality but unstable for deep stacks, often needing learning rate warmup.
• Pre-norm (GPT-2 onward): normalize first, $x+\text{Sublayer}(\text{LN}(x))$. The residual path stays clean, so gradients flow freely and very deep models train without warmup.

Pitfall

Do not put LayerNorm on the residual (skip) path. Normalizing the identity branch destroys the clean gradient highway that makes deep residual networks trainable.

RMSNorm

RMSNorm drops the mean subtraction and the $\beta$ shift, normalizing only by the root mean square:

$$y=\frac{x}{\sqrt{\frac{1}{H}{\sum }_{i=1}^H{x}_i^2+ϵ}}\odot \gamma$$

It is cheaper (no mean, no bias) yet matches LayerNorm quality, so modern LLMs such as LLaMA use it.

class RMSNorm(nn.Module):
    def __init__(self, dim, eps=1e-6):
        super().__init__()
        self.gamma = nn.Parameter(torch.ones(dim))
        self.eps = eps
 
    def forward(self, x):
        rms = x.pow(2).mean(dim=-1, keepdim=True).add(self.eps).rsqrt()
        return x * rms * self.gamma

Applications

• Layer Normalization is commonly used in transformer architectures, RNNs, and other deep learning models to improve training stability and performance.

Info

Benefits of Using Layer Normalization:

• Faster convergence.
• Improved training stability.
• Robust performance across varying input distributions.