Normalization Layers and Training Behavior

The chosen axes determine which examples or features shape one another.

The fourth card is no longer a research curiosity. Drop LayerNorm's re-centering step and the layer gets cheaper. Zhang and Sennrich tested that in 2019, starting from the hypothesis “that re-centering invariance in LayerNorm is dispensable”. What they found: “Extensive experiments on several tasks using diverse network architectures show that RMSNorm achieves comparable performance against LayerNorm but reduces the running time by 7%~64% on different models.”

PyTorch now ships the layer as torch.nn.RMSNorm, defined as y = x / RMS(x) * gamma with RMS(x) = sqrt(eps + mean(x^2)) and elementwise_affine=True. There is no mean-subtraction term anywhere in that formula. Meta runs it in production. Its own disclosure of LLaMA, in 2023, says: “To improve the training stability, we normalize the input of each transformer sub-layer, instead of normalizing the output. We use the RMSNorm normalizing function, introduced by Zhang and Sennrich (2019).”

That is what the card's “efficient transformer variants” actually names: a measured 7%–64% runtime reduction, a shipped framework layer, and a production model that normalizes this way.

A complete checkpoint must preserve both categories.

The two categories are not stored the same way. Gamma and beta are learned parameters that gradients update. The running mean and variance are buffers that no gradient touches. They are advanced by a fixed rule on each training forward pass, then substituted for the batch statistics at evaluation time. Restore the first category, reset the second, and you have loaded a model whose weights are right and whose evaluation behavior is not. Nothing in the training loss reports it.

These cases can change predictions without changing the learned matrix weights.

Use small, reproducible comparisons before changing architecture.

Step 4 is worth a number rather than a hope. Schneider and colleagues replaced the training-set BatchNorm statistics with statistics of the corrupted test images and reported: “Using the corrected statistics, ResNet-50 reaches 62.2% mCE on ImageNet-C compared to 76.7% without adaptation.” No retraining. No change to a single learned weight. The whole gain came from recomputing buffers, and a second group at Google Brain reproduced the effect independently.

Step 5 says architecture-specific because the same layer in a different position changes what the optimizer requires. Xiong and colleagues showed in 2020 that in Post-LN Transformers “the expected gradients of the parameters near the output layer are large”. That is why the warm-up stage is needed at all. In Pre-LN Transformers “the gradients are well-behaved at initialization”, and “Pre-LN Transformers without the warm-up stage can reach comparable results with baselines while requiring significantly less training time”.

Nguyen and Salazar had reported the same thing the year before: “First, we show that pre-norm residual connections (PRENORM) and smaller initializations enable warmup-free, validation-based training with large learning rates.” They measured an average +1.1 BLEU over bilingual baselines on five low-resource pairs, and 32.8 BLEU on IWSLT '15 English–Vietnamese. They also noted that in the high-resource WMT '14 English–German setting “PRENORM degrades performance”. The winner reverses with the setting. That is the whole reason step 5 is an experiment and not a recommendation.

Each forward pass computes BatchNorm statistics from the micro-batch present on that device unless synchronized normalization is used. Accumulating gradients later does not retroactively combine those activation statistics.

A global batch of 512 can therefore behave like BatchNorm batches of eight. Inspect the actual forward-pass group.

This is not folklore, and it is not an outsider's complaint. BatchNorm's own author conceded it in print two years after the original paper. Sergey Ioffe's 2017 paper on Batch Renormalization opens with the limit: “However, its effectiveness diminishes when the training minibatches are small, or do not consist of independent samples.” The cause it names is “the dependence of model layer inputs on all the examples in the minibatch, and different activations being produced between training and inference”. That is the micro-batch problem and the mode bug, in one sentence, from the person who introduced the layer.

Wu and He measured the size of it a year later: “On ResNet-50 trained in ImageNet, GN has 10.6% lower error than its BN counterpart when using a batch size of 2”. It holds because “GN's computation is independent of batch sizes” — the statistic is computed per example, so batch size never enters it.

Synchronization across devices is an opt-in that an engineer has to name. PyTorch offers no flag on the layer. Cross-device statistics live in a separate class, torch.nn.SyncBatchNorm (eps=1e-05, momentum=0.1), whose statistics are “calculated per-dimension over all mini-batches of the same process groups”. The documentation says it must be applied through torch.nn.SyncBatchNorm.convert_sync_batchnorm() before the network is wrapped in DistributedDataParallel. Skip that conversion and every worker normalizes with its own local micro-batch.

Input standardization uses statistics from the training data to place raw features on suitable scales. A normalization layer transforms intermediate activations using a rule embedded in the model.

The layer may contain learned affine parameters and mode-dependent state. Replacing one with the other changes both the function and the training dynamics.

Normalization alters scale, parameterization, noise, and gradient flow. Its benefit may not come from one simple story such as “reducing internal covariate shift.”

Focus on observable consequences. These include stable rates, faster progress, batch sensitivity, evaluation consistency, and resource cost. Each one is something you can measure on your own model, rather than something you accept from a paper's title.