Intermediate
14 min read
#Normalization#Batch Normalization#Layer Normalization#RMSNorm#GroupNorm#Transformers

09. Normalization (BatchNorm, LayerNorm, RMSNorm, GroupNorm)

Stabilizing internal covariate shift: Batch Normalization, Layer Normalization in Transformers, Instance Normalization, Group Normalization, and RMSNorm.

Normalization: Complete Notes (Beginner to Advanced)


Introduction#

Normalization techniques transform data (either the raw input features or the intermediate activations inside a network) so that they have a more consistent, well-behaved numerical scale. This consistency has a direct, significant effect on how easily and stably a neural network can be trained.

The general problem normalization solves: neural networks are trained using gradient-based optimization, which is sensitive to the scale of the numbers it operates on. If different features (or different layers' activations) have wildly different scales, ranges, or shifting distributions, the loss surface becomes distorted (elongated, skewed, or uneven), making it much harder for gradient descent to navigate efficiently. Normalization reshapes the data to have consistent statistical properties (typically a mean of 0 and a standard deviation/variance of 1), producing a smoother, more symmetric loss surface that is easier and faster to optimize.


1. Feature Normalization#

Feature normalization refers to scaling the raw input features of a dataset before they are ever fed into the network, so that all features share a similar scale and range. This is distinct from the normalization layers (Batch Norm, Layer Norm, etc.) covered later in this note, which normalize activations inside the network during training.

Two common approaches:

1.1 Standardization (Z-score normalization)#

x^=xμσ+ϵ\hat{x} = \frac{x - \mu}{\sigma + \epsilon}

Transforms each feature to have a mean of 0 and a standard deviation of 1, based on statistics computed from the training set.

1.2 Min-Max Scaling#

x_normalized = (x - min) / (max - min)

Transforms each feature to fit within a fixed range, typically [0, 1].

Why this matters, concretely: imagine a dataset with two features: "age" (ranging roughly 0-100) and "annual income in dollars" (ranging roughly 0-1,000,000). Without normalization, the income feature's raw values are thousands of times larger than the age feature's values. Since the weighted sum computed by a neuron (w1*age + w2*income) is dominated by whichever feature has the larger raw scale, the optimizer needs a much smaller learning rate to avoid huge, unstable updates for the income-related weight, while that same tiny learning rate would make progress on the age-related weight painfully slow. This creates an elongated, poorly-conditioned loss surface (steep in one direction, very shallow in another), similar to the "ravine" problem discussed in the Optimizers notes.

Without feature normalization: training can converge extremely slowly, or require careful, feature-specific tuning of the learning rate, since features on vastly different scales distort the loss surface's geometry, making it harder for a single global learning rate to work well for every weight simultaneously.

With feature normalization: all input features contribute to the weighted sums on a comparable scale, resulting in a more evenly-shaped loss surface that a single, reasonably-chosen learning rate can navigate efficiently for every feature at once.

Important practical note: the mean, standard deviation (or min/max) used for normalization must be computed only from the training set, and then that exact same transformation (using the training set's statistics) must be applied to the validation and test sets. Computing separate statistics for the test set would leak information about the test data's distribution into the preprocessing step, and would also mean the model is evaluated under different scaling conditions than it was trained on.

🐍 Python
import numpy as np def standardize(X_train, X_test): mean = X_train.mean(axis=0) std = X_train.std(axis=0) X_train_norm = (X_train - mean) / std X_test_norm = (X_test - mean) / std # uses TRAINING set's mean and std, not the test set's return X_train_norm, X_test_norm X_train = np.array([[25, 50000], [40, 80000], [35, 65000]]) X_test = np.array([[30, 55000]]) X_train_norm, X_test_norm = standardize(X_train, X_test) print("Normalized training data:\n", X_train_norm) print("Normalized test data:\n", X_test_norm)
🐍 Python
# Using scikit-learn's StandardScaler (a common practical tool) from sklearn.preprocessing import StandardScaler scaler = StandardScaler() X_train_norm = scaler.fit_transform(X_train) # computes mean/std from training data, then transforms it X_test_norm = scaler.transform(X_test) # reuses the SAME mean/std learned from training data

2. Batch Normalization#

Batch Normalization (BatchNorm) normalizes the activations of a layer across the batch dimension, meaning it computes the mean and variance of each individual feature/channel across all the samples currently in the mini-batch, then uses those statistics to normalize that feature/channel.

Formula (for a single feature/channel):

code
mu_B = (1/m) * sum(x_i for i in the batch) (batch mean) sigma_B^2 = (1/m) * sum((x_i - mu_B)^2 for i in the batch) (batch variance) x_hat_i = (x_i - mu_B) / sqrt(sigma_B^2 + epsilon) (normalize) y_i = gamma * x_hat_i + beta (scale and shift)

Where:

  • m is the batch size
  • epsilon is a tiny constant added for numerical stability (avoiding division by zero)
  • gamma and beta are learnable parameters, one pair per feature/channel, allowing the network to learn to undo the normalization if that turns out to be beneficial for a specific feature

Why gamma and beta are necessary: simply forcing every activation to have a mean of 0 and variance of 1 could actually restrict what the layer is able to represent (for example, some activation functions perform best with inputs in a specific range that isn't necessarily centered at 0). Adding learnable scale (gamma) and shift (beta) parameters lets the network learn the ideal scale and mean for each normalized feature, effectively giving it the option to "undo" the normalization (by learning gamma = sqrt(sigma_B^2 + epsilon) and beta = mu_B) if the un-normalized version actually works better for that particular feature, while defaulting to the stabilizing effects of normalization otherwise.

Behavior at inference/test time: at inference, you often process a single sample (or a small batch that isn't representative of the training distribution), so calculating a fresh mean/variance from the current batch wouldn't be reliable or even well-defined for a single sample. Instead, Batch Normalization keeps a running average of the mean and variance calculated during training (updated at every training step), and uses these stored running statistics at inference time instead of computing new ones.

Without Batch Normalization: as data flows through a deep network, the distribution of each layer's inputs can shift and change during training as the weights of earlier layers are updated, a phenomenon researchers termed "internal covariate shift." This forces each layer to continuously re-adapt to a constantly changing input distribution, which slows down training significantly. Deep networks also become more sensitive to weight initialization and require more careful learning rate tuning.

With Batch Normalization: each layer's inputs are kept at a consistently stable scale and distribution throughout training, which allows for faster training, the use of higher learning rates, and reduces sensitivity to the specific weight initialization scheme chosen, since BatchNorm actively re-centers and re-scales activations at every layer regardless of what the incoming weight values happen to be.

Where BatchNorm is typically placed: usually inserted between the linear transformation and the activation function of a layer: Linear -> BatchNorm -> Activation.

A key limitation: Batch Normalization's statistics depend on having a reasonably large batch size to compute a meaningful mean and variance. With very small batch sizes, the batch statistics become noisy and unreliable estimates of the true data distribution, which can hurt BatchNorm's effectiveness. This limitation is part of the motivation for Layer Normalization, Instance Normalization, and Group Normalization, covered next.

🐍 Python
import numpy as np def batch_norm_forward(X, gamma, beta, epsilon=1e-5): # X shape: (batch_size, num_features) -- normalize across the batch dimension (axis 0) mu = X.mean(axis=0) var = X.var(axis=0) X_hat = (X - mu) / np.sqrt(var + epsilon) return gamma * X_hat + beta X = np.array([[1.0, 2.0], [3.0, 4.0], [5.0, 6.0]]) # 3 samples, 2 features gamma = np.array([1.0, 1.0]) beta = np.array([0.0, 0.0]) output = batch_norm_forward(X, gamma, beta) print("Batch normalized output:\n", output)
🐍 Python
# Using PyTorch's built-in BatchNorm import torch import torch.nn as nn model = nn.Sequential( nn.Linear(10, 20), nn.BatchNorm1d(20), # normalizes across the batch dimension for each of the 20 features nn.ReLU() ) model.train() # uses current batch statistics, and updates the running mean/variance X = torch.randn(8, 10) # batch of 8 samples output_train = model(X) model.eval() # uses the stored RUNNING mean/variance instead of the current batch's statistics output_eval = model(X)

3. Layer Normalization#

Layer Normalization (LayerNorm) normalizes activations across the feature dimension, independently for each individual sample, rather than across the batch dimension like BatchNorm.

Formula (for a single sample, across all its features):

code
mu = (1/H) * sum(x_j for all features j in this sample) (mean over the sample's own features) sigma^2 = (1/H) * sum((x_j - mu)^2 for all features j in this sample) (variance over the sample's own features) x_hat_j = (x_j - mu) / sqrt(sigma^2 + epsilon) y_j = gamma_j * x_hat_j + beta_j

Where H is the number of features (hidden units) in that layer for a single sample, and the mean/variance are computed per sample, not across the batch.

Key difference from Batch Normalization:

AspectBatch NormalizationLayer Normalization
Statistics computed acrossThe batch dimension (across samples), per featureThe feature dimension (across features), per sample
Depends on batch size?Yes, needs a reasonably large batch for stable statisticsNo, works identically regardless of batch size, even batch size of 1
Behavior at inferenceUses stored running statistics from trainingComputes statistics fresh, the same way as training (no running average needed)
Common use caseConvolutional networks (CNNs), computer visionRecurrent networks (RNNs), and especially Transformers (used in models like BERT, GPT)

Why LayerNorm is preferred for Transformers and sequence models over BatchNorm: in sequence models, especially with variable-length sequences (like sentences of different lengths in NLP) or very small batch sizes, BatchNorm's batch-dependent statistics become unreliable or awkward to define correctly across a batch of differently-shaped or differently-padded sequences. Since LayerNorm normalizes each individual sample independently based only on its own features, it is completely unaffected by batch size or the characteristics of other samples in the batch, making it far more robust and stable in these settings.

Without Layer Normalization (using BatchNorm in a Transformer, for example): the model's behavior at inference would depend on stored running statistics from training that may not represent variable-length inputs well, and training could be less stable with the very small effective batch sizes sometimes needed for large models or long sequences.

With Layer Normalization: since each sample is normalized independently based only on its own feature values, the technique behaves identically regardless of batch size or the makeup of other samples in the batch, which is precisely why it became the standard normalization choice in Transformer-based architectures.

🐍 Python
import numpy as np def layer_norm_forward(X, gamma, beta, epsilon=1e-5): # X shape: (batch_size, num_features) -- normalize across the FEATURE dimension (axis 1), per sample mu = X.mean(axis=1, keepdims=True) var = X.var(axis=1, keepdims=True) X_hat = (X - mu) / np.sqrt(var + epsilon) return gamma * X_hat + beta X = np.array([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]]) # 2 samples, 3 features each gamma = np.array([1.0, 1.0, 1.0]) beta = np.array([0.0, 0.0, 0.0]) output = layer_norm_forward(X, gamma, beta) print("Layer normalized output:\n", output)
🐍 Python
# Using PyTorch's built-in LayerNorm import torch import torch.nn as nn layer_norm = nn.LayerNorm(normalized_shape=20) # normalizes across the last dimension of size 20 X = torch.randn(8, 20) # batch of 8 samples, 20 features each output = layer_norm(X) # Unlike BatchNorm, this behaves identically whether X has 1 sample or 1000 samples

4. Instance Normalization#

Instance Normalization (InstanceNorm) is primarily used for image data (with shape batch, channels, height, width) and normalizes each individual sample's each individual channel separately, computing statistics only across that channel's spatial dimensions (height and width), independently for each sample.

Formula (for one specific sample and one specific channel):

Mathematical Formulation
mu = mean of all pixel values in this sample's this channel, across height and width
sigma^2 = variance of all pixel values in this sample's this channel, across height and width

x_hat = (x - mu) / sqrt(sigma^2 + epsilon)
y = gamma * x_hat + beta

Key distinction from the other normalization types:

NormalizationStatistics computed across (for an image batch of shape [N, C, H, W])
Batch NormAcross N (batch), H, and W -- separately per channel C
Layer NormAcross C, H, and W -- separately per sample N
Instance NormAcross H and W only -- separately per sample N AND per channel C

Why Instance Normalization is especially useful for style transfer: Instance Normalization was originally developed for and remains particularly popular in style transfer tasks (generating an image that combines the content of one image with the artistic style of another). The reasoning is that an image's specific contrast and per-channel brightness statistics often encode a large part of its "style," while the actual arrangement of content (shapes, edges, objects) is comparatively less sensitive to these exact statistics. By normalizing away each individual image's own specific per-channel contrast and mean brightness (independently, without mixing in information from other images in the batch, unlike BatchNorm), Instance Normalization effectively strips out much of an individual image's low-level style information, which is exactly the type of transformation useful for tasks where you want to swap or blend styles between different images.

Without Instance Normalization (relying on BatchNorm for a style-transfer-focused task): BatchNorm's statistics would be computed jointly across multiple different images in a batch, mixing together stylistic information from unrelated images rather than treating each image's own style characteristics independently, which is not well-suited to per-image style manipulation.

With Instance Normalization: each image's own contrast and per-channel statistics are normalized completely independently of every other image in the batch, isolating and removing each image's individual low-level style signal in a way that's directly useful for style-transfer and certain image generation tasks.

🐍 Python
import numpy as np def instance_norm_forward(X, gamma, beta, epsilon=1e-5): # X shape: (batch_size, channels, height, width) # Normalize across height and width (axes 2, 3), separately for each sample and channel mu = X.mean(axis=(2, 3), keepdims=True) var = X.var(axis=(2, 3), keepdims=True) X_hat = (X - mu) / np.sqrt(var + epsilon) return gamma * X_hat + beta np.random.seed(0) X = np.random.randn(2, 3, 4, 4) # 2 images, 3 channels, 4x4 spatial size gamma = np.ones((1, 3, 1, 1)) beta = np.zeros((1, 3, 1, 1)) output = instance_norm_forward(X, gamma, beta) print("Output shape:", output.shape)
🐍 Python
# Using PyTorch's built-in InstanceNorm import torch import torch.nn as nn instance_norm = nn.InstanceNorm2d(num_features=3) # 3 channels X = torch.randn(2, 3, 4, 4) # batch of 2 images, 3 channels, 4x4 spatial size output = instance_norm(X)

5. Group Normalization#

Group Normalization (GroupNorm) is a middle ground between Layer Normalization and Instance Normalization. It divides a sample's channels into a fixed number of groups, and computes normalization statistics across the spatial dimensions and the channels within each group, independently per sample.

Formula (conceptually, for one sample and one group of channels):

Mathematical Formulation
For each sample, divide the C channels into G groups (each group has C/G channels).
For each group:
    mu = mean of all values across the group's channels AND the spatial dimensions (H, W)
    sigma^2 = variance across that same set of values

    x_hat = (x - mu) / sqrt(sigma^2 + epsilon)
    y = gamma * x_hat + beta   (gamma, beta are still typically per-channel, applied after normalization)

How Group Normalization relates to Layer Norm and Instance Norm as special cases:

  • If the number of groups G = 1 (all channels in a single group), Group Normalization becomes mathematically equivalent to Layer Normalization applied across channels and spatial dimensions.
  • If the number of groups G = C (each channel is its own separate group), Group Normalization becomes mathematically equivalent to Instance Normalization.

This makes Group Normalization a flexible, tunable generalization that sits between these two other techniques, with the number of groups G as a hyperparameter you choose.

Why Group Normalization was developed (the specific problem it solves): Batch Normalization's reliance on batch statistics works well with large batch sizes, but its performance degrades noticeably when the batch size is small (e.g., in tasks like object detection or video processing, where large images or many video frames per sample make it computationally infeasible to fit a large batch into GPU memory at once). Group Normalization was specifically designed to be independent of batch size entirely (since, like Layer Norm and Instance Norm, it computes statistics per-sample rather than across the batch), while still, unlike pure Layer Norm, respecting the fact that in convolutional networks, groups of related channels (e.g., channels detecting related types of visual features) often benefit from being normalized together rather than either individually (Instance Norm) or all together as one giant group across every channel (Layer Norm).

Without Group Normalization (using Batch Normalization with a very small batch size): the batch statistics computed from just a few samples become noisy and unreliable estimates of the true data distribution, which can degrade both training stability and the model's final accuracy.

With Group Normalization: since statistics are computed per-sample (not per-batch) but still across a meaningful, tunable grouping of related channels (not just a single channel like Instance Norm, or all channels like Layer Norm), the technique remains stable and effective regardless of batch size, which has made it a popular choice in computer vision tasks like object detection and segmentation, where small batch sizes are common due to memory constraints.

🐍 Python
import numpy as np def group_norm_forward(X, num_groups, gamma, beta, epsilon=1e-5): # X shape: (batch_size, channels, height, width) N, C, H, W = X.shape G = num_groups X_grouped = X.reshape(N, G, C // G, H, W) mu = X_grouped.mean(axis=(2, 3, 4), keepdims=True) var = X_grouped.var(axis=(2, 3, 4), keepdims=True) X_hat = (X_grouped - mu) / np.sqrt(var + epsilon) X_hat = X_hat.reshape(N, C, H, W) return gamma * X_hat + beta np.random.seed(0) X = np.random.randn(2, 8, 4, 4) # 2 images, 8 channels, 4x4 spatial size gamma = np.ones((1, 8, 1, 1)) beta = np.zeros((1, 8, 1, 1)) output = group_norm_forward(X, num_groups=4, gamma=gamma, beta=beta) # 4 groups of 2 channels each print("Output shape:", output.shape)
🐍 Python
# Using PyTorch's built-in GroupNorm import torch import torch.nn as nn group_norm = nn.GroupNorm(num_groups=4, num_channels=8) # 8 channels divided into 4 groups of 2 X = torch.randn(2, 8, 4, 4) output = group_norm(X)

6. RMS Normalization#

RMS Normalization (Root Mean Square Normalization, or RMSNorm) is a simplified, more computationally efficient variant of Layer Normalization. It normalizes based only on the root mean square of the values, entirely skipping the mean-centering step that LayerNorm performs.

Formula:

code
RMS(x) = sqrt( (1/H) * sum(x_j^2 for all features j) + epsilon ) x_hat_j = x_j / RMS(x) y_j = gamma_j * x_hat_j

Where H is the number of features, and gamma is a learnable per-feature scale parameter (RMSNorm typically does not include a learnable shift/bias parameter beta, unlike LayerNorm and BatchNorm).

Key difference from Layer Normalization: LayerNorm first subtracts the mean (re-centering the values around 0) before dividing by the standard deviation (re-scaling). RMSNorm skips the mean subtraction step entirely, and instead divides directly by the root mean square of the values, which combines the effects of both the values' scale and their deviation from zero into a single computation.

Mathematical Formulation
LayerNorm:  x_hat = (x - mean(x)) / sqrt(variance(x) + epsilon)      [re-centers AND re-scales]
RMSNorm:    x_hat = x / sqrt(mean(x^2) + epsilon)                     [re-scales only, no re-centering]

Why skipping the mean-centering step is a reasonable simplification: the original authors of RMSNorm found, through empirical research, that the primary benefit of LayerNorm largely comes from re-scaling the values to control their magnitude, rather than specifically from re-centering them around a mean of zero. By removing the mean-calculation and mean-subtraction steps, RMSNorm requires noticeably less computation than LayerNorm (skipping an entire pass to compute the mean, and the subtraction operation itself), while empirically achieving comparable performance and training stability in practice.

Without RMSNorm (using standard LayerNorm everywhere for a very large model): every single normalization step across every layer requires computing both a mean and a variance, and performing a mean-subtraction, which adds up to meaningful extra computational cost and processing time when repeated across the billions of parameters and countless normalization layers found in very large modern models.

With RMSNorm: the computation is simplified to just a root-mean-square calculation and a single division, removing the mean-related computation and subtraction step entirely, providing a meaningful efficiency improvement in very large-scale models with minimal to no loss in model quality, which is why RMSNorm has been adopted in several prominent modern large language model architectures (such as LLaMA) in place of standard LayerNorm.

🐍 Python
import numpy as np def rms_norm_forward(X, gamma, epsilon=1e-8): # X shape: (batch_size, num_features) rms = np.sqrt(np.mean(X ** 2, axis=-1, keepdims=True) + epsilon) X_hat = X / rms return gamma * X_hat X = np.array([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]]) gamma = np.array([1.0, 1.0, 1.0]) output = rms_norm_forward(X, gamma) print("RMS normalized output:\n", output)
🐍 Python
# Using PyTorch (RMSNorm is available as a built-in module in recent PyTorch versions) import torch import torch.nn as nn rms_norm = nn.RMSNorm(normalized_shape=3) X = torch.tensor([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]]) output = rms_norm(X) print(output)

Summary Table: Normalization Techniques#

TechniqueStatistics Computed AcrossDepends on Batch Size?Typical Use Case
Feature NormalizationEntire training dataset (once, before training)No (applied to raw inputs, not during training)Preprocessing raw input features
Batch NormalizationBatch dimension, per channel/featureYesCNNs, computer vision
Layer NormalizationFeature dimension, per sampleNoTransformers, RNNs, NLP
Instance NormalizationSpatial dimensions, per sample AND per channelNoStyle transfer, image generation
Group NormalizationSpatial dimensions + a group of channels, per sampleNoObject detection/segmentation with small batch sizes
RMS NormalizationFeature dimension (scale only, no mean-centering), per sampleNoLarge language models (e.g., LLaMA), efficient Transformers

Quick Recap (Beginner to Advanced Flow)#

  1. Feature normalization scales raw input features (using statistics from the training set only) before training even begins, preventing features on vastly different scales from distorting the loss surface.
  2. Batch Normalization normalizes activations across the batch dimension per feature/channel, using batch statistics during training and stored running statistics during inference; it reduces internal covariate shift but depends on a reasonably large batch size.
  3. Layer Normalization normalizes across the feature dimension independently per sample, making it completely independent of batch size, which is why it's the standard choice in Transformer architectures.
  4. Instance Normalization normalizes each sample's each channel separately across only its spatial dimensions, effectively isolating and removing per-image style information, making it especially useful for style transfer.
  5. Group Normalization divides channels into groups and normalizes within each group per sample, acting as a tunable middle ground between Layer Norm (G=1) and Instance Norm (G=C), and remains effective even with very small batch sizes.
  6. RMS Normalization simplifies Layer Normalization by skipping the mean-centering step, normalizing purely based on the root mean square of the values, offering a computationally cheaper alternative widely adopted in modern large language models.
Knowledge Checkpoint

09. Normalization Checkpoint

Q1.Across which tensor dimensions does Batch Normalization compute mean and variance?
AAcross the batch dimension (N) and spatial dimensions (H, W), normalizing each channel independently.
BAcross all channels and features for a single sample.
COnly across time steps in an RNN sequence.
DAcross model parameters during weight decay.
Q2.Why is Layer Normalization (LayerNorm) used in Transformers and RNNs instead of Batch Normalization?
ALayerNorm computes statistics across all feature/channel dimensions for each individual sample independently of batch size and sequence length.
BBatchNorm is mathematically incompatible with GPUs.
CLayerNorm requires zero learnable parameters.
DBatchNorm cannot run on floating point data.
Q3.What computational simplification does RMSNorm (used in LLaMA and modern LLMs) make compared to standard LayerNorm?
AIt removes mean centering and normalizes activations strictly by the Root Mean Square of inputs, achieving 10-50% speedup with equivalent stability.
BIt replaces floating point math with integer approximations.
CIt skips the scale parameter gamma.
DIt computes variance across the entire vocabulary.
Track Your Learning

Finished studying this notebook?

Mark this guide as completed to update your course progress roadmap.