Batch Normalization

Batch normalization rescales activations using statistics computed across a mini-batch, then applies a learned scale and shift. For an input shaped (N, C), it treats each feature column separately.

Its most important distinction is training versus inference. Training uses batch statistics; evaluation normally uses stored running estimates. Confusing those modes can change predictions even when the model weights stay fixed.

For a batch of activations shaped (N, C) (N examples, C features) BN computes the mean and variance down each column (across the N examples) and standardizes that feature. Before the learned affine transform, each feature has zero batch mean and variance close to one (depending on $\epsilon$). Learned $\gamma$ and $\beta$ then set its scale and offset. They can undo normalization for fixed statistics, but cannot exactly undo every batch's changing statistics with one fixed pair of parameters.

Contrast with Layer Norm, which normalizes across the features of a single example (along C, per row): independent of batch size, which is why Transformers use LN, not BN.

For a mini-batch $\mathcal{B} = \{x_1, \dots, x_N\}$ (per feature), BN computes:

$\epsilon$ is a small constant for numerical stability. $\gamma$ and $\beta$ are learned per feature.

Train vs. eval: the critical distinction. At training time, $\mu_\mathcal{B}$ and $\sigma_\mathcal{B}^2$ are computed from the current batch. The layer also maintains running estimates via an exponential moving average:

With PyTorch's default track_running_stats=True, evaluation uses fixed running statistics, making this layer independent of the other examples in the batch. If tracking is disabled, it uses batch statistics in evaluation too. Forgetting to switch to eval mode (model.eval()) is a classic bug: predictions then depend on whatever else is in the batch.

Key takeaway: Batch Norm is two different behaviors wearing one layer: batch statistics while training, frozen running statistics at inference. The model.eval() switch is not optional; it changes what the layer computes.

The normalization-family landscape, and when each applies:

For a two-dimensional (N, C) input, this matches the default PyTorch distinction: the training forward pass uses population variance, while the running estimate uses the unbiased sample variance. Training needs at least two examples here.

1import torch
2from torch import Tensor, nn
3
4
5class BatchNorm1d(nn.Module):
6    def __init__(self, num_features: int, eps: float = 1e-5, momentum: float = 0.1) -> None:
7        super().__init__()
8        self.eps, self.momentum = eps, momentum
9        self.gamma = nn.Parameter(torch.ones(num_features))   # learnable scale
10        self.beta = nn.Parameter(torch.zeros(num_features))   # learnable shift
11        self.register_buffer("running_mean", torch.zeros(num_features))
12        self.register_buffer("running_var", torch.ones(num_features))
13
14    def forward(self, x: Tensor) -> Tensor:  # x: (N, num_features)
15        if self.training:
16            if x.size(0) < 2:
17                raise ValueError("Training requires at least two examples")
18            mean = x.mean(dim=0)
19            var = x.var(dim=0, unbiased=False)               # biased var, as in BN
20            with torch.no_grad():                            # update running stats
21                self.running_mean.mul_(1 - self.momentum).add_(self.momentum * mean)
22                unbiased_var = var * x.size(0) / (x.size(0) - 1)
23                self.running_var.mul_(1 - self.momentum).add_(self.momentum * unbiased_var)
24        else:
25            mean, var = self.running_mean, self.running_var  # frozen at inference
26        x_hat = (x - mean) / torch.sqrt(var + self.eps)
27        return self.gamma * x_hat + self.beta
28
29
30bn = BatchNorm1d(16)
31out = bn(torch.randn(32, 16))
32assert out.shape == (32, 16)
33bn.eval()                                                    # switch to running stats
34assert bn(torch.randn(1, 16)).shape == (1, 16)              # works for batch size 1

1. Conceptual: What does Batch Norm actually normalize, and along which axis? (Each feature, across the examples in the batch: per-column for (N, C) input.)
2. Implementation: What changes between training and inference, and why? (Train uses batch stats + updates running averages; eval uses frozen running stats so output is independent of batch composition.)
3. Applied: Why does BN struggle with very small batch sizes, and what would you use instead? (Batch stats become noisy/unreliable; use Group Norm or Layer Norm, which don't depend on batch size.)
4. Systems-level: Why do Transformers use Layer Norm rather than Batch Norm? (LN normalizes per-example across features: independent of batch size and sequence length, and well-behaved for variable-length sequences and small/streaming batches.)
5. Failure modes: Your model trains well but predictions are unstable in production. What's a likely BN-related cause? (Forgot model.eval(), so BN uses batch stats at inference: outputs depend on what else is batched together.)

Without looking: write the four BN equations (mean, variance, normalize, scale-shift), name what $\gamma$ and $\beta$ are for, and explain in one sentence exactly what differs at inference time. Then compare against Stage 3.

Sources and further reading

• Ioffe and Szegedy, Batch Normalization: normalization and learned affine parameters.
• PyTorch BatchNorm1d documentation: running statistics, variance estimators, and evaluation behavior.