Normalisation Layers
Batch norm normalises each feature over the batch; layer norm normalises each example over its features. The difference decides whether the layer behaves the same at training and inference — and batch norm does not, which makes it a train/serve skew source with a name.
The problem, the obvious approach, and why it breaks
Every lesson starts where the work starts: someone has a problem, and the first model that comes to mind looks fine offline.
The model evaluated well in the training framework and behaves differently in the serving container, with identical weights. Which layer in the network has a different definition at inference time?
A visual-search team's image encoder produces slightly different embeddings in production than in evaluation for the same images. The weights are byte-identical. The gap is small per image and large enough to shuffle nearest-neighbour results, and it got worse after a change to how requests are batched at serving time.
Normalisation is a training trick that keeps activations in a healthy range so the network trains faster. Put a batch-norm layer after every convolution, as the reference architectures do, and treat it as part of the architecture.
At training time batch norm uses the current batch's statistics; at inference it uses running averages accumulated during training. The layer computes a different function in the two modes. If serving forgets to switch modes, each request is normalised by its own tiny batch, and the output depends on which other images happened to arrive at the same time.
- At training time batch norm uses the current batch's statistics; at inference it uses running averages accumulated during training. The layer computes a different function in the two modes. If serving forgets to switch modes, each request is normalised by its own tiny batch, and the output depends on which other images happened to arrive at the same time.
- When serving does switch modes correctly, the running statistics are the ones accumulated on the training distribution. Production images with a different colour profile — a new camera, a new upload pipeline — are normalised by the wrong mean and variance, and every downstream layer sees shifted inputs. Identical weights, different function: Train / Serve Skew built into the architecture.
- The serving change that made it worse was a switch from single-request inference to dynamic batching; a build that had been accidentally running in training mode now normalised over a batch of unrelated users' images and produced embeddings that varied with traffic.
What is being predicted, and from what data
This domain leads with these two. A target nobody defined precisely is a label nobody can trust, and a dataset nobody can describe is a model nobody can debug.
- Keep the forward pass numerically the same function of the input at training time and at serving time; a layer whose output depends on which other examples are in the batch violates that by construction.
- The surrounding system embeds images for retrieval; its target — neighbour quality — depends on the embedding being a function of the image alone (Embeddings).
- Training batches of images drawn at random, so the per-channel mean and variance over a batch are stable estimates of the population statistics.
- Serving requests that arrive one at a time or in small batches of whatever came together, so a batch statistic computed at serving time would be a statistic of one user's images, not of the population.
How it actually works
Precisely enough to predict its behaviour — not a framework API.
- Both layers subtract a mean and divide by a standard deviation, then apply a learned scale and shift. They differ in what the statistics are computed over. Batch norm computes them per feature (per channel) across the examples in the batch; layer norm computes them per example across the features. Batch norm's output for one example therefore depends on the other examples; layer norm's does not.
- Because batch norm needs a batch, it cannot run the same way on a single example at inference. The standard solution is to keep an exponential running average of the batch means and variances during training and use those fixed numbers at inference. That gives a deterministic, per-example function — but a different function from the one trained, and one that encodes the training distribution's statistics as frozen constants inside the model.
- Layer norm has no such split. Its statistics are computed from the example itself, so the layer is identical in both modes and independent of batch size. That, together with variable sequence lengths that make per-position batch statistics ill-defined, is why transformers use layer norm rather than batch norm (Transformer Fundamentals).
- Why normalisation helps training at all is less settled than its practice: it keeps pre-activations in a range where derivatives are healthy (Vanishing and Exploding Gradients), makes the loss surface smoother so a larger learning rate is stable, and reduces the sensitivity of each layer to changes in the ones before it. The practical effect — faster, more stable training of deep stacks — is robust; the explanation is contested.
Over the batch, or over the features
The two layers look identical in code except for one axis argument, and that axis decides whether the layer is a function of one example or of the batch. The consequence follows immediately: a function of the batch cannot be evaluated the same way on a batch of one, so batch norm needs a second definition for inference and layer norm does not.
1def batch_norm(x, gamma, beta, state, training, eps=1e-5, momentum=0.1):2 # x: [batch, features]; statistics per feature, across the batch3 if training:4 mu, var = x.mean(axis=0), x.var(axis=0)5 state["mu"] = (1 - momentum) * state["mu"] + momentum * mu # running stats6 state["var"] = (1 - momentum) * state["var"] + momentum * var # saved with the weights7 else:8 mu, var = state["mu"], state["var"] # frozen: the training distribution9 return gamma * (x - mu) / (var + eps) ** 0.5 + beta10 11def layer_norm(x, gamma, beta, eps=1e-5):12 # statistics per example, across features; no mode, no state, no batch dependence13 mu, var = x.mean(axis=1, keepdims=True), x.var(axis=1, keepdims=True)14 return gamma * (x - mu) / (var + eps) ** 0.5 + betastate["mu"] and state["var"] are part of the model and must ship with it. The training flag is a serving contract: get it wrong and every request is normalised by its neighbours in the batch.
A skew source with a name
The serving lesson describes skew as a difference between what the model saw in training and what it receives in production (Train / Serve Skew). Batch norm produces it inside the model: the same input tensor is normalised by batch statistics in training and by frozen population statistics at inference, and if the population moves, the frozen statistics are wrong for every request after that.
Embeddings computed in the training framework in evaluation mode match the checkpointed run exactly; retrieval quality on the held-out set is unchanged from the previous model.
Embeddings for the same images differ slightly from the offline ones; nearest-neighbour results shuffle; after the serving batcher change, the same image returns different neighbours depending on the traffic it arrived with.
- 1The serving build did not switch the normalisation layers to inference mode, so each request was normalised by its own batch — a batch of one before the batcher change, a batch of strangers after it.
- 2Even with the mode correct, the running statistics were accumulated on the training distribution; images from a newer upload pipeline with a different colour profile are normalised by the wrong constants.
- 3A fine-tune on a small internal set had updated the running statistics towards that set without meaningfully changing the weights, so "the same weights" was almost true and the model was not the same.
What the frozen statistics assume
The running mean and variance are the training distribution, compressed to two numbers per channel and welded into the model. Every assumption the domain makes about feature distributions staying put applies to them, with the twist that they are not listed as features anywhere and no feature monitor will watch them.
The per-channel statistics of the pre-normalisation activations at serving time match the running statistics frozen into the artifact, and every normalisation layer runs in inference mode.
holds when The serving path sets inference mode and loads the statistics saved with these weights; the input distribution — camera, preprocessing, resolution, colour handling — is the one training saw.
breaks when A serving build runs in training mode; a fine-tune updates the statistics; a new input source changes the channel statistics; a preprocessing step upstream is changed "harmlessly".
respond Fix the mode or the loaded statistics first. If the input population has genuinely moved, recalibrate the statistics on the new distribution or retrain; do not adjust the weights to compensate for a normalisation mismatch.
How to build it
Most important first.
- Treat the train/inference mode switch as a serving contract, tested: a request replayed through the training framework in evaluation mode and through the serving path must produce the same output (Serving Contract Tests).
- Treat batch norm's running statistics as part of the artifact and version them with the weights (What a Model Artifact Contains); they are learned constants and a model without them is a different model.
- Prefer layer norm, or a batch-independent variant, wherever inference batches are small, variable or made of unrelated requests, and wherever the input distribution is expected to move.
- Monitor the per-channel mean and variance of pre-normalisation activations at serving time against the frozen running statistics; a drift between them is a drift in the input distribution the layer was calibrated for (Feature Drift).
What to measure
Which number actually maps to the decision — and which numbers look relevant and are not.
- Output difference between the training framework in evaluation mode and the serving path, on a replayed sample. Zero, to numerical tolerance, is the only acceptable value; a nonzero value locates a mode or statistics mismatch.
- Distance between the serving-time activation statistics and the stored running statistics, per batch-norm layer, over time — the early-warning signal for input drift through this specific mechanism.
- Do not measure this with an accuracy number on a held-out set evaluated in the training framework. That path never exercises the serving mode and cannot see the gap.
What must stay true after deployment
The field this whole domain exists for. A model is a set of assumptions with weights attached; these are the ones a monitor or a test should be checking.
- The serving path runs every normalisation layer in inference mode, and the running statistics loaded are the ones saved with these weights.
- The input distribution at serving time has per-channel statistics close to the frozen running statistics; a camera, pipeline or preprocessing change that moves them changes the model's function without touching a weight.
- The batch size at serving time does not affect the output, which is true only if no layer computes statistics across the batch.
- Offline: a contract test that replays a fixed sample through the training framework in evaluation mode and through the serving container, and asserts equal outputs to a tolerance; a second test that asserts the serving output for one request is unchanged by the other requests in its batch.
- Online: log pre-normalisation activation statistics per layer at serving time on a sample of traffic and compare against the stored running statistics; alert on a sustained shift.
- Over time: when the model is fine-tuned or retrained, re-run both tests and record the running statistics alongside the weights in the registry (The Model Registry).
What can go wrong
- The running statistics were accumulated over the last few batches of training with a momentum that never converged, so they are noisy estimates of the training distribution and the evaluation-mode model is slightly wrong even on training data.
- Fine-tuning on a small dataset updates the running statistics towards the small set's distribution while the weights barely move; the fine-tuned model behaves differently on the original distribution for a reason nobody edited.
- A distributed training run computes batch statistics per device, so the effective batch for normalisation is the per-device micro-batch, and the model is silently trained with much noisier normalisation than the global batch size suggests (Data Parallelism).
- Layer norm avoids the mode split and the batch dependence, but for convolutional models batch norm often trains better, and switching costs a retraining and a re-tuning of the recipe.
- Monitoring activation statistics at serving time is another set of counters per layer, on the latency-critical path.
- Treating running statistics as part of the artifact is the right model and it means every fine-tune produces a new artifact even when the weights are frozen — more versions to manage.
- "The weights are identical, so the model is identical." A batch-norm model is its weights plus its running statistics plus its mode flag. Change any of the three and the function changes.
- "Batch norm is just a training trick and has no effect at inference." It has a permanent effect at inference: a frozen normalisation by the training distribution's statistics, which is the mechanism by which input drift reaches every layer after it.
- "We use layer norm, so we have no train/serve skew." Layer norm removes this one source. The feature pipeline, the tokeniser and the preprocessing are still separate implementations with all the usual skew (Train / Serve Skew).
Where this applies
ML advice is stated as universal far more often than it is. These labels say what each claim is specific to — and where CONTESTED appears, the note gives the strongest form of the opposing view.
- MODEL-SPECIFICThe mode-split failure is specific to batch norm and other layers that compute statistics across examples; layer norm, group norm and RMS norm compute per example and do not have it, which is one reason transformers standardised on per-example normalisation.
- CONTESTEDWhy normalisation helps optimisation is genuinely unsettled. The original account — reducing "internal covariate shift" — has been challenged by work showing the benefit persists when the shift is deliberately reintroduced, with a smoother loss surface and permission for a larger learning rate offered instead. The lesson states the robust practical effect and leaves the mechanism open, because the engineering consequences do not depend on which explanation is right.
Where the depth lives
This domain teaches the model and hands the rest off by name.
- — Testing & Reliability Engineering — an equivalence test between the training framework and the serving container is a contract test on the model itself, and keeping it green across fine-tunes is a testing discipline this domain assumes rather than teaches.