Normalization: Controlling Representation Scale, Stabilizing Deep Optimization

The optimization of deep networks depends not only on what functions the model can express, but also on the range in which intermediate values fall. Even if two parameterizations have the same function expressiveness, the activation scales they produce differ, and the optimization difficulty may also differ significantly.

Suppose a layer receives an activation vector x. If the overall magnitude of x is too large, after being passed to subsequent linear layers it will be further amplified, and may make the attention logits too large, or push nonlinear units into the saturation region; in severe cases, numerical overflow may also occur. If the magnitude of x is too small, effective signals and gradients may be rounded away in finite-precision computation, or may become too weak relative to the residual branch and be drowned out. The problem is not just that a particular value is too large or too small at one time, but that the scales across different samples, tokens, and network layers drift continuously during training.

This scale drift changes the local conditioning that the optimizer actually faces. At the current position, the loss surface may have some directions that are very steep and some that are very flat; for a fixed learning rate, the same step size may be too large in steep directions and too small in flat directions. When the scale of intermediate representations keeps changing, this difference in local curvature also changes, making it harder for the optimizer to advance training with a stable set of step sizes.

The input to normalization is some activation or representation to be processed, and the output is a representation with more controlled scale. It first limits the numerical range through deterministic statistics and scaling rules, and then uses a learnable parameter γ to restore the magnitude required by the task; some forms of normalization also use a learnable parameter β to restore the required offset. The causal chain is: activation scale drift → unstable workspace and local curvature for subsequent operators → fixed learning rate finds it hard to accommodate all directions → normalization controls the scale → γ and β preserve the ability of the task to reselect magnitude and offset.

Therefore, normalization is more accurately a reparameterization and conditioning improvement technique. It makes parameter updates occur in a coordinate system where numerical values are more controllable, but it does not guarantee that every layer's output always follows a standard normal distribution, nor can it alone eliminate all saturation, overflow, or gradient problems. It solves the optimization conditioning problem caused by representation scale, rather than making a permanent probabilistic distribution commitment about the output distribution.

LayerNorm computes statistics independently for each token's hidden vector. Given a d-dimensional vector x, where xᵢ is the i-th coordinate, first compute the mean of all coordinates of the vector:

μ = (1/d)Σxᵢ

Here Σ means summing all d coordinates, and μ describes the overall magnitude around which this token's hidden representation is centered. Next compute the variance:

σ² = (1/d)Σ(xᵢ − μ)²

Each coordinate first subtracts the mean, then is squared, and the average is taken over all d coordinates. σ² measures the dispersion of the coordinates around the mean. LayerNorm then transforms the i-th coordinate as:

LN(x)ᵢ = γᵢ(xᵢ − μ)/$\sqrt{\text{σ² + ε}}$ + βᵢ

Subtracting μ performs centering, expressing each coordinate relative to the current token's own mean level; dividing by $\sqrt{\text{σ² + ε}}$ scales by the standard deviation, keeping the representation's dispersion under control. ε is a small positive number used to prevent the division from becoming unstable when the variance is very small. Finally, the per-dimension learnable parameters γᵢ adjust the magnitude of the i-th dimension, and βᵢ adjust the offset of the i-th dimension, so normalization does not permanently lock down the scale and center position the model can express.

μ and σ² are both computed from all hidden dimensions of the current token, and each token row independently obtains its own statistics. γᵢ and βᵢ are not quantities temporarily computed for each token; instead, they are parameters shared across samples and learned per hidden dimension. Therefore the input and output are still the same d-dimensional vector, but each output coordinate simultaneously depends on all coordinates of the input vector.

The variance formula uses denominator d, not the d − 1 common for unbiased sample variance in statistics. This is not estimating an unknown population variance from a sample; rather, it defines a deterministic vector transformation, so the population variance form is exactly part of the algorithm.

The statistical axis determines the semantics of the transformation. For a representation whose shape can be understood as “token × hidden”, LayerNorm computes statistics row by row along the hidden dimension; each token uses only its own hidden coordinates. If it is mistakenly computed along the sequence dimension, the tensor shape may still be correct and the program may run normally, but different tokens will participate in each other's statistics, and what is executed is no longer the LayerNorm defined here.

Take the input vector x = [1, 2, 3], first ignore ε, and set the learnable scale γ = 1 and offset β = 0. This allows us to observe only the centering and standard-deviation scaling themselves.

First step: calculate the mean of the three coordinates:

μ = (1 + 2 + 3)/3 = 2

Second step: subtract the mean from each coordinate to get the deviations from the mean:

x − μ = [1 − 2, 2 − 2, 3 − 2] = [−1, 0, 1]

This step shifts the center of the vector to 0 while preserving the relative high-low relationship among the three coordinates. Third step: calculate the population variance:

σ² = (1² + 0² + 1²)/3 = 2/3 ≈ 0.667

The standard deviation is the square root of the variance:

σ = $\sqrt{\text{2/3}}$ ≈ 0.816

Finally, divide each deviation by the standard deviation:

(x − μ)/σ = [−1, 0, 1]/0.816 ≈ [−1.225, 0, 1.225]

Therefore, under these simplified conditions, LayerNorm transforms [1, 2, 3] into approximately [−1.225, 0, 1.225]. The output has mean 0, and the three coordinates have a mean square of about 1; this corresponds to first removing the overall offset and then adjusting the dispersion scale to a uniform level.

This example also reveals what the transformation preserves and removes. If we add 10 to all inputs, the vector becomes [11, 12, 13]; the new mean also increases by 10, and after centering we still get [−1, 0, 1], so the normalization result is unchanged. If we multiply the input by a positive number 5, both the deviations from the mean and the standard deviation are multiplied by 5, and after division the result is also approximately unchanged. The "approximate" here reminds us that the actual implementation also includes ε, which causes the strict scale cancellation to change in the low-variance region.

The choices γ = 1 and β = 0 are only for convenience in hand calculation. In actual LayerNorm, after standardization, it also uses per-dimension affine parameters to relearn suitable scales and offsets for different hidden dimensions, so while obtaining a controlled numerical scale, it still retains the per-dimension adjustment capability required by the task.

RMSNorm does not perform centering; instead, it directly controls the root-mean-square scale of the original vector. Given a d-dimensional vector x, its root mean square is:

RMS(x) = $\sqrt{\text{(1/d)Σxᵢ² + ε}}$

where xᵢ is the i-th coordinate, Σ sums the squares of all coordinates, dividing by d gives the mean square, and taking the square root yields an overall scale of the same order of magnitude as the original coordinates. ε is used for numerical stability when the scale is very small. The output of the i-th dimension is:

RMSNorm(x)ᵢ = γᵢxᵢ/RMS(x)

γᵢ is a scaling parameter learned per hidden dimension. Compared with LayerNorm, this formula has no xᵢ − μ, so it does not first subtract the vector mean; it usually has only γ and no β for restoring the offset. Omitting centering means it only constrains the overall root mean square and does not force the output mean to 0, so the mean information in the original representation is preserved.

Taking x = [1, 2, 3] as an example, set γ = 1 and ignore ε. The sum of squares of the three coordinates is 1 + 4 + 9 = 14, so:

RMS(x) = $\sqrt{\text{14/3}}$ ≈ 2.160

After dividing coordinate-wise by this scale, the output is approximately:

[1, 2, 3]/2.160 ≈ [0.463, 0.926, 1.389]

The mean square of the three output coordinates is 1, but the mean is not 0. This is exactly the core difference between RMSNorm and LayerNorm: LayerNorm controls the variance after centering, while RMSNorm controls the root mean square of the original vector without centering.

Both are approximately insensitive to overall positive proportional scaling. If x is multiplied by a positive number, the numerator and the corresponding scale will increase together, and roughly cancel after division; whether they cancel exactly is also affected by ε. Overall translation is different: LayerNorm's centering removes the offset added uniformly to all coordinates, while RMSNorm does not subtract the mean, so translation changes the vector's root mean square, direction, and final output.

Therefore, the input and output shapes of RMSNorm are the same as the original hidden vector, but the structure it preserves differs from LayerNorm. Choosing it means you only need to control the overall energy scale of the original representation, while accepting that common offsets still affect the result; it does not provide the property of zero output mean.

The names of normalization methods are similar, but what actually determines the algorithm's semantics is "along which axes the statistics are computed." The statistical axis determines which input positions an output position depends on, and also determines whether external statistical information is needed during training and inference.

In BatchNorm for convolutional scenarios, statistics are usually computed for each channel, jointly along the batch dimension and spatial dimensions. A sample in the same channel therefore depends on other samples in the same batch. Training uses the current batch statistics, while inference commonly uses moving statistics accumulated during training, so the rules differ between the two stages.

LayerNorm computes statistics along the hidden dimension of a single sample or a single token. A token's mean and variance are determined only by its own hidden coordinates and do not depend on other samples in the same batch. Both training and inference can compute them from the current input using the same rules. RMSNorm has a similar dependency scope and also computes along the hidden dimension of a single token, but the statistic is the root mean square of the original vector. GroupNorm, within a single sample, computes statistics along pre-partitioned channel groups and corresponding spatial positions, likewise not depending on other samples, with consistent training and inference rules.

These dependency differences explain why LayerNorm is suitable for variable-length sequences and small batches: not because it is necessarily "better" on all tasks, but because each token's statistics do not change with the number, length, or content of samples in the same batch. Conversely, as long as the statistical axis includes the batch, a sample's normalization result may be affected by other samples in the same batch.

During implementation, you cannot judge an algorithm as correct merely by the "Norm" name or by a legally shaped tensor. If the model definition requires joint normalization over multiple dimensions, normalized_shape must precisely cover those dimensions; if you habitually select only the last dimension, you may get a computation with a completely correct shape but wrong dependency relationships. When validating a normalization implementation, first clarify the semantics of each axis of the input tensor, then confirm that the statistical axes, parameter shapes, and training and inference rules match the algorithm definition.

Normalization gains scale stability by removing certain global changes, but those removed changes may also be exactly the information the task needs. The key to understanding its boundary is distinguishing between “relative patterns across dimensions” and “the absolute mean or magnitude of the current sample.”

Ignoring ε and the affine parameters, for any positive number a, and for a shift b·1 added to all dimensions, LayerNorm satisfies:

LN(ax + b·1) = LN(x)

Here 1 denotes the vector whose coordinates are all 1. The common shift b enters both the input and the mean and cancels during centering; the proportional scaling a enlarges both the deviations from the mean and the standard deviation, and also cancels when divided. Therefore, LayerNorm preserves the relative high/low structure of each dimension compared with this token’s mean, but no longer preserves the token’s original overall mean or proportional scale.

Under the same simplified conditions, RMSNorm satisfies for proportional scaling:

RMSNorm(ax) = RMSNorm(x)

The numerator and the root mean square are both multiplied by a, so the proportion cancels. But RMSNorm does not subtract the mean, so if a constant is added to all coordinates, both the vector direction and the root mean square change, and the output changes accordingly. It removes the overall proportional magnitude, but remains sensitive to a common shift.

The learnable γ and β can set a fixed scale and shift shared across samples for each channel after normalization. They can reshape the typical magnitude and center of each hidden dimension, but cannot recover information that has been completely removed by the current sample statistics: if two inputs differ only by the overall mean or scale that LayerNorm removes, they already yield the same result before entering the affine transform, and afterward the shared γ and β cannot tell which original input was which.

Thus, scale insensitivity is both an optimization advantage and an information choice. If the absolute magnitude in a task itself represents confidence, energy, or counts, direct normalization may weaken or delete this signal. In that case you must confirm whether the architecture has a bypass that does not go through this normalization to preserve it; otherwise subsequent layers cannot reconstruct the vanished sample-level overall scale or mean from the normalization output alone.

Pre-Norm and Post-Norm use the same residual branch and normalization operator, but they are placed in different orders, so gradients take different paths when passing through deep networks. Let F denote a residual sublayer; Pre-Norm can be written as:

y = x + F(Norm(x))

Normalization occurs before the residual sublayer. The input x has an identity path that is directly added to the output y; during backpropagation, the gradient can travel along this backbone across the current layer without first passing through the derivative transformation of Norm. The F branch still depends on the normalized input, but the residual backbone preserves a more direct gradient channel.

Post-Norm can be written as:

y = Norm(x + F(x))

The input and the residual branch are added first, and the result is then normalized. This way, every layer's output must pass through Norm, and the backward gradient depends more on Norm's Jacobian. The Jacobian is a matrix composed of derivatives of the output with respect to each input coordinate; it describes how a local perturbation or gradient is scaled, mixed, and rotated as it passes through this transformation. The deeper the network, the more important the joint influence of consecutive layers' Jacobians on gradient propagation becomes.

This causal chain explains why even just moving the position of Norm can significantly change training difficulty: Pre-Norm provides a more direct identity residual path, generally makes deep Transformers easier to optimize, and reduces how sensitive training is to learning-rate warmup and initialization; Post-Norm immediately normalizes the result after each residual addition, but deep gradients must repeatedly pass through the normalization transformation.

This difference is not just a training trick. The position of normalization also affects the scale of the final representation and how the network uses depth, so the two structures define different computational processes. A trained Post-Norm model cannot become Pre-Norm losslessly merely by moving Norm before the residual branch; the original parameters were learned under the original gradient path and representation scale, and moving the position changes each layer's input, output, and overall function.

Normalization requires dividing by the scale of the representation. When a vector has almost identical values in all dimensions, the variance becomes very small, and the denominator also approaches 0. At this point, differences that were originally just rounding errors or tiny perturbations may be amplified by the division into noticeable changes.

For example, if the variance is 10⁻¹², the corresponding standard deviation is only $\sqrt{\text{10⁻¹²}}$ = 10⁻⁶. Dividing directly by this number is equivalent to applying a gain of about one million. Even if there is only extremely small rounding noise among the input coordinates, the normalized output may fluctuate wildly. If the formula adds ε = 10⁻⁵ to the variance, the denominator is approximately:

$\sqrt{\text{10⁻¹² + 10⁻⁵}}$ ≈ $\sqrt{\text{10⁻⁵}}$ ≈ 0.00316

In this case, the maximum gain is controlled by ε and is no longer arbitrarily determined by the near-zero input variance. Therefore, ε is not only used to avoid exact division by zero; it essentially sets the scaling upper bound and numerical sensitivity in low-variance regions.

The choice of ε involves a trade-off. A larger ε makes the lower bound of the denominator larger and the computation more stable, but the output variance will also deviate more noticeably from the ideal unit variance. When the variance is much larger than ε, the effect of ε is small; when the variance is of the same order as ε or smaller, the output scale is mainly determined by ε.

The exact position of the formula is also important.$\sqrt{\text{variance + ε}}$ and $\sqrt{\text{variance}}$ + ε are not the same transformation in low-variance regions: the former adds ε to the variance before taking the square root, while the latter adds ε after obtaining the standard deviation, resulting in different denominators and gain upper bounds. Therefore, when migrating weights, performing quantization, or replacing the normalization operator, you cannot only check the layer name. The original formula form, the specific value of ε, and the statistical axis must all match; otherwise, an implementation that seems close on ordinary inputs may show noticeable differences on low-variance inputs.

The most insidious errors in normalization implementations often do not break tensor shapes. For a representation of shape [B,S,H], both LayerNorm and RMSNorm should have each token compute statistics only along its own H dimension; the output still remains [B,S,H]. If S or B is mistakenly included in the statistics, the program may run as usual, but one token's result will incorrectly depend on other positions.

Affine parameters also have clear semantics. LayerNorm's γ and β and RMSNorm's γ are learned along hidden coordinates, so the parameters should correspond element-wise to the normalized feature dimension and then be reused across batch and sequence positions. The same γᵢ should act on the i-th hidden coordinate of all tokens, and should not change with the token, nor be misaligned to the sequence axis.

Dependency experiments can identify statistical-axis errors: fix the target token's hidden vector and change only other samples in the same batch or neighboring tokens. If the target token's LayerNorm or RMSNorm output changes accordingly, the implementation has mixed in input positions that should not participate. Then, set mutually distinguishable γ values for different hidden dimensions to directly observe whether the affine parameters broadcast correctly along the H dimension.

Therefore, “no NaN” or “output shape is correct” are only minimum conditions. Whether implementations are equivalent also depends on the statistical axis, the position of ε in the formula, and the parameter shapes and broadcasting directions of γ and β; these details collectively determine which inputs and parameters each output coordinate depends on.

The presence of normalization layers in a Transformer does not mean that numerical values at all positions will automatically remain stable. Norm only constrains the representations it actually receives and outputs; other links such as learning rate, initialization, loss scaling, residual accumulation over time, and attention softmax can still become uncontrolled. Therefore, when a loss spike appears, you cannot rule out numerical issues merely because the model “has Norm,” nor can you assume by default that the fault must come from the normalization layer.

If the input RMS to Norm is observed to increase continuously with network depth while the Norm output remains controlled, the likely link is residual branches accumulating in the same direction over time. Each layer’s normalization can only handle the input at its current position; it cannot prevent the residual backbone from continuously accumulating magnitude after many layers of addition. In this case, you should compare the residual backbone RMS with the branch update RMS to determine how large each update is relative to the backbone and whether it keeps pushing the representation higher in the same direction.

If the Norm output is normal but NaN appears in the attention stage, the fault is more likely to lie in the QK logits, the mask, or the softmax. You should record the maximum logits per attention head and check whether any row is entirely masked. Finite normalized inputs do not guarantee that subsequent dot products, mask processing, and exponential operations are all within safe ranges.

If spikes appear suddenly only in specific batches, you should correlate the numerical trajectory with the data. Abnormal sequence lengths, unusual sample content, or loss scaling can all trigger the problem; you need to record sample IDs, token counts, and gradient norms to determine whether the spike begins in the forward representations or only appears at the loss and backpropagation stage.

If quality degrades after switching to a different inference backend after training, rather than producing NaN directly, you should first check whether ε, the normalization axis, the formula form, and fused operators match the original implementation. Such differences often produce finite values but accumulate layer by layer into a shift in model behavior. Comparing the maximum output error between the two backends layer by layer can locate where the deviation first becomes noticeable.

A reliable diagnostic sequence is to measure segment by segment along the computation chain: data → layer input → Norm output → sublayer output → residual sum → loss. This can distinguish scale drift before normalization, differences in normalization implementations, internal overflow in sublayers, residual accumulation, and anomalies at the loss end, rather than attributing all explosion phenomena to Norm.

When validating a normalization implementation, first prove that it computes the intended formula, not merely that it outputs finite values with the correct shape. The most direct starting point is to use a small vector [1, 2, 3] that can be computed by hand: check the mean, variance, standard deviation, and final output of LayerNorm respectively, as well as the root mean square and final output of RMSNorm. A small example can quickly reveal whether it incorrectly subtracts the mean, omits the square root, uses the wrong denominator, or places ε incorrectly.

Next, check the algorithm's semantics using transformation invariance. When the input is multiplied by a positive number, the outputs of LayerNorm and RMSNorm should exhibit approximate scale insensitivity according to their respective definitions; when the same constant is added to all coordinates, LayerNorm's centering should eliminate this translation, while RMSNorm's output should change. If observed results contradict these properties, even if the error on ordinary random inputs is small, it indicates that the implemented formula or statistical axis may be wrong.

Degenerate and extreme inputs are used to check numerical boundaries. An all-identical vector and a zero vector will make the variance or root mean square close to its minimum, which can verify whether ε participates in the computation correctly; inputs with very small variance can test whether the gain is controlled; inputs with very large values can expose overflow or precision issues in squaring, accumulation, and reciprocal square root. Testing should not only judge whether NaN appears but also compare the output scale with reference results.

Statistical axes can be verified through dependency experiments. Fix one token's hidden vector and change only other samples in the same batch or adjacent tokens; a LayerNorm that computes statistics independently over the hidden dimension should not change with these irrelevant positions. If the result changes, the implementation may incorrectly aggregate along the batch or sequence dimension. At the same time, confirm how γ and β are broadcast: they must act element-wise along the intended hidden dimension, not accidentally align across tokens, batch, or other axes.

Finally, compare forward and backward passes with a trusted reference. Compare outputs and input gradients across various tensor shapes, sequence lengths, and ε settings to simultaneously check the formula, statistical axis, and affine broadcasting. If using a fused operator or an alternative backend, also compare errors layer by layer starting from the first normalization layer, and check the final metrics: layer-by-layer errors locate where deviations begin, and final metrics confirm whether the local differences have already changed model behavior.

Representations in deep networks undergo residual addition and sublayer transformations, so their scale can continually change across network layers, samples, and tokens. Scale drift pushes subsequent computations into different numerical regions and also changes the local conditions faced by the optimizer. Normalization, in this chain, takes the current representation, computes statistics for the current token along the model-specified hidden dimension, and outputs a representation with more controlled scale.

LayerNorm first subtracts the mean and then scales by the standard deviation after centering, thereby reducing the representation's sensitivity to common translation across dimensions and to overall proportional scaling. RMSNorm does not perform centering; it scales only by the root mean square of the original vector, thus controlling overall magnitude while preserving the response to common translation. Both preserve the relative patterns among dimensions, but they weaken certain absolute scale information.

Normalization is not the end of the computation. Per-feature γ can relearn the output magnitude, and β in forms such as LayerNorm can also relearn the offset. They allow the model to continue choosing the typical scale and center position of each feature in a controlled coordinate system, but they cannot restore sample-level mean or magnitude that has already been completely removed by the current sample statistics.

Whether Norm is placed before or after the residual branch also changes the path of backpropagation. When placed before the branch, the residual backbone has a more direct identity gradient path; when placed after the addition, each layer's gradient depends more on the Jacobian of the normalization transform. Therefore, the effect of normalization is determined not only by the formula but also by its position in the residual structure.

Numerical stability forms another layer of conditions. ε sets a lower bound on the denominator for small-variance inputs, limiting the extent to which tiny noise is amplified; whether it is inside or outside the square root also changes the actual gain in the small-variance region. The statistical axis, ε, formula form, and affine broadcasting are all part of the implementation semantics, not merely details that can be replaced arbitrarily.

Even if Norm itself is correct, attention logits, softmax, long-term residual accumulation, and loss computation can still go out of control independently. Normalization only constrains where it is located; it cannot substitute for stability measures across the whole model chain. Therefore, the actual contract of model weights includes the normalization axis, the exact formula, ε, the broadcasting method of affine parameters, and the residual position; when migrating or replacing operators, this contract must be preserved through layer-by-layer forward and gradient equivalence tests.

To continue understanding normalization, you can proceed along five dependency threads: statistical axis, residual backbone, overall architecture, training stability, and low-precision deployment.

Batch normalization illustrates another choice of statistical axis. It includes the batch in the statistical scope, so one sample's output depends on other samples in the same batch; the training phase uses batch statistics, and the inference phase often switches to moving statistics. Comparing it with LayerNorm, which computes statistics along the hidden dimension of a single token, clarifies why the statistical axis changes dependency relationships and the training and inference rules.

Residual Connection is the direct prerequisite concept for understanding Pre-Norm. The identity path lets the input bypass the residual branch and reach the output, and whether Norm is placed before or after the branch determines whether this path must pass through the Jacobian of normalization. Only by combining the residual structure with the normalization position can you explain why deep gradients differ depending on the choice of Pre-Norm and Post-Norm.

Transformer provides the overall architectural context in which normalization operates. Attention layers and feed-forward layers are both inside residual structures, and the placement of Norm around these sublayers determines what scale of input each sublayer receives and whether the representation after the residual addition is immediately normalized. After understanding the formula for a single Norm, you also need to trace its connections with attention, feed-forward layers, and the residual backbone across the full network.

Training spikes need to be connected with gradient clipping. Norm controls the representation scale at a specific position; it cannot directly guarantee that parameter updates are always controlled. Gradient clipping limits risk from the perspective of abnormal updates. Distinguishing the two helps avoid conflating representation stability with update stability.

Low-precision deployment also depends on quantization concepts. Normalization involves scale estimation, ε, formula form, and affine parameters; fused operators or quantization backends must preserve the semantic equivalence of these computations. Deployment acceptance cannot just check whether the final program runs; it must confirm that scale handling and per-layer results before and after fusion remain consistent with the original implementation.

After mastering these dependencies, you should be able to given any tensor, clearly identify the normalization axis, compute by hand the statistics and outputs of LayerNorm and RMSNorm, explain how the identity residual path in Pre-Norm propagates gradients, and design an implementation acceptance plan covering degenerate inputs, statistical axis, forward-backward comparison, and per-layer alignment.