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)#
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.
🐍 PythonInteractive WebAssemblyimport 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)
🐍 PythonInteractive WebAssembly# 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):
codemu_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:
mis the batch sizeepsilonis a tiny constant added for numerical stability (avoiding division by zero)gammaandbetaare 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.
🐍 PythonInteractive WebAssemblyimport 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)
🐍 PythonInteractive WebAssembly# 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):
codemu = (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:
| Aspect | Batch Normalization | Layer Normalization |
|---|---|---|
| Statistics computed across | The batch dimension (across samples), per feature | The feature dimension (across features), per sample |
| Depends on batch size? | Yes, needs a reasonably large batch for stable statistics | No, works identically regardless of batch size, even batch size of 1 |
| Behavior at inference | Uses stored running statistics from training | Computes statistics fresh, the same way as training (no running average needed) |
| Common use case | Convolutional networks (CNNs), computer vision | Recurrent 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.
🐍 PythonInteractive WebAssemblyimport 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)
🐍 PythonInteractive WebAssembly# 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 Formulationmu = 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:
| Normalization | Statistics computed across (for an image batch of shape [N, C, H, W]) |
|---|---|
| Batch Norm | Across N (batch), H, and W -- separately per channel C |
| Layer Norm | Across C, H, and W -- separately per sample N |
| Instance Norm | Across 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.
🐍 PythonInteractive WebAssemblyimport 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)
🐍 PythonInteractive WebAssembly# 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 FormulationFor 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.
🐍 PythonInteractive WebAssemblyimport 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)
🐍 PythonInteractive WebAssembly# 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:
codeRMS(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 FormulationLayerNorm: 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.
🐍 PythonInteractive WebAssemblyimport 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)
🐍 PythonInteractive WebAssembly# 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#
| Technique | Statistics Computed Across | Depends on Batch Size? | Typical Use Case |
|---|---|---|---|
| Feature Normalization | Entire training dataset (once, before training) | No (applied to raw inputs, not during training) | Preprocessing raw input features |
| Batch Normalization | Batch dimension, per channel/feature | Yes | CNNs, computer vision |
| Layer Normalization | Feature dimension, per sample | No | Transformers, RNNs, NLP |
| Instance Normalization | Spatial dimensions, per sample AND per channel | No | Style transfer, image generation |
| Group Normalization | Spatial dimensions + a group of channels, per sample | No | Object detection/segmentation with small batch sizes |
| RMS Normalization | Feature dimension (scale only, no mean-centering), per sample | No | Large language models (e.g., LLaMA), efficient Transformers |
Quick Recap (Beginner to Advanced Flow)#
- 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.
- 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.
- 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.
- 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.
- 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. - 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.
09. Normalization Checkpoint
Finished studying this notebook?
Mark this guide as completed to update your course progress roadmap.