Batch Normalization

By the end of this lesson, you will be able to...

• trace BatchNorm's normalization, learned scale, and learned shift on one batch;

• distinguish batch statistics during training from running statistics during inference;

• diagnose a train/eval or small-batch BatchNorm failure.

Idea in one sentence: BatchNorm normalizes each feature using batch evidence during training, then restores useful representational freedom with learned scale and shift while storing a separate inference-time estimate.

Core Insight

Consider a classifier whose training loss looks healthy, but its one-image inference prediction changes when the model is accidentally left in training mode. The weights did not change. The difference is that BatchNorm used the one-image batch's statistics instead of the running statistics accumulated during training.

BatchNorm is therefore not generic “normalization inside a network.” It is a stateful layer with two modes. It can make activation scale easier to manage during training, but that benefit creates an operational contract: train and inference must use the intended statistics.

One Feature Through a Training Batch

Take one preactivation feature for a batch of four examples:

x = [2, 4, 6, 8]
batch mean = 5
batch variance = 5

With a small epsilon for numerical safety:

x_hat = (x - mean) / sqrt(variance + epsilon)
      ≈ [-1.342, -0.447, 0.447, 1.342]

This batch has centered, unit-scale values after normalization. If BatchNorm stopped here, it would force every later representation into that fixed form. It does not. It learns one gamma and one beta per feature:

y = gamma * x_hat + beta

For illustrative gamma = 2 and beta = 1, the output is approximately [-1.684, 0.106, 1.894, 3.684]. The model can recover a useful scale and offset; normalization provides a stable reference, not a permanent ban on scale.

The forward contract during training is:

input batch
-> batch mean and variance
-> normalized values
-> learned gamma and beta
-> output plus updated running estimates

Backward propagation differentiates through all of these operations, so gamma, beta, and earlier layer weights receive gradients. A framework normally owns this bookkeeping, but the mechanism explains why batch size and mode affect the output.

Read the Axes Before You Normalize

The phrase “normalize a batch” hides an important choice: BatchNorm normally treats each feature independently and reduces over examples. If an activation matrix has shape (batch, features), its mean and variance have one value per feature, not one value for the whole matrix.

For example, let two examples reach a dense layer with two features:

X = [[ 2, 10],
     [ 4, 14]]

mean per feature     = [3, 12]
variance per feature = [1,  4]

Ignoring epsilon to keep the arithmetic visible, normalization produces:

X_hat = [[-1, -1],
         [ 1,  1]]

The first column was compared only with the first column's batch values; the second column was compared only with the second column's. Then each feature gets its own learned parameters. If gamma = [2, 0.5] and beta = [1, -1], the output is:

Y = [[-1, -1.5],
     [ 3, -0.5]]

This dimension check catches a common implementation mistake. Averaging over both examples and features would mix quantities that mean different things: a feature measured in hundreds could dominate a feature measured near zero. For a convolutional tensor, the usual BatchNorm variant similarly keeps separate statistics per channel while reducing over the batch and spatial positions. The exact axis convention depends on the layer type and framework, so inspect the expected tensor shape rather than copying a formula blindly.

The same rule explains the parameter shapes. gamma, beta, running mean, and running variance must broadcast across the batch dimension while retaining one entry per feature or channel. When debugging a custom layer, print those shapes first. A numerically plausible output can still be wrong when the reduction axis or broadcasting axis is wrong.

The Running-State Contract

Using the current batch at inference would make one prediction depend on which other examples happened to arrive beside it. BatchNorm therefore keeps running mean and variance during training. A simplified running-mean update is:

running_mean = momentum * running_mean +
               (1 - momentum) * batch_mean

At inference, it uses stored running values instead of the new request's batch values:

training:  normalize with current batch mean and variance
inference: normalize with stored running mean and variance

Suppose training has stored mean 5 and variance 5. A single inference input x = 8 should normalize to about 1.342, then apply gamma and beta. If the model remains in training mode, that one-item batch has mean 8 and variance 0; it normalizes the same item close to zero instead. The prediction can change because the layer has followed a different contract.

This is the evidence behind the practical rule: set evaluation mode for validation and inference, and checkpoint running statistics with parameters. Weight values alone do not capture a BatchNorm model's behavior.

Train and Eval Are Different State Transitions

It is useful to describe BatchNorm as two related computations, not as one formula with a minor flag. In training mode, the current batch both determines the forward output and updates the layer's stored estimate. In evaluation mode, the stored estimate determines the forward output and is not updated.

A generic moving-average update can be written as running = (1 - alpha) * running + alpha * batch_stat. A small alpha preserves more history; a larger one follows recent batches more closely. Libraries use different names and conventions for their momentum argument, and may use a slightly different variance estimator for stored state. That is why production debugging should read the framework documentation rather than assume an argument called momentum means the same formula everywhere.

Do not use exact equality between a training-mode and evaluation-mode forward pass as a health check. They are expected to differ, even on the same input, because they intentionally use different statistics. A better check is: evaluation outputs should be stable for the same example regardless of which unrelated examples share a request batch. If they are not, inspect the mode switch and whether a custom wrapper bypasses it.

Saving and restoring is part of the same contract. A checkpoint that excludes running buffers can reload without an obvious shape error yet produce changed validation predictions. Treat learned parameters and BatchNorm's non-learned running state as one deployable model state.

What BatchNorm Helps—and What It Costs

By controlling per-feature activation scale on each training batch, BatchNorm often makes optimization less sensitive to initialization and enables useful update sizes. It does not prove that every model trains faster, nor does it replace correct gradients, loss semantics, or validation.

The trade-off is batch dependence. Small batches give noisy mean and variance estimates; a batch of one has no meaningful within-batch variance. Distributed or tiny-batch settings may therefore need synchronized statistics, accumulated statistics, or a normalization method that does not normalize over the batch axis, such as LayerNorm or GroupNorm.

The boundary signal is visible: if train metrics look good while eval metrics collapse, verify train/eval mode and running statistics before changing optimizer or architecture. If behavior varies sharply with batch size, inspect the normalization axis and batch statistics.

Placement, Gradients, and the Small-Batch Boundary

In many feed-forward blocks, BatchNorm is placed after a learned affine or convolutional operation and before the nonlinearity. That convention gives the layer access to the preactivation distribution, but it is an architectural pattern rather than a law: residual blocks, pretrained models, and framework modules may use a different ordering. When reading code, trace the actual tensor sequence instead of inferring it from the class names.

The layer does not merely rescale a detached activation. Its batch mean and variance depend on every example in the batch, so changing one example can alter another example's normalized value and gradient. This coupling is often harmless or useful at ordinary batch sizes, but it makes BatchNorm a poor fit when batches are extremely small, highly variable, or semantically grouped in a way that should not influence one another.

Suppose a detection pipeline has one image per device. Training-mode BatchNorm sees very little evidence for each channel's mean and variance, so its estimates fluctuate from step to step. Increasing the global batch, synchronizing statistics across devices, freezing a pretrained BatchNorm layer, or using GroupNorm can each be reasonable responses. They solve different causes: more samples improve the estimate; synchronization aggregates evidence; freezing stops state drift; another normalization family removes dependence on the batch axis. Choose after identifying which failure is present.

BatchNorm also is not a substitute for regularization. It may change optimization dynamics and sometimes has a regularizing side effect because statistics vary by batch, but it does not replace an explicit decision about data augmentation, weight decay, dropout, or validation design. Keep the causal story narrow: it normalizes feature-wise using a particular source of statistics and maintains state for inference.

A Small Debugging Table

Do not diagnose BatchNorm from loss alone. Print the current mode, batch size, running mean/variance, and one feature's activation range. Those values make the otherwise hidden state inspectable.

Checks

Check: Why are gamma and beta necessary after normalization?

Answer: They let the network learn the scale and offset useful for the task; BatchNorm normalizes a reference representation but does not need to keep the output fixed there.

Check: Why can an inference request of one example produce a different result in training mode?

Answer: Its own batch mean is the example and its batch variance is near zero, unlike the running statistics learned from training batches.

Check: Does BatchNorm remove the need for a sensible initializer?

Answer: No. It can make optimization more forgiving, but initialization, input scale, loss, and learning rate still shape the first forward and backward passes.

Trace It Yourself

For x = [1, 3, 5, 7], compute the batch mean, variance, normalized values, then output with gamma = 0.5 and beta = -1. Finally explain what changes if inference uses running mean 5 and variance 5 instead of this batch's values.

Model answer: Mean is 4, variance is 5, so normalized values are about [-1.342, -0.447, 0.447, 1.342]; the output is about [-1.671, -1.224, -0.776, -0.329]. With running statistics, the same raw values are normalized relative to the training population rather than this four-item batch, so the outputs can differ. That is intended when the running values represent training data; it is a bug when eval mode is accidentally omitted.

Resources

• [PAPER] Batch Normalization — Focus: normalization plus learned scale and shift.
• [DOCUMENTATION] PyTorch: BatchNorm1d — Focus: running statistics and train/eval behavior.
• [TUTORIAL] CS231n: Batch Normalization — Focus: implementation intuition.

Key Takeaways

• BatchNorm uses current batch statistics during training, then learned gamma and beta restore useful scale and shift.
• Running statistics make inference independent of the request batch, but only when eval mode is used correctly.
• Small batches expose BatchNorm's core trade-off: less reliable batch statistics and more operational state to inspect.