{"slug": "batch-layer-and-group-norm-side-by-side-on-mnist", "title": "Batch, layer, and group norm, side by side on MNIST", "summary": "A developer implemented batch normalization, layer normalization, and group normalization from scratch and compared their performance on a simple multi-layer perceptron (MLP) classifying the MNIST dataset, finding that each method stabilizes training but makes different assumptions about data distribution. The code, available in a GitHub notebook, uses a three-hidden-layer network with 128 activations per layer and ReLU non-linearity, and the author discusses the underlying assumptions of each normalization technique.", "body_md": "The code below is trimmed to the essentials; the full, runnable version lives in the original notebook:\n\ngithub.com/litlig/notebooks/normalization.ipynb\n\nNormalization 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?\n\nHere 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.\n\n## Set up[#](#set-up)\n\nWe use a basic fully-connected neural network to classify digits from MNIST:\n\n**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.\n\nThe layer and network are kept generic so we can swap in a normalized layer type without touching the rest of the code:\n\n``` python\nclass MLPLayer(nn.Module):\n  def __init__(self, input_dim, hidden_dim, activation):\n    super().__init__()\n    self.linear = nn.Linear(input_dim, hidden_dim)\n    self.activation = activation()\n\n  def forward(self, x):\n    return self.activation(self.linear(x))\n\nclass SimpleNN(nn.Module):\n  def __init__(self, n_layer, layerNN, activationNN=nn.ReLU):\n    super().__init__()\n    self.flatten = nn.Flatten()\n    self.mlp = nn.Sequential()\n    self.mlp.append(layerNN(input_dim, dim_hidden, activationNN))\n    for _ in range(n_layer - 1):\n      self.mlp.append(layerNN(dim_hidden, dim_hidden, activationNN))\n    self.mlp.append(nn.Linear(dim_hidden, 10))\n\n  def forward(self, x):\n    return self.mlp(self.flatten(x))\n```\n\n## Batch normalization[#](#batch-normalization)\n\nDeep 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.\n\nBatch normalization brings the covariates back to a favorable domain by shifting and scaling:\n\n\\[\\widehat{x} = \\frac{x - \\mu}{\\sigma}\\]where \\(\\mu\\) is the mean and \\(\\sigma\\) is the standard deviation over the mini-batch.\n\nIn 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\\).\n\nNote 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.\n\n``` python\nclass BatchNormMLP(nn.Module):\n  def forward(self, x):\n    x = self.linear(x)  # B, hidden_dim\n    if self.training:\n      mean = x.mean(dim=0)                    # over the batch\n      var = x.var(dim=0, unbiased=False)\n      # update running stats for eval time ...\n      x = (x - mean) / torch.sqrt(var + eps)\n    else:\n      x = (x - self.running_mean) / torch.sqrt(self.running_var + eps)\n    x = x * self.gamma + self.beta\n    return self.activation(x)\n```\n\n## Layer normalization[#](#layer-normalization)\n\nBatch 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:\n\n\\[\\widehat{x_i} = \\frac{x_i - \\mu_i}{\\sigma_i}\\]Here \\(x_i\\) is a sample vector with \\(d\\) dimensions. For each sample it:\n\n- projects \\(x\\) onto the \\((d-1)\\)-dimensional subspace where \\(\\sum_j x_{ij} = 0\\);\n- divides by \\(\\sigma\\) to fix the norm to 1.\n\nUnlike 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.\n\nLayer 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.\n\n``` python\nclass LayerNormMLP(nn.Module):\n  def forward(self, x):\n    x = self.linear(x)                          # B, hidden_dim\n    mean = x.mean(dim=1, keepdim=True)          # over features, per sample\n    var = x.var(dim=1, unbiased=False, keepdim=True)\n    x = (x - mean) / torch.sqrt(var + eps)\n    x = x * self.gamma + self.beta\n    return self.activation(x)\n```\n\n## Group normalization[#](#group-normalization)\n\nGroup 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.\n\nGroup 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.\n\n``` python\nclass GroupNormMLP(nn.Module):\n  def forward(self, x):\n    B = x.shape[0]\n    x = self.linear(x).view(B, self.n_group, self.hidden_dim // self.n_group)\n    mean = x.mean(dim=2, keepdim=True)          # within each group\n    var = x.var(dim=2, unbiased=False, keepdim=True)\n    x = (x - mean) / torch.sqrt(var + eps)\n    x = x.view(B, self.hidden_dim)\n    x = x * self.gamma + self.beta\n    return self.activation(x)\n```\n\n## Results across all models[#](#results-across-all-models)\n\nWe compare the four models — one vanilla MLP and three MLPs with different normalizations — on training speed, test-set accuracy, and neuron activation.\n\nFor 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.\n\nAll three normalized models train faster and reach better accuracy than the vanilla MLP, but there is no significant difference in training speed among themselves:\n\n| Model | Test accuracy (%) |\n|---|---|\n| MLP | 84.09 |\n| BatchNorm MLP | 96.58 |\n| LayerNorm MLP | 95.37 |\n| GroupNorm MLP | 96.33 |\n\nThere 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.", "url": "https://wpnews.pro/news/batch-layer-and-group-norm-side-by-side-on-mnist", "canonical_source": "https://blog.strayforge.com/posts/normalization-notebook/", "published_at": "2026-07-31 00:00:00+00:00", "updated_at": "2026-07-31 18:26:25.476594+00:00", "lang": "en", "topics": ["machine-learning", "neural-networks"], "entities": ["MNIST", "BatchNorm", "LayerNorm", "GroupNorm", "GitHub", "PyTorch"], "alternates": {"html": "https://wpnews.pro/news/batch-layer-and-group-norm-side-by-side-on-mnist", "markdown": "https://wpnews.pro/news/batch-layer-and-group-norm-side-by-side-on-mnist.md", "text": "https://wpnews.pro/news/batch-layer-and-group-norm-side-by-side-on-mnist.txt", "jsonld": "https://wpnews.pro/news/batch-layer-and-group-norm-side-by-side-on-mnist.jsonld"}}