Batch and Layer Normalization

Normalization layers transform intermediate activations, but their axes and state differ. Batch normalization uses mini-batch statistics during training and running estimates during evaluation. Layer normalization computes statistics over the configured trailing dimensions for each example and uses the same input-dependent calculation in both modes.

Prerequisites

Complete Regularization and Dropout. You should understand feature columns, batches, and model.train()/model.eval().

Learning goals

• identify which axes batch and layer normalization reduce.
• explain running statistics and learned affine parameters separately.
• recognize small-batch and train/eval mismatch risks.

Worked example

The (4, 3) input represents four examples and three features. BatchNorm1d(3) normalizes each feature across the batch; LayerNorm(3) normalizes the three features within each row.

import torch
from torch import nn

x = torch.tensor(
    [[1.0, 2.0, 3.0], [3.0, 4.0, 5.0], [5.0, 6.0, 7.0], [7.0, 8.0, 9.0]]
)
batch_norm = nn.BatchNorm1d(3, affine=False, momentum=1.0)
layer_norm = nn.LayerNorm(3, elementwise_affine=False)

batch_norm.train()
batch_training = batch_norm(x)
running_mean = batch_norm.running_mean.clone()
batch_norm.eval()
batch_evaluation = batch_norm(x)
layer_output = layer_norm(x)

def rounded_means(values):
    return [0.0 if abs(value) < 5e-7 else round(value, 6) for value in values]


feature_means = rounded_means(batch_training.mean(dim=0).tolist())
row_means = rounded_means(layer_output.mean(dim=1).tolist())
print(f"batch_train_feature_mean={feature_means}")
print(f"batch_running_mean={running_mean.tolist()}")
print(f"layer_row_mean={row_means}")
print(f"eval_differs={not torch.allclose(batch_training, batch_evaluation)}")

Output:

batch_train_feature_mean=[0.0, 0.0, 0.0]
batch_running_mean=[4.0, 5.0, 6.0]
layer_row_mean=[0.0, 0.0, 0.0, 0.0]
eval_differs=True

The display maps rounded values very close to zero to 0.0, avoiding platform-dependent signed-zero formatting. Both modules normally include learned scale and bias parameters; affine=False and elementwise_affine=False remove them here to expose the normalization calculation.

PyTorch's BatchNorm momentum controls the running-statistic update and is not optimizer momentum. With tiny or nonrepresentative batches, estimates can be noisy or systematically mismatched with deployment. Validation must use eval() so it does not update running statistics from validation examples. LayerNorm avoids cross-example dependence, but choosing the wrong normalized shape can erase structure you intended to preserve.

Normalization may smooth optimization and permit different learning rates, yet it is not a substitute for feature reasoning, split integrity, or evaluation. Its ordering relative to affine layers, activations, residual paths, and dropout is an architectural decision.

Common mistakes

• Assuming every normalization layer standardizes over the batch.
• Updating BatchNorm running statistics on validation or test data.
• Interpreting BatchNorm momentum like SGD momentum.
• Forgetting that a batch of one can make batch-statistic behavior unsuitable.
• Adding normalization to fix exploding loss before checking scale, learning rate, and invalid values.

Practice

Compute the three batch feature means and the first row mean before normalization. Then change LayerNorm(3) to normalize an incompatible shape and read the resulting error as a shape-contract diagnostic.

Next steps

Continue with Debugging Training Runs to combine loss, gradient, activation, data, and mode evidence.

Primary sources

Checked 2026-10-16: PyTorch BatchNorm1d, PyTorch LayerNorm, Batch Normalization paper, and Layer Normalization paper.