LayerNorm vs BatchNorm vs RMSNorm

Comparing Normalization Dimensions

The core difference between normalization techniques comes down to which axis you average over.

Given an activation tensor with dimensions Batch (N), Sequence/Height (L), Channels/Features (C):

┌──────────────────────────┬──────────────────────────┬──────────────────────────┐
│  1. BATCH NORMALIZATION  │  2. LAYER NORMALIZATION  │  3. RMSNORM              │
├──────────────────────────┼──────────────────────────┼──────────────────────────┤
│ Averages across Batch N  │ Averages across Features │ Scales by Root Mean      │
│ for each feature channel.│ C for each sample token. │ Square of features       │
│ Best for Computer Vision │ Best for Transformers    │ without computing mean.  │
│ (ResNet, ConvNeXt).      │ (BERT, GPT, ViT).        │ Faster LLM training.     │
└──────────────────────────┴──────────────────────────┴──────────────────────────┘

Visualizing the Axes

Imagine a 3D block of numbers where:

• N is the number of images or sentences in your batch.
• L is the sequence position or image spatial grid.
• C is the feature embedding dimension.

  BATCH NORM:           Normalizes ALONG the Batch dimension N.
                        Requires large batch sizes to get accurate mean and variance.

  LAYER NORM:           Normalizes ACROSS the Feature dimension C for each single sample.
                        Works identically whether batch size is 1 or 1000.

Why Transformers Use Layer Normalization

1. Variable Sequence Lengths: Text sentences in a batch have different word lengths. Batch Normalization struggles when padding tokens are mixed into batch statistics. Layer Normalization treats every token independently.
2. Batch Size Independence: Large language models often train or serve inference with small batch sizes due to memory limits. Layer Normalization works perfectly even when batch size is 1.

RMSNorm (Root Mean Square Normalization)

Modern Large Language Models like LLaMA 2 and LLaMA 3 replace Layer Normalization with RMSNorm.

Standard Layer Normalization calculates both Mean and Variance:

x_norm = ( x - mean ) / sqrt( variance + epsilon )

Research showed that the stabilizing benefit comes almost entirely from scaling by magnitude, while subtracting the mean adds compute overhead.

RMSNorm skips mean calculation:

x_rms = x / sqrt( mean( x^2 ) + epsilon )

Removing mean subtraction cuts normalization GPU compute time by up to 50 percent while maintaining model training stability.

Say this out loud

Batch Normalization averages across mini batch samples for each feature, making it great for computer vision but dependent on large batch sizes. Layer Normalization averages across features independently for each sample, making it ideal for variable length text in transformers. RMSNorm simplifies Layer Normalization by scaling by root mean square magnitude without calculating the mean, saving GPU time in modern LLMs like LLaMA.

Followups to expect

1. What is Instance Normalization? Normalizes across spatial pixels for each channel in a single image independently, commonly used in style transfer models.
2. What is Group Normalization? Splits feature channels into smaller groups and normalizes within each group, bridging the gap between Layer Norm and Batch Norm for vision models with small batch sizes.