Layer normalization is the small, unglamorous operation that makes deep transformers trainable at all. For every token independently, it takes that token’s feature vector, subtracts the mean, divides by the standard deviation, and then re-scales and re-shifts with two learned vectors. That is the whole operation — four lines of arithmetic — yet it is the reason gradients flow cleanly through dozens of stacked blocks, the reason the same forward pass works at training time and during single-token generation, and a quiet but real slice of the compute budget on a CPU-hosted small model. This piece builds it from first principles: what axis it normalizes and why, the exact forward math with a worked numeric example, the roles of the gain γ and bias β, why LayerNorm displaced BatchNorm for sequence models, the backward pass and its three coupled gradient terms, and how it stabilizes training. RMSNorm and pre-vs-post placement are covered briefly as siblings.
What axis LayerNorm normalizes
The first thing to pin down is which numbers get averaged together. A batch of transformer activations is a tensor of shape [B, N, d] — B sequences, N tokens each, and d features (the model dimension) per token. Layer normalization computes its statistics along the feature axis d, separately for every one of the B × N tokens. Each token vector x ∈ ℝ^d gets its own mean and its own variance, computed from just its d coordinates.
This is the defining choice, and everything else follows from it. Because the statistics come only from a single token’s own features, there is no coupling across tokens and no coupling across the batch — token 5 in sequence 2 is normalized using nothing but token 5’s own d numbers. Contrast this with BatchNorm, which normalizes each feature across the batch (mixing different examples together). LayerNorm never looks sideways at other examples, which is exactly what makes it safe for variable-length sequences, tiny batches, and one-token-at-a-time generation. Keep the shapes in mind: reduce over d, broadcast back over d.
The forward pass, line by line
For one token vector x = (x_1, …, x_d), LayerNorm is four steps:
μ = (1/d) · Σ_i x_i # mean over the d features
σ² = (1/d) · Σ_i (x_i - μ)² # (biased) variance over d
x̂ = (x - μ) / √(σ² + ε) # normalize: mean 0, var 1
y = γ · x̂ + β # scale and shift (elementwise)The first two lines collapse the d features into two scalars, μ and σ². The third line is the actual normalization: subtract the mean so the vector is centered, divide by the standard deviation √(σ² + ε) so it has unit variance. After this step x̂ has mean 0 and variance 1 by construction, regardless of what scale the incoming activations had. The fourth line applies the two learnable per-feature vectors γ, β ∈ ℝ^d elementwise. Note the variance uses 1/d, not the Bessel-corrected 1/(d-1) — this is the biased estimator, and it matches what PyTorch’s nn.LayerNorm actually computes. The whole thing is differentiable and cheap: a couple of reductions and a couple of elementwise ops.
A worked numeric example
Take a tiny model dimension d = 4 and one token x = [1, 2, 3, 4]. Walk the four steps by hand.
μ = (1+2+3+4)/4 = 2.5
dev = x - μ = [-1.5, -0.5, 0.5, 1.5]
σ² = (2.25+0.25+0.25+2.25)/4 = 5.0/4 = 1.25
√(σ²+ε) ≈ √1.25 ≈ 1.1180 (ε = 1e-5 is negligible here)
x̂ = dev / 1.1180 ≈ [-1.342, -0.447, 0.447, 1.342]Sanity-check the result: the mean of x̂ is (-1.342 - 0.447 + 0.447 + 1.342)/4 = 0, and the variance is (1.8 + 0.2 + 0.2 + 1.8)/4 = 1.0 — centered and unit-variance, exactly as promised. Now apply an affine transform, say γ = [2,2,2,2] and β = [1,1,1,1]: the output is y = 2·x̂ + 1 ≈ [-1.684, 0.106, 1.894, 3.684]. The shape of the vector (its pattern of relative values) is preserved from x̂; only its overall scale and offset changed. That separation — normalization fixes the distribution, γ/β choose the scale — is the entire design.
Why subtract the mean and divide by the deviation
The point of normalization is to stop activation magnitudes from drifting as they pass through many layers. Without it, each block’s output distribution depends on the block before it; small systematic biases compound multiplicatively across depth, and activations can blow up or collapse toward zero. This ‘internal covariate shift’ forces later layers to constantly re-adapt to a moving input distribution, which slows and destabilizes training.
Centering (subtracting μ) removes the arbitrary DC offset of the vector; scaling (dividing by √(σ²+ε)) fixes its overall magnitude. Geometrically, the operation projects each token vector onto a region of roughly constant radius — every normalized token has the same length in a statistical sense, so the layer that reads it sees inputs of a predictable scale no matter what the previous layer produced. The activations become self-similar in magnitude across depth and across tokens. That predictability is what lets you stack 12, 32, or 96 blocks and still get gradients of a usable size at the bottom of the stack. The normalization does not change what information a token carries — only its scale — so it is nearly free in representational terms but hugely stabilizing in optimization terms.
The learnable gain γ and bias β
Forcing every token to mean 0 and variance 1 is a strong constraint, and sometimes it is the wrong one — a given feature might genuinely need a larger spread, or a nonzero baseline, for the next layer to work well. The affine parameters γ (gain/scale) and β (bias/shift) give that capacity back. They are per-feature vectors in ℝ^d, learned by gradient descent like any other weights, and applied elementwise: y_i = γ_i · x̂_i + β_i.
Crucially, γ and β can undo the normalization if that is what minimizes the loss: setting γ_i = √(σ²+ε) and β_i = μ would recover the original activation. So normalization does not remove representational power; it re-parameterizes it, moving the choice of scale and shift out of the fragile, coupled forward computation and into two clean, independently learned vectors. Typical initialization is γ = 1, β = 0, i.e. start as pure normalization and let training deviate from it as needed. Because there are only 2d of these parameters per norm layer, they are a rounding error in the total parameter count but matter a lot for what the model can express right after each normalization.
LayerNorm vs BatchNorm: the batch-independence win
BatchNorm normalizes each feature using statistics computed across the batch — for feature j, it averages x_j over all examples in the mini-batch. That works beautifully for large-batch vision training but breaks in three ways that matter for sequence models. First, the statistics depend on batch composition, so the output for one example changes depending on which other examples happen to share its batch — a strange coupling. Second, it needs a decently large batch to estimate the per-feature mean/variance reliably; tiny batches give noisy statistics. Third, it must maintain running averages at training time to use at inference, creating a train/eval discrepancy.
LayerNorm sidesteps all three because its statistics come from a single token’s own features. There is no batch dependence at all: the same token produces the same normalized output regardless of what else is in the batch, or whether there is a batch at all. The forward pass is identical at training and inference — no running averages, no mode switch, no model.eval() subtlety for the norm. For autoregressive generation, where you often process one token with an effective batch of 1, this is not a convenience but a requirement.
Why sequences break BatchNorm specifically
Sequence models pile on extra reasons BatchNorm is a poor fit. Sequences have variable length, so a naive batch-statistic over the [B, N, ·] tensor mixes real tokens with padding, and the number of valid tokens per feature-column varies from batch to batch — the statistics become a function of your padding scheme. Masking around this is fiddly and error-prone.
Worse, the whole premise of BatchNorm — that examples in a batch are interchangeable samples from one distribution — is shaky for tokens. Position 0 and position 500 in a sequence are not exchangeable draws; averaging a feature across positions and across sequences blends genuinely different distributions. Autoregressive decoding makes it impossible anyway: you generate token t+1 before token t+2 exists, so there is no batch of future tokens to normalize against, and using batch statistics would leak information across positions in a way that violates causality. LayerNorm’s per-token, per-position independence dodges every one of these problems — it treats each token as its own self-contained unit, which is precisely the right granularity for a sequence of variable length generated left to right.
The backward pass, and why it has three terms
Backprop through LayerNorm is more interesting than the forward pass because μ and σ² both depend on all the features, so perturbing one input x_i ripples everywhere. Let x̂ be the normalized vector and dx̂ = dy · γ the upstream gradient after undoing the affine step. The gradient with respect to the input is:
dx = (1/d) · (1/√(σ²+ε)) ·
( d · dx̂ - Σ_j dx̂_j - x̂ · Σ_j (dx̂_j · x̂_j) )
dγ = Σ_tokens dy · x̂ # summed over all tokens in the batch
dβ = Σ_tokens dy # summed over all tokens in the batchThe three terms inside the parentheses have clean meanings. d · dx̂ is the direct effect of x_i on its own normalized value. −Σ_j dx̂_j is the correction because changing x_i shifts the mean, which moves every output. − x̂ · Σ_j(dx̂_j x̂_j) is the correction because x_i also shifts the variance, again affecting all positions. The two subtracted terms are what make the gradient sum to zero along d — the Jacobian projects out the mean and variance directions. Frameworks compute this for you, but knowing the shape of it is invaluable when debugging exploding or vanishing gradients around a norm layer.
Effect on gradient flow and training stability
The subtracted terms in the backward pass are not incidental — they are the source of LayerNorm’s stabilizing effect on optimization. Because the input gradient is projected orthogonal to the mean and variance directions, the norm layer automatically discards the components of the upstream gradient that would merely rescale or shift the activations. Those components are exactly the ones the affine parameters absorb, so the input path only receives the gradient that changes the token’s direction, not its scale.
The practical payoff is that the effective scale of gradients is decoupled from the scale of activations. In an un-normalized deep net, a layer that happens to produce large activations passes large gradients backward, and the product across many layers explodes or vanishes; LayerNorm re-standardizes the forward signal at every block, so the backward signal stays in a usable range no matter the depth. This is why transformers can be trained with relatively large learning rates and still converge, why they are far less sensitive to weight initialization than they would otherwise be, and why removing the norm layers from a deep transformer typically makes it diverge within a few steps. Normalization does not add capacity — it makes the capacity that is there trainable.
Pre-norm vs post-norm placement (brief)
Where you put the norm relative to the residual connection changes training dynamics substantially. The original transformer used post-norm: x → x + Sublayer(x), then LayerNorm(·) on the sum. Modern models overwhelmingly use pre-norm: normalize first, then add — x → x + Sublayer(LayerNorm(x)).
The difference is what the residual highway carries. In pre-norm the raw x flows through the skip connection untouched by any normalization, so there is a clean identity path from input to output and gradients reach early layers directly — this is what makes very deep stacks trainable without a learning-rate warmup babysitter, at the cost of a slightly less expressive per-block transformation. Post-norm normalizes the combined signal, which can yield marginally better final quality but is notoriously touchy to train deep, usually needing warmup and careful initialization. This is a genuine sibling topic with its own trade-offs; here it is enough to know that the math of the norm itself is identical in both — only its position in the residual block moves.
RMSNorm: the cheaper cousin (brief)
RMSNorm is LayerNorm with the mean-centering step removed. It skips computing μ and skips subtracting it, dividing instead by the root-mean-square of the raw features:
rms = √( (1/d) · Σ_i x_i² + ε )
y = γ · ( x / rms ) # no μ, and often no βThe empirical finding behind RMSNorm is that most of LayerNorm’s benefit comes from the scaling, not the centering — re-scaling to a fixed magnitude is what stabilizes training, and subtracting the mean adds comparatively little. Dropping the mean removes one reduction and one subtraction over d, and the bias β is frequently dropped too, which is a small but real compute and memory saving that compounds across every norm in a deep model. Many recent large models (the LLaMA family, for instance) use RMSNorm for exactly this reason. It is a sibling article’s worth of detail; the one-line takeaway is that RMSNorm is a deliberate simplification of the LayerNorm math — keep the variance-like scaling, drop the mean.
CPU-SLM cost: where the FLOPs and the memory go
On a GPU, LayerNorm is a rounding error next to the big matrix multiplies. On a CPU-hosted small language model the accounting is less forgiving, and it is worth knowing why. Each LayerNorm touches every activation twice in the forward pass — once to compute the reductions for μ and σ², once to apply the normalization and affine — so it is a few passes over an [N, d] tensor per layer. The FLOP count is modest (O(N · d) per norm, linear in both), but the operation is memory-bandwidth-bound, not compute-bound: it reads and writes the full activation tensor while doing very little arithmetic per element.
On a CPU, where memory bandwidth is the scarce resource and there is no army of cores to hide latency, these bandwidth-bound elementwise-and-reduce passes take a real, measurable slice of per-token latency — and a transformer has two norms per block, so a 24-block model runs 48 of them per token. This is precisely why fused kernels matter: computing mean and variance in one pass (Welford or a sum/sum-of-squares fusion) and folding the affine into the same pass avoids re-reading the tensor from memory. It is also part of RMSNorm’s appeal on CPU — one fewer reduction and one fewer elementwise op per norm is bandwidth you get back on every single token you generate.
Common pitfalls and practical notes
A handful of mistakes recur. Normalizing the wrong axis: LayerNorm must reduce over the feature dimension d; accidentally reducing over the token axis N silently turns it into something like a batch/instance norm and couples tokens that should be independent. Always check that normalized_shape matches the last dimension(s). Dropping or mis-scaling ε: it exists to keep √(σ²+ε) away from zero for near-constant vectors; too small and you risk division blow-ups in low precision, too large and you damp the normalization. The default 1e-5 (sometimes 1e-6) is usually right.
Precision: compute the mean and variance in float32 even when the model runs in bfloat16/float16 — the sum-of-squares is where half precision loses the most, and a low-precision reduction can produce a visibly wrong variance. Forgetting the affine: if γ/β are absent or frozen at 1/0, the layer can only produce unit-variance outputs, which occasionally hurts. And remember that γ, β are learned parameters that need gradients and an optimizer slot — small in count, but real. Get the axis, the ε, and the reduction precision right and LayerNorm is one of the most robust components in the whole stack.