Batch Norm vs Layer Norm vs Group Norm, Compared on MNIST

The code below is trimmed to the essentials; the full, runnable version lives in the original notebook: github.com/litlig/notebooks/normalization.ipynb

Normalization is a ubiquitous part of modern transformer architectures. I frequently see different normalization algorithms (BatchNorm, LayerNorm) in neural net definitions. Even though I know it’s there to improve training performance and avoid issues like covariate shift and non-linearity saturation, I always wondered: what are the assumptions behind these different approaches? How can such a transformation not lose input information? And what scenarios call for one approach over another?

Here I wrote a from-scratch implementation of the three — batch norm, layer norm, and group norm — applied them to a simple MLP on the MNIST dataset, and discuss the underlying assumptions as I understand them.

Batch norm vs layer norm vs group norm at a glance

#

The rest of the post builds each one from scratch and explains where these differences come from.

Setup: an MLP on MNIST

#

We use a basic fully-connected neural network to classify digits from MNIST:

• Input: 28×28 image, flattened to 784 features.
• Hidden layers: 3 fully-connected layers, each with 128 activations and ReLU non-linearity.
• Output layer: fully-connected layer with 10 activations (one per class) and softmax for classification.
• Loss: categorical cross-entropy.

The layer and network are kept generic so we can swap in a normalized layer type without touching the rest of the code:

class MLPLayer(nn.Module):
  def __init__(self, input_dim, hidden_dim, activation):
    super().__init__()
    self.linear = nn.Linear(input_dim, hidden_dim)
    self.activation = activation()

  def forward(self, x):
    return self.activation(self.linear(x))


class SimpleNN(nn.Module):
  def __init__(self, n_layer, layerNN, activationNN=nn.ReLU):
    super().__init__()
    self.flatten = nn.Flatten()
    self.mlp = nn.Sequential()
    self.mlp.append(layerNN(input_dim, dim_hidden, activationNN))
    for _ in range(n_layer - 1):
      self.mlp.append(layerNN(dim_hidden, dim_hidden, activationNN))
    self.mlp.append(nn.Linear(dim_hidden, 10))

  def forward(self, x):
    return self.mlp(self.flatten(x))

Batch normalization

#

Reference paper

Deep NN training often runs into a problem called covariate shift together with non-linearity saturation. During training, the covariates passed to the non-linear activation function can drift into the saturation regime, leading to vanishing gradients and slow training.

Batch normalization brings the covariates back to a favorable domain by shifting and scaling:

where \(\mu\) is the mean and \(\sigma\) is the standard deviation over the mini-batch.

In neural nets there is always a linear layer aggregating information across the hidden dimension before it is fed to the activation, \(z = Wx + b\). Scaling and shifting \(x\) does not lose information, because it can be absorbed by \(W\) and \(b\). However, the net adjusts \(W\) and \(b\) by gradient descent for the task; stabilizing the covariate distribution is not part of its objective, and this often leads to poor training performance. Batch norm keeps the covariate statistics stable by construction. To restore the scale and shift that the net actually needs, it introduces trainable parameters, so the input to the activation becomes \(\gamma \widehat{x} + \beta\).

Note that batch norm assumes the mini-batch statistics represent the overall population fairly well, so it often performs poorly when the batch size is too small.

class BatchNormMLP(nn.Module):
  def forward(self, x):
    x = self.linear(x)  # B, hidden_dim
    if self.training:
      mean = x.mean(dim=0)                    # over the batch
      var = x.var(dim=0, unbiased=False)
      # update running stats for eval time ...
      x = (x - mean) / torch.sqrt(var + eps)
    else:
      x = (x - self.running_mean) / torch.sqrt(self.running_var + eps)
    x = x * self.gamma + self.beta
    return self.activation(x)

Layer normalization

#

Reference paper

Batch norm normalizes across batch samples for each neuron in a layer; layer norm instead treats each sample individually and normalizes across all neurons for that sample:

Here \(x_i\) is a sample vector with \(d\) dimensions. For each sample it:

1. projects \(x\) onto the \((d-1)\)-dimensional subspace where \(\sum_j x_{ij} = 0\);
2. divides by \(\sigma\) to fix the norm to 1.

Unlike batch norm, the layer norm transformation is not fully recoverable by a linear transformation nor by the affine step. It reduces the feature dimension by 2: as the example below shows, the 3-D \(x\) is now confined to a circle.

Layer norm assumes that the (centered) direction of a sample carries the information, and that the mean and scale are redundant. That is a genuine inductive bias, but in practice it is often a good bet.

class LayerNormMLP(nn.Module):
  def forward(self, x):
    x = self.linear(x)                          # B, hidden_dim
    mean = x.mean(dim=1, keepdim=True)          # over features, per sample
    var = x.var(dim=1, unbiased=False, keepdim=True)
    x = (x - mean) / torch.sqrt(var + eps)
    x = x * self.gamma + self.beta
    return self.activation(x)

Group normalization

#

Reference paper

Group norm takes a step further from layer norm: instead of normalizing across all features in a layer, it splits the features into groups and normalizes each group by its own mean and standard deviation. The idea is that different groups can have genuinely different statistics, so they should not be pooled together when computing the mean and variance.

Group norm further reduces the degrees of freedom to \(d - 2g\), where \(d\) is the total feature dimension and \(g\) is the number of groups, in exchange for not letting heterogeneous groups pollute each other’s statistics. It was introduced for vision models with small batch sizes: channel groups (e.g. from grouped convolutions, or filters tuned to different frequencies or orientations) can have genuinely different natural scales, and pooling all of them into one LayerNorm-style mean/variance would wash out that structure.

class GroupNormMLP(nn.Module):
  def forward(self, x):
    B = x.shape[0]
    x = self.linear(x).view(B, self.n_group, self.hidden_dim // self.n_group)
    mean = x.mean(dim=2, keepdim=True)          # within each group
    var = x.var(dim=2, unbiased=False, keepdim=True)
    x = (x - mean) / torch.sqrt(var + eps)
    x = x.view(B, self.hidden_dim)
    x = x * self.gamma + self.beta
    return self.activation(x)

Results: BatchNorm vs LayerNorm vs GroupNorm

#

We compare the four models — one vanilla MLP and three MLPs with different normalizations — on training speed, test-set accuracy, and neuron activation.

For neuron activation, the big white chunks for the vanilla MLP mean there are neurons that never activate for any sample in the batch. With normalization the chart shows a more mixed black-and-white texture: each neuron fires for some samples and not others, so its output is actually input-dependent — which is what we want if the network is going to use that neuron to help separate digit classes.

All three normalized models train faster and reach better accuracy than the vanilla MLP, but there is no significant difference in training speed among themselves:

There is no significant difference among the three normalizations here because the example is simple enough that all of them perform well. There is no convolutional layer for feature extraction, so group norm has no advantage over layer norm in this setting.

When to use which

#

• Batch norm when batches are large and fixed-size, as in classic CNN image classification. It is the only one of the three that uses information across samples, and the noise in its batch statistics acts as a mild regularizer. Avoid it with small or variable batches, and remember that training and inference compute different things.
• Layer norm when each sample has to be normalized on its own: sequence models, small per-device batches, or inference on one example at a time. It is the default in transformers.
• Group norm when features come in groups with different natural scales (typically convolutional channels) and the batch is too small for batch norm, e.g. detection and segmentation with a few high-resolution images per GPU. With a single group it reduces to layer norm; with one channel per group it becomes instance norm.

Why do transformers use layer norm instead of batch norm?

#

Mostly because batch statistics are a poor fit for sequences:

1. Variable lengths and padding. Sequences in a batch have different lengths, so batch statistics would mix real tokens with padding and change with how each batch happens to be composed.
2. Small effective batches. Large models often fit only a few sequences per device, which makes batch statistics noisy, the same weakness noted above.
3. Generation one token at a time. Autoregressive inference often runs with a batch size of 1, where batch norm has to fall back on running averages that may not match what it saw in training. Layer norm computes exactly the same thing in training and inference.

Many recent LLMs go one step further and use RMSNorm, which drops the mean-centering and only rescales each sample by its root mean square.