Layer Normalization — Stabilizing Deep Transformer Training

Why activations need normalizing

A deep network is a long composition of linear maps and nonlinearities. Each layer multiplies its input by a weight matrix, and small mismatches in scale compound multiplicatively: if every layer inflates the norm of its activations by even a modest factor, forty layers later the values are enormous; if each shrinks them, they vanish. The same compounding hits the gradients on the way back. This drift is often called internal covariate shift — the distribution each layer must learn from keeps moving as the layers below it update.

Normalization attacks the problem directly. Instead of hoping initialization and learning rates keep every layer’s activations in a sane range, we force the range at each layer: subtract the mean, divide by the standard deviation, and hand the next layer a distribution with known first and second moments. The network no longer has to spend capacity chasing a moving scale; it can spend it on the actual function.

The definition

Let x = (x_1, …, x_H) be the activation vector for a single token, where H is the model width (the feature dimension). Layer normalization computes the mean and variance across those H features:

μ   = (1/H) Σ_i x_i
σ^2 = (1/H) Σ_i (x_i - μ)^2
x̂_i  = (x_i - μ) / √(σ^2 + ε)
y_i  = γ_i · x̂_i + β_i

The normalized value x̂_i (read ‘x-hat’) has zero mean and unit variance by construction. The small constant ε (typically 1e-5) guards the square root against a near-zero variance. Then two learnable vectors, the scale γ and the shift β — each of length H — are applied element-wise. The key structural fact: normalization happens per token, over the feature axis. Every token in the sequence, and every example in the batch, is normalized independently using only its own H numbers.

Shapes and where the reduction happens

Getting the axis right is most of understanding LayerNorm. A transformer activation tensor is typically X: [B, N, H] — batch, sequence length, hidden size. LayerNorm reduces over the last axis, H, producing per-position statistics of shape [B, N, 1] that broadcast back over the features.

Concretely, for every one of the B × N token positions we collapse H numbers into a single mean and a single variance, normalize that position’s vector, then scale and shift. The parameters γ and β are shared across all positions and all examples — they are properties of the layer, not of any token — so the whole operation adds just 2H parameters. There is no mixing between tokens and none between examples: contrast this with batch norm, which reduces over B (and N) and therefore couples examples together. That independence is exactly what makes LayerNorm behave identically at training time and at inference, and robust to any batch size, including a batch of one.

The scale and shift: gamma and beta

Forcing zero mean and unit variance is a blunt instrument. Some features genuinely want a large dynamic range; a downstream nonlinearity like a sigmoid or GELU behaves very differently at magnitude 0.5 than at magnitude 5. If normalization always clamped every feature to unit variance, we would be throwing away representational freedom the network might need.

The learnable γ and β give it back. After normalizing, the layer can rescale each feature by γ_i and re-offset it by β_i, learning the scale and center that suit the task rather than the ones normalization happened to impose. In principle the network can even undo the normalization entirely: setting γ_i = σ and β_i = μ recovers the original activation. That is the elegant part — normalization does not remove expressiveness, it changes the parameterization so that the well-conditioned setting (mean 0, variance 1) is the default and any departure from it is learned explicitly through two cheap parameters.

A worked numeric example

Take a token with four features, x = [2, 4, 6, 8], so H = 4. The mean is μ = (2+4+6+8)/4 = 5. The squared deviations are [9, 1, 1, 9], so σ^2 = (9+1+1+9)/4 = 5 and σ = √5 ≈ 2.236 (the ε is negligible here).

x̂ = (x - 5) / 2.236
   = [ -3/2.236, -1/2.236, +1/2.236, +3/2.236 ]
   ≈ [ -1.342, -0.447, +0.447, +1.342 ]

Check the result: the mean of x̂ is 0 and its variance is 1, exactly as promised, regardless of the original scale. With the default γ = 1 and β = 0 the output equals x̂. Feed in [20, 40, 60, 80] instead — ten times larger — and you get the same normalized vector: LayerNorm is invariant to a global rescaling of a token’s features, and to adding a constant to all of them. Those two invariances are precisely what stop layer-to-layer scale drift in its tracks.

LayerNorm versus batch norm

Batch normalization, which came first, normalizes each feature across the batch dimension: for feature i it uses the mean and variance of x_i over all examples in the mini-batch. That works well for convolutional vision models but is a poor fit for transformers on three counts.

First, sequence models have variable length and are frequently run at batch size one (a single prompt at inference), where a batch statistic is meaningless or unstable. Second, batch norm must maintain running averages to use at inference, creating a train/test discrepancy that LayerNorm simply does not have — its statistics are computed the same way in both regimes. Third, batch norm couples examples: the output for one token depends on the other tokens that happened to share its batch, which is awkward for autoregressive generation. LayerNorm sidesteps all three by reducing over features within a single token, making it batch-independent, length-independent, and identical at train and test time. That is why essentially every transformer uses it.

The backward pass

Training needs gradients through the normalization. Writing x̂ for the normalized vector and g_i = ∂L/∂y_i for the upstream gradient, the parameter gradients are simple sums over all token positions:

∂L/∂γ_i = Σ_tokens  g_i · x̂_i
∂L/∂β_i  = Σ_tokens  g_i

The gradient with respect to the input is where the coupling lives. Because μ and σ each depend on every x_j, changing one input nudges the normalized value of all the others, so the Jacobian is not diagonal. Let h_i = g_i · γ_i be the gradient after the affine step. Then, with H features:

∂L/∂x_i = (1/σ) ·
   [ h_i  -  (1/H) Σ_j h_j  -  x̂_i · (1/H) Σ_j h_j x̂_j ]

Read the three terms as corrections: the raw gradient, minus its mean (the path through μ), minus a component along x̂ (the path through σ).

Reading the Jacobian

The backward formula is more than bookkeeping — it explains the stabilizing effect directly. Look at what those two subtracted terms do to the incoming gradient h. The first removes its mean, so the gradient that reaches the input has zero mean across features. The second removes its projection onto x̂, so the gradient is also orthogonal to the current normalized activation.

Together these mean LayerNorm passes back only the part of the gradient that actually changes the shape of the activation vector, not its overall scale or offset — the two directions the forward pass is invariant to anyway. A component that merely tries to rescale or shift every feature uniformly is projected out, because it cannot change the output. The overall 1/σ factor then ties gradient magnitude to activation scale: large activations damp gradients, and vice versa. This self-regulation is a big part of why gradients neither explode nor vanish across depth.

Why it stabilizes training

Pulling the threads together, LayerNorm stabilizes training through several reinforcing mechanisms. The forward pass guarantees each layer receives inputs with fixed first and second moments, so no layer ever sees activations that have drifted to a tiny or gigantic scale. The invariance to per-token rescaling and shifting means an unlucky weight update in one layer cannot blow up the activation norm feeding the next.

On the backward side, the mean-subtraction and orthogonalization we just derived keep gradient magnitudes bounded and remove the useless scale/shift directions, which smooths the loss landscape and lets you use a larger, more aggressive learning rate without divergence. Empirically the payoff is that networks train faster, tolerate a wider range of hyperparameters, and — most importantly — remain trainable at depths where an unnormalized stack would simply diverge or stall. LayerNorm does not change what a transformer can represent; it changes the optimization so that the good solutions are actually reachable by gradient descent.

Epsilon and numerical care

The ε inside the square root is not cosmetic. When a token’s features are nearly identical, the variance approaches zero and 1/√σ^2 would blow up, sending a small amount of noise to an enormous normalized value and destabilizing both the forward and backward passes. Adding ε (commonly 1e-5 or 1e-6) bounds that ratio.

Precision matters too. The sum of squares in the variance is prone to cancellation and overflow in low precision, so mean and variance are almost always accumulated in fp32 even when the activations are stored in bf16 or fp16. Skipping that is a classic source of subtle, slow-to-diagnose divergence at large width — a careless implementation quietly corrupts the very stability it is meant to provide.

Practical and CPU-SLM notes

On a CPU-served small language model, LayerNorm is cheap in FLOPs — a couple of passes over H numbers per token — but it is memory-bandwidth bound, because it forces a reduction (the mean and variance) before any output can be written. Without a GPU’s bandwidth, those reductions can take a surprising slice of per-token latency.

So LayerNorm is a prime candidate for fusion: folding it into the adjacent residual-add or the following matmul avoids re-reading the activation tensor from memory. Its cost is also one motivation for RMSNorm, which drops the mean-subtraction and one reduction for a measurable speedup at similar quality. Both are covered in the companion articles; the point here is that this tiny operation is worth optimizing precisely because it runs in every layer.