Batch, layer, instance and group normalization

Feed a linear layer y = Wx two raw features: a pixel intensity somewhere in 0 to 255 and a binary flag in 0 to 1. For both features to have a similar influence on y, the two corresponding entries of W must differ by about two orders of magnitude. And because neither feature is centered at zero, the layer needs a large Bias just to bring its outputs near zero, where the next activation function does useful work.

Those are the two troubles every layer faces. Inputs that are not centered force a large bias, and inputs with different scales per element force the entries of W to vary a lot. Both make gradient descent slow, because one learning rate has to serve directions of very different steepness, so the update zig-zags along the steep direction while creeping along the shallow one. Goodfellow, Bengio and Courville motivate the same problem with a deep chain of layers: every update to an early layer changes the statistics that every later layer sees, so the layers keep re-adapting to each other.

The idea, due to Ioffe and Szegedy in 2015, is to stop hoping the inputs are nice and to force them to be. For each dimension k, subtract the mean and divide by the standard deviation, so that the dimension has zero mean and unit variance:

The part that makes this a layer rather than a preprocessing step is that the map is differentiable. It sits inside the network, at any hidden layer, and gradients flow through it exactly as through a convolution. The CS231n notes call it differentiable preprocessing built into the network itself. On slide 27 this is why Normalization appears as the fifth component of a CNN, beside convolution, pooling, fully connected layers and activations, and slide 39 later returns to the same map with the activation function circled instead.

Start with one column of numbers. Take the mini-batch from slide 38 after its ReLU, and look only at feature 1 across the four samples: (3, 1, 3, 2). Standardizing that column is the entire batch norm computation in miniature.

Now the general rule. The input to a batch norm layer after a fully connected layer is a matrix x of shape N x D: N samples in the mini-batch (the Batch size) by D features. Every average runs down the N axis, and that is the whole story of the axes: one statistic per feature column, averaged over the sample rows.

The epsilon under the root is a tiny constant, 1e-5 by default in PyTorch, that prevents division by zero when a feature happens to be constant across the batch. It changes nothing numerically otherwise, but it appears in every formula and in every fusion derivation, so keep it.

Why gamma and beta exist

Is forcing zero mean and unit variance always good? It is a hard constraint, and a network is not always best served by activations pinned to a standard distribution. So the layer gets two learnable vectors of shape D, gamma and beta, and outputs y = gamma x_hat + beta. Ioffe and Szegedy point out that setting gamma = sqrt(Var[x]) and beta = E[x] recovers the original activations, if that were the optimal thing to do. The network can therefore undo normalization entirely, which means normalization can never make the model less expressive; it can only change what is easy to learn.

Goodfellow, Bengio and Courville explain why the new parametrization learns more easily even though it can represent the same functions. In the original network the mean of a hidden unit is set by a long chain of interacting weights in all the layers below. After BN the mean is set by beta alone and the scale by gamma alone, two parameters with a direct, local effect on the quantity gradient descent needs to move.

The tanh plot on slide 31 makes the point concrete. The derivative 1 - tanh squared (x) peaks at 1 at the origin and falls below 0.1 once |x| exceeds about 1.8. If normalized values feed a tanh, then gamma decides how far they reach into the flat, saturated tails and beta decides where the center of the band sits. A small gamma keeps the layer almost linear; a large one lets it saturate on purpose. The model chooses the amount of saturation instead of inheriting it from whatever scale the previous layer produced.

A keyword-spotting model on a microcontroller hears one audio clip at a time. A batch of one has zero variance, so the statistics of the "batch" are meaningless, and even a batch of ten would make the answer for one clip depend on which nine clips happened to arrive with it. Ioffe and Szegedy state the requirement plainly: at inference the output should depend only on the input, deterministically.

So the layer changes its source of statistics. During training it keeps a running average of the mini-batch means and variances it has seen. At inference it uses those stored values and ignores the batch entirely. PyTorch updates the buffers as running = (1 - momentum) running + momentum x_batch with momentum 0.1, normalizes with the biased variance during training, and stores the unbiased variance (divide by N - 1) in running_var, following the paper's m/(m - 1) correction. Goodfellow, Bengio and Courville note the payoff: with running averages the model can be evaluated on a single example.

Frozen statistics turn BN into an affine map

Once mu and sigma squared are constants, batch norm on feature j is nothing more than multiply by one number and add another. And if the layer before it is a linear layer z = Wx + b, two affine maps in a row are one affine map. Write it out:

Read the first formula per output row j: the whole row of W is multiplied by the scalar gamma_j / sqrt(sigma_j squared + eps), and the bias becomes that scalar times (b_j - mu_j), plus beta_j. The same holds for a convolution: each output filter is scaled by its channel's factor. Ioffe and Szegedy already say this in the paper: the normalization is a linear transform that may be composed with the scaling by gamma and shift by beta into a single linear transform. PyTorch ships it as torch.nn.utils.fusion.fuse_conv_bn_eval, which requires both modules in eval mode with their running buffers computed. This is Batch norm fusion.

For embedded deployment this is the reason the slide says zero overhead at test time. A fused network runs no extra multiply-accumulate operations, moves no extra activations through memory, launches one fewer kernel per layer, and stores four fewer vectors per layer in flash. In later lectures on pruning, the very gamma factors that get fused away are also what network slimming reads to decide which channels matter.

What BN buys during training, and the one way it bites

• Much easier training and higher learning rates: normalizing each layer's input stops small parameter changes from amplifying into large, suboptimal changes in activations and gradients as they pass through many layers (Ioffe and Szegedy, section 3.3). On ImageNet the paper reaches the same accuracy as its baseline in 14 times fewer training steps.
• Better gradient flow: the tanh argument from the previous concept applies at every layer. Activations are kept out of the saturated tails unless the model wants them there, so derivatives stay usable.
• Robustness to initialization: scaling the weights of a layer by a constant scales its output by the same constant, and normalization divides that constant back out. The scale of the initial weights largely stops mattering.
• Regularization during training, the "How?" on slide 33: each sample is normalized with statistics that depend on the other random samples in its mini-batch, so the network no longer produces deterministic values for a given training example (section 3.4). That injected noise acts like a mild regularizer, and the paper reports that dropout can be removed or reduced in strength. Dive into Deep Learning notes that batch sizes around 50 to 100 inject roughly the right amount of noise.

The red line on slide 33 is the hazard that follows from everything above. The layer computes something different in training mode (batch statistics) and evaluation mode (running statistics). Forget to call model.eval() and your deployed model's answer for one input depends on the other inputs in the batch and degrades silently; for a batch of one, x_hat is 0 by construction so every feature collapses to beta (PyTorch's BatchNorm1d refuses to run at all in that case). It is one of the most common bugs in deep learning code precisely because, for any larger batch, nothing crashes.

Take one activation tensor from a convolution layer, N x C x H x W = 8 x 64 x 32 x 32. It holds about half a million numbers. Every normalization in this part does the same arithmetic, subtract a mean and divide by a standard deviation, then scale and shift. The only thing that distinguishes them is the answer to one question: which of those numbers are averaged together to make one mean?

Wu and He write all four as one formula. Each value x_i is normalized with a mean and standard deviation computed over a set S_i, and different normalizations simply use different definitions of that set:

Two patterns in that table carry most of the marks. First, the shape of mu is the tensor shape with a 1 wherever you averaged and the full size everywhere else. Batch norm averages over N, H, W and leaves 1 x C x 1 x 1; instance norm averages over H, W only and leaves N x C x 1 x 1. Second, gamma and beta do not follow the mu shape: they are per channel for BN, IN and GN. Ioffe and Szegedy learn a pair per feature map rather than per activation, and Wu and He write that BN, LN, IN and GN all learn a per-channel linear transform, the convention the table and the explorer below use. A BatchNorm2d over 64 channels therefore holds 128 learnable values, whatever the image size.

Why the batch-free variants exist

Layer norm was introduced by Ba, Kiros and Hinton for exactly the cases where a batch is awkward. Its statistics come from all of the summed inputs to the neurons in a layer on a single training case, so it performs exactly the same computation at training and test time and is straightforward to apply to recurrent networks by computing the statistics separately at each time step. Slide 35 lists both facts. Because it never looks at another sample, it also works at batch size one, and the Transformer block in part 06 uses it in every Add and Norm step: Vaswani and colleagues define each sub-layer output as LayerNorm(x + Sublayer(x)).

Instance norm came from style transfer. Ulyanov, Vedaldi and Lempitsky found that replacing batch norm with instance norm, applied both at training and testing, sharply improved stylization, because the per-image contrast statistics are exactly what a style transfer network needs to discard. PyTorch's InstanceNorm2d defaults to no affine parameters and no running statistics.

Group norm is the embedded and detection story. Wu and He show that batch norm's error increases rapidly when the batch size becomes smaller: a ResNet-50 on ImageNet at batch size 2 reaches 34.7 percent error with BN and 24.1 percent with GN, while GN's computation is independent of batch size. With G = 32 as the default, GN becomes LN when G = 1 and IN when G = C, which is what the explorer's two slider ends show: one channel per group is instance norm, all channels in one group is layer norm. When you fine-tune a model on a device with a memory budget that allows one or two samples per step, a batch-free normalization is the difference between a model that trains and one that does not. Wu and He note that dropping the batch size constraint can free considerably more memory, sixteen times or more.

Tom Yeh's batch norm worksheet runs a mini-batch of four samples with three features each passes through a linear layer, a ReLU, then batch norm. Filling in every cell is the single best preparation for the computational question on this topic, so here is every number, recomputed rather than copied from the heavily rounded original.

The layout hides two conventions. The mini-batch is written as columns x1 = (1, 0, 2), x2 = (0, 3, 1), x3 = (3, 1, 0), x4 = (0, 1, 2) with a row of ones appended, so the red 3 x 4 "Linear Layer" block is [W | b] with W = [[1, 0, 1], [1, 1, 0], [0, 2, -1]] and b = (0, -1, 0). Likewise the "Scale and Shift" block is [diag(gamma) | beta] with gamma = (2, 3, -1) and beta = (0, 0, 1), applied to the normalized matrix with its own row of ones. The small two-column "Trainable Parameters" box (drawn as a 2 x 2 grid in the blank template) stands for the columns [gamma | beta]; for this three-feature layer it holds a 3 x 2 matrix.

Switch the simulator to eval mode and the same four inputs produce different outputs with the simulator's placeholder running mean and variance of one (PyTorch itself starts running_mean at 0 and running_var at 1): sample 1, feature 1 gives 4.00 instead of 1.81. That gap is the slide 33 bug made visible. To reproduce the fusion table of the previous concept, switch the placement to the paper order, press "load this batch's statistics" so the running mean becomes 2.25, 1.25, 1.25 and the running variance 0.6875, 1.6875, 6.6875, then press fuse: W' and b' appear computed live, with the check that they reproduce BN(Wx + b). With the placeholder running values of one you get a different, equally valid fused layer with scales 2, 3, -1, because the fused weights depend on whatever statistics training froze.