Normalization layers

At 2 a.m., your image classifier has a training loss that falls beautifully for 400 steps. Then it becomes nan. The first layers produce activations around 0.3; twenty layers later, some values are 80,000. A different model with the same data, optimizer, and learning rate trains normally.

Or your language model trains, but only when every GPU gets exactly 8,192 tokens. Change the sequence lengths and the loss starts lurching. Run validation with one example instead of a full batch and accuracy falls from 91 percent to 12 percent.

These can be symptoms of poorly controlled activation statistics, but they are not specific to normalization. Also check:

• the learning rate,
• loss scaling,
• masks,
• precision, and
• train/evaluation mode.

A normalization layer rescales activations using statistics of those activations. The aim is mostly to make optimization better behaved: gradients become less sensitive to arbitrary changes in scale, and the loss surface often becomes easier for an optimizer to navigate.

The axis over which you compute the statistics decides almost everything.

The basic operation

Suppose one neuron receives four activation values:

[2, 4, 6, 8]

Their mean is 5. Their population variance is 5, so their standard deviation is approximately 2.236. Standardizing them gives:

[(2 - 5) / 2.236, (4 - 5) / 2.236, (6 - 5) / 2.236, (8 - 5) / 2.236]

That is approximately:

[-1.342, -0.447, 0.447, 1.342]

The values now have mean near zero and standard deviation near one. A normalization layer usually applies a learned scale and shift:

output = gamma * normalized_value + beta

gamma and beta let the network restore a useful scale or offset. Normalization constrains the signal during optimization, not the final representation.

Why it helps

Why does this help? If an earlier weight change makes activations ten times larger, a downstream sigmoid can saturate near zero or one, where its gradient is tiny.

A softmax can become nearly one-hot before the model has learned useful boundaries. Keeping activation scales controlled makes these failures less likely. This connects directly to vanishing and exploding gradients.

Normalization can still change generalization, because optimization affects generalization. But it does not replace dropout or weight decay.

BatchNorm: effective, but batch-dependent

Batch normalization computes statistics across examples in a minibatch. For a convolutional tensor shaped [batch, channels, height, width], it normally computes one mean and variance per channel across:

• batch positions,
• height positions, and
• width positions.

The original explanation focused on internal covariate shift: earlier layers change the distribution received by later layers. Thus, normalizing each layer’s inputs should help.

That is useful intuition, but not a settled causal explanation. BatchNorm’s main benefit is better optimization geometry: it often makes loss and gradients smoother with respect to parameter changes, so larger learning rates become practical and descent is less brittle.

The benefit depends on having useful statistics and matching training with deployment.

Training is not evaluation

BatchNorm has two modes:

• During training, it uses the current minibatch’s mean and variance and updates running estimates.
• During evaluation, it uses those running estimates instead of the current batch.

That makes single-example inference stable. Without it, an image’s prediction could depend on which other images happened to be processed beside it.

import torch

from torch import nn

torch.manual_seed(0)

# Two features, four examples.
batch = torch.tensor([
    [1.0, 10.0],
    [3.0, 14.0],
    [5.0, 18.0],
    [7.0, 22.0],
])

bn = nn.BatchNorm1d(num_features=2)

bn.train()
train_output = bn(batch)
print("training batch means:", train_output.mean(dim=0))
print("running means:", bn.running_mean)

bn.eval()
eval_output = bn(batch)
print("evaluation output:", eval_output)

The training output has feature means close to zero because it uses this batch. The evaluation output differs because eval() switches to the running means and variances.

With only one small batch, those estimates are still close to initialization, so the difference is especially visible.

Use the corresponding modes explicitly:

model.train()
loss = model(inputs)

model.eval()
with torch.no_grad():
    predictions = model(inputs)

model.eval() also changes dropout. It does not disable gradients by itself, which is why inference commonly uses torch.no_grad() too.

Small batches and sequence lengths

Forgetting model.eval() can make validation output change with validation batch size. BatchNorm is still using each batch’s statistics, and dropout may still be active.

Small batches make BatchNorm estimates noisy. A batch of 1 can fail for some shapes because there are not enough values to estimate statistics. Gradient accumulation does not fix this: it accumulates gradients across forward passes, while ordinary BatchNorm still computes statistics separately in each small pass.

Possible fixes include:

• a larger true per-forward-pass batch,
• synchronized statistics across devices,
• frozen BatchNorm, or
• a different normalization method.

Variable-length sequences are another mismatch: padding and time positions can contaminate statistics unless the operation is carefully masked. This is why BatchNorm is uncommon inside transformer blocks, though it remains effective for many convolutional vision models with reasonably large batches.

LayerNorm: statistics from one example

Layer normalization computes statistics within each example, usually across the feature dimension.

A transformer hidden state might have shape [batch, sequence length, hidden size], such as [32, 512, 768]. LayerNorm(768) normalizes the 768 features of each token separately. It does not compare tokens across examples.

Given the same vector and learned parameters, LayerNorm behaves the same in training and evaluation. It has no running means or batch-size dependence. Calling eval() does not change its normalization rule.

That suits transformers: batches may contain one sequence or many, sequence lengths vary, and each token has a feature vector to normalize. For [batch, time, hidden], LayerNorm(hidden) is normally correct.

Normalizing across time would make a token depend on the rest and length of the sequence. If the dimensions happen to match, this can fail silently.

LayerNorm is not free. Its feature reduction occurs in every transformer block, and removing each token’s mean and controlling its scale may be too strong for an architecture that relies on precise magnitude information.

RMSNorm: control scale without centering

RMSNorm uses the root mean square:

RMS(x) = sqrt(mean(x_i^2) + epsilon)

It divides by that value but does not subtract the mean. For [2, 4, 6, 8], the mean squared value is 30, so the RMS is approximately 5.477; the normalized values are approximately:

[0.365, 0.730, 1.095, 1.461]

The positive offset remains. RMSNorm controls vector scale while preserving its mean direction.

In many residual transformer stacks, controlling activation magnitude matters more than exact centering. RMSNorm therefore offers a simple, stable per-example operation.

It is deterministic like LayerNorm and has no running statistics. The trade-off is that architectures needing centered activations may prefer LayerNorm.

import torch

from torch import nn


class RMSNorm(nn.Module):
    def __init__(self, hidden_size, eps=1e-8):
        super().__init__()
        self.weight = nn.Parameter(torch.ones(hidden_size))
        self.eps = eps

    def forward(self, x):
        mean_square = x.pow(2).mean(dim=-1, keepdim=True)
        inverse_rms = torch.rsqrt(mean_square + self.eps)
        return x * inverse_rms * self.weight


x = torch.tensor([
    [[2.0, 4.0, 6.0, 8.0]],
])  # shape: batch 1, time 1, hidden 4

norm = RMSNorm(hidden_size=4)
print(norm(x))

GroupNorm: vision without a large batch

Group normalization divides the channels of each image into groups and computes statistics within each group, including spatial positions. With 32 channels and 8 groups, each group contains 4 channels. One image’s statistics cannot be changed by another image.

This makes GroupNorm a useful small-batch vision choice:

• BatchNorm uses other examples and spatial positions.
• LayerNorm uses one example across features at a location.
• GroupNorm uses one example across groups of channels and spatial positions.

It is common in detection, segmentation, and high-resolution vision tasks where a batch of 1–8 is unavoidable.

The group count is a hyperparameter: too few groups approaches LayerNorm over many channels, while too many make groups narrow. GroupNorm is designed for convolutional feature maps, not sequence normalization.

Pre-norm and post-norm residual blocks

A residual block transforms x with a sublayer F and adds the original input back. Normalization can go before or after that operation:

• Post-norm: output = Norm(x + F(x))
• Pre-norm: output = x + F(Norm(x))

In pre-norm, the residual branch leaves an explicit identity path from input to output. During backpropagation, part of the gradient can travel through that addition without passing through every normalization Jacobian. This usually makes very deep stacks easier to optimize.

In post-norm, every block’s output passes through normalization. It can work with suitable initialization and learning-rate schedules, but deep transformers are generally more sensitive during early training.

Pre-norm with LayerNorm or RMSNorm is common in deep transformer language models. A final normalization before the output head is also common because the residual stream is not normalized after every addition.

Pre-norm is not always superior: it can produce less tightly controlled final representations, and some recipes get better quality from post-norm. For a deep, fragile stack, it is usually the safer stability-first choice.

Choosing a normalization layer

Situation	Usually choose	Why it fits	Main cost or risk
Convolutional vision, batch of 32 or more	BatchNorm	Batch statistics are reliable and channel-wise scaling fits feature maps	Train/eval asymmetry and batch dependence
Convolutional vision, batch of 1 to 8	GroupNorm	Each image supplies its own statistics	Group count needs tuning
Transformer or variable-length sequence	LayerNorm	Each token is normalized independently across hidden features	More reduction work; removes mean information
Deep transformer where scale control matters most	RMSNorm	Per-example scale normalization is stable and simple	Does not center activations
Deliberate generalization pressure	Dropout or weight decay, alongside a suitable norm	These target overfitting or parameter complexity	Can hurt underfitting models

Inspect:

• the tensor shape,
• the batch size that fits in memory, and
• deployment conditions.

Do not treat the table as a rule.

Diagnosing failures

Validation collapses or changes with batch size. Check model.eval(). BatchNorm may still be using validation minibatch statistics, and dropout may still be active.

Loss becomes noisy with small batches. BatchNorm has too little data for reliable statistics. Gradient accumulation alone is not enough.

Use one of these responses:

• Increase the true forward-pass batch.
• Synchronize devices.
• Freeze BatchNorm.
• Choose GroupNorm for vision and LayerNorm/RMSNorm for sequences.

A language model depends on padding or sequence length. Generic BatchNorm may be mixing tokens and padding into its statistics. Normalize the hidden dimension with LayerNorm or RMSNorm.

A layer runs but performs poorly. Check the normalized axis. For [batch, time, hidden], it is usually hidden. Also verify that the checkpoint and code agree on pre-norm versus post-norm.

Deployment fails after export. BatchNorm’s running statistics may be stale, poorly estimated from tiny batches, or mismatched with a new camera, sensor, or population. Re-estimate or freeze them deliberately and test at the real inference batch size.

What normalization cannot do

Normalization cannot rescue:

• a wrong learning rate,
• bad labels,
• a broken loss, or
• an architecture with no useful gradient path.

It also cannot guarantee safe values everywhere. Separate failure modes include:

• learned gamma,
• residual accumulation,
• mixed-precision overflow,
• attention logits, and
• optimizer instability.

Monitor activation ranges and gradient norms when training is fragile.

The deeper limitation is that each method encodes an assumption:

• BatchNorm treats a minibatch as a useful statistical population.
• LayerNorm treats a token’s mean and scale as removable.
• RMSNorm preserves the mean but controls magnitude.
• GroupNorm assumes local channel groups are meaningful.

Those assumptions are often useful, but they remain assumptions.

What to remember

• Normalization changes activation scale and optimization geometry; it is not interchangeable with dropout or weight decay.
• BatchNorm uses minibatch statistics during training and saved running statistics during evaluation. Forgetting model.eval() is a classic bug.
• Small batches and variable-length sequences weaken BatchNorm.
• LayerNorm and RMSNorm normalize each example independently, which suits transformers. RMSNorm controls scale without subtracting the mean.
• GroupNorm suits small-batch vision, and pre-norm residual blocks usually make deep transformer stacks easier to train.

Quick check

0/3

Q1

BatchNorm has fewer learned parameters than LayerNorm.BatchNorm uses noisy minibatch statistics, while LayerNorm is deterministic for each example.LayerNorm always includes dropout internally.BatchNorm directly penalizes large weights.

Check answer

Q2

BatchNorm has unreliable small-batch statistics; try GroupNorm or increase the true per-forward-pass batch.LayerNorm is using too many features; replace it with dropout.RMSNorm is reading the wrong running mean; call model.eval() during training.The optimizer cannot work with convolutional layers; remove normalization.

Check answer

Q3

BatchNorm after each residual addition, because more documents always make token statistics reliable.LayerNorm or RMSNorm before each sublayer, because statistics stay per-example and the residual path remains easier to optimize.GroupNorm after each attention head, because heads are image channels.No normalization, because residual connections already guarantee stable gradients.

Check answer