🎓 A1Learner
🔍

PyTorch Batch Normalization

Batch normalization is a layer that rescales each feature to mean 0 and std 1 across the batch. The rescaling happens inside the network, on every pass.

The drift, before the fix

The plain stack in the last lesson let the std collapse. Batch normalization enters as the fix for that table.

h = torch.relu(linear(h))  # std fell below 0.10 by layer 3

One batch, rescaled per column

The three columns here start at very different scales. After the layer, every column has mean 0 and std 1, measured the way the layer computes it.

import torch

torch.manual_seed(0)
x = torch.randn(4, 3) * torch.tensor([1.0, 10.0, 100.0])
norm = torch.nn.BatchNorm1d(3)
norm.train()
out = norm(x)
print("Column means:", [round(v, 4) for v in out.mean(dim=0).tolist()])
print("Column stds:", [round(v, 4) for v in out.std(dim=0, unbiased=False).tolist()])

Running stats build up

Each training pass nudges the stored running stats toward the batch stats. The momentum setting controls how big each nudge is.

import torch

torch.manual_seed(0)
norm = torch.nn.BatchNorm1d(3)
x = torch.randn(4, 3) * 10
norm.train()
norm(x)
print("Running mean:", [round(v, 2) for v in norm.running_mean.tolist()])
print("Running var:", [round(v, 2) for v in norm.running_var.tolist()])

Example

import torch

torch.manual_seed(0)
x = torch.randn(64, 16)
pairs = [(torch.nn.Linear(16, 16), torch.nn.BatchNorm1d(16))
         for _ in range(10)]

print(f"Input: mean {x.mean():.4f}, std {x.std():.4f}")
h = x
for i, (linear, norm) in enumerate(pairs):
    h = torch.relu(norm(linear(h)))
    print(f"Layer {i + 1}: mean {h.mean():.4f}, std {h.std():.4f}")

Key points

  • Batch: the group of rows passed through the model together in one step
  • Feature: one column of the activations, tracked separately from the other columns
  • Use torch.nn.BatchNorm1d(features) to normalise a 2D batch, one column at a time
  • Running stats: the layer's stored running averages of mean and variance, updated on each training pass
  • In training mode the layer uses the current batch's own mean and std
  • In eval mode the layer uses the running stats, so one row gives the same answer in any company
  • Batch normalization pins the scale per feature, so the drift from the last lesson mostly disappears

In this example, the same 10-layer stack from the last lesson gains a BatchNorm1d after each Linear layer. The input and the seed match the earlier run. The std no longer collapses toward zero. Every layer after ReLU sits between 0.5455 and 0.6262, close to the first layer's 0.5861. The mean holds near 0.39 at every depth. The values after ReLU are not std 1, because ReLU clips the negative half after the layer rescales. The point is the steadiness. Depth no longer decides the scale.

Note

BatchNorm1d expects a (batch, features) input, and a training batch of one row cannot give a batch std.

Try it Yourself

Edit the code, then press Run.

Running is simulated · Python Playground

Exercise

Which values does BatchNorm1d use to rescale a batch in eval mode?