Layer Normalization

Given a single training example with a feature (activation) vector $x \in \mathbb{R}^H$ — for instance, the $H$ hidden units of one layer for one token — layer normalization computes the mean and variance across the $H$ features of that one example:

It then standardizes and applies a learned affine transform, per feature:

Here $\gamma, \beta \in \mathbb{R}^H$ are learned parameters (one scale and one shift per feature), and $\epsilon$ is a small constant for numerical stability. Crucially, $\mu$ and $\sigma^2$ are computed independently for every example in the batch, using only that example's own $H$ numbers. There is no dependence whatsoever on other rows in the batch. Contrast this with batch normalization, which for a mini-batch of size $N$ computes, for each feature $i$, statistics across the $N$ examples:

Batch normalization reduces variance across examples for a fixed feature; layer normalization reduces variance across features for a fixed example. This single transposition of the reduction axis has large practical consequences: layer normalization needs no running mean/variance buffers, behaves identically in train and eval mode, and its output for one example is completely unaffected by which other examples happen to be in the same batch — or whether there even is a batch (batch size 1 is fine). The terminology "layer" refers to normalizing over all the summed inputs to the neurons of a layer for one training case — historically framed for fully connected layers, later generalized (in Transformers) to normalize over the last dimension of a tensor, e.g., the $d_{\text{model}}$ embedding dimension for each token independently of sequence position and batch index.

Forward pass. For each example (each row of a batch, or more generally each position along all axes except the normalized one), layer normalization: 1. Computes the mean of that example's feature vector. 2. Computes the variance of that example's feature vector. 3. Standardizes: subtract the mean, divide by the standard deviation (plus $\epsilon$). 4. Applies a learned per-feature scale $\gamma$ and shift $\beta$, shared across all examples but distinct per feature index. Because the reduction is over features, not batch, this is a genuinely per-example operation: you could run it one example at a time and get an identical result to running it on a whole batch. This is exactly why it needs no running statistics for inference — at test time it does precisely the same computation it did at training time. ```text

# Forward pass of Layer Normalization

# x : matrix of shape (N, H) -- N examples, H features each

# gamma : vector of shape (H,) -- learned scale, shared across examples

# beta : vector of shape (H,) -- learned shift, shared across examples

# eps : small constant, e.g. 1e-5 function LAYER_NORM_FORWARD(x, gamma, beta, eps): N, H = shape(x) y = zeros(N, H) cache = [] # store per-example intermediates for backward for n in 1..N: row = x[n, :] # this example's H-dim feature vector mu = mean(row) # scalar, O(H) var = mean((row - mu) ** 2) # scalar, O(H) (biased/population variance) std_inv = 1 / sqrt(var + eps) x_hat = (row - mu) std_inv # standardized vector, shape (H,) y[n, :] = gamma x_hat + beta # elementwise scale + shift cache[n] = (x_hat, std_inv, mu, var) return y, cache ``` ```text

# Backward pass of Layer Normalization

# dY : upstream gradient, shape (N, H)

# cache : per-example (x_hat, std_inv, mu, var) from forward

# gamma : shape (H,)

# Returns : dX (N,H), dGamma (H,), dBeta (H,) function LAYER_NORM_BACKWARD(dY, cache, gamma): N, H = shape(dY) dGamma = zeros(H) dBeta = zeros(H) dX = zeros(N, H) for n in 1..N: x_hat, std_inv, mu, var = cache[n] dy = dY[n, :] # shape (H,) dGamma += dy x_hat # accumulate over examples dBeta += dy dx_hat = dy gamma # shape (H,), chain through scale

# standard "normalization backward" identity (derived below):

# dx = std_inv/H ( Hdx_hat - sum(dx_hat) - x_hat sum(dx_hat x_hat) ) sum_dxhat = sum(dx_hat) sum_dxhat_xhat = sum(dx_hat x_hat) dX[n, :] = (std_inv / H) (H dx_hat - sum_dxhat - x_hat sum_dxhat_xhat) return dX, dGamma, dBeta ``` The backward formula looks intricate but follows mechanically from applying the chain rule through $\mu$, $\sigma^2$, and the division — it is algebraically identical in structure to batch normalization's backward pass, with the sums running over the feature axis $H$ instead of the batch axis $N$. This symmetry is exactly what one would expect: layer norm is "batch norm with the reduction axis transposed," so its computational graph, and hence its gradient derivation, transposes accordingly.

Time complexity. For a batch of $N$ examples each with $H$ features, computing the mean and variance for one example costs $O(H)$ (one pass to sum, one pass for squared deviations, or a single fused pass using the sum-of-squares identity). Doing this for all $N$ examples costs $O(NH)$. The standardization and affine transform are elementwise, also $O(NH)$. So the forward pass is $\Theta(NH)$ time — asymptotically identical to batch normalization, and identical to the cost of simply reading the input once (any normalization scheme must touch every element at least a constant number of times, so this is optimal up to constant factors). The backward pass reuses the same cached per-example scalars and is likewise $\Theta(NH)$: each example's gradient computation involves two more reductions over its $H$ features (`sum_dxhat`, `sum_dxhat_xhat`), each $O(H)$, times $N$ examples.

Space complexity. The cache stores, per example, $x̂ \in \mathbb{R}^H$ plus two or three scalars ($\mu$, $\sigma^2$ or its inverse std), so the cache is $\Theta(NH)$ — dominated by storing $\hat x$, which is the same order as the input itself. The learned parameters $\gamma, \beta$ are $\Theta(H)$, independent of $N$ — a key structural fact: capacity does not grow with batch size.

Why it is correct — invariance argument. The purpose of any normalization layer is to control the scale of activations flowing into subsequent layers, combating internal covariate shift and vanishing/exploding gradients. The correctness claim to check is: does layer normalization actually make outputs invariant to certain nuisance rescalings of its input, in a way that helps optimization? Layer normalization is invariant to per-example rescaling of the pre-normalization input, in the following precise sense. Suppose $x' = a \cdot x + b\mathbf{1}$ for a per-example scalar $a > 0$ and vector of constant shift $b\mathbf{1}$ (the same constant added to every feature of that example). Then

so

This means layer normalization's output is (approximately, exactly as $\epsilon \to 0$) unchanged if the entire feature vector of one example is uniformly rescaled and shifted — for instance, if the weight matrix feeding into that layer is rescaled by a constant, or if a bias shifts all pre-activations of one example equally. This is exactly the invariance batch normalization achieves across the batch axis; layer normalization achieves the analogous invariance within each example, which is precisely the axis relevant to phenomena like re-parameterizing per-example activation magnitude in a recurrent network as it evolves over time steps — the internal state at time $t$ can drift in scale independent of other time steps or other sequences in the batch, and layer norm renormalizes each one back to unit scale on its own terms. A second correctness point concerns representational capacity: normalizing forces the pre-activation distribution (per example) to have mean 0 and variance 1, which could in principle destroy useful information (e.g., if the network needs large activations for some feature). The learned $\gamma, \beta$ restore full expressiveness: since $\gamma, \beta$ are free parameters, the layer can represent the identity mapping ($\gamma = \sqrt{\sigma^2+\epsilon}$, $\beta=\mu$, recovering $x$ exactly) or any other affine reshaping of the standardized vector. So layer normalization is a reparameterization of the space of achievable outputs, not a strict restriction — the network loses nothing in principle, but gains a smoother, more scale-controlled optimization landscape, because now every layer's output is guaranteed to sit in a fixed, well-conditioned range regardless of how weights evolve during training.

Tradeoffs. Layer normalization removes the batch-size dependency of batch normalization but pays a price: it does not remove feature-wise covariate shift the way batch norm removes cross-example variability for a fixed feature, so on tasks where different features have very different natural scales (e.g., some CNN feature maps), layer norm can be a worse fit than batch norm — this is part of why convolutional vision architectures historically kept batch norm, while sequence and attention architectures moved to layer norm. Layer norm also cannot exploit any regularizing "noise" that batch statistics introduce (batch norm's dependence on which examples land in a batch acts as a mild stochastic regularizer); layer norm's deterministic per-example computation offers no such implicit regularization.

Take one example (one token's hidden state) with $H = 4$ features, before normalization:

Suppose $\gamma = [1, 1, 1, 1]$ and $\beta = [0, 0, 0, 0]$ (freshly initialized, so layer norm currently behaves as pure standardization), and $\epsilon = 10^{-5}$.

Confirm the invariance claim from the previous section: rescale this same example by $a = 3$ and shift by $b = 10$, giving $x' = [16, 22, 22, 28]$.

Up to the $\epsilon$-induced rounding (matching to 4 decimal places), $\hat x' = \hat x$ — exactly the invariance predicted: multiplying the whole feature vector by 3 and adding 10 to every entry left the standardized output unchanged, because layer norm only cares about the shape of the per-example vector, not its absolute scale or offset. Now suppose training has adjusted $\gamma = [0.5, 0.5, 2.0, 2.0]$, $\beta = [1, 1, -1, -1]$. Applying these to the original $\hat x = [-1.41419, 0, 0, 1.41419]$:

This shows how the learned affine parameters let the network place each feature at whatever scale and offset is useful downstream, even though the pre-affine vector $\hat x$ is forced to be zero-mean, unit-variance for every example.

Layer normalization was introduced by Jimmy Lei Ba, Jamie Ryan Kiros, and Geoffrey E. Hinton in 2016, explicitly framed as a fix for batch normalization's dependence on batch size and its awkwardness in recurrent networks, where activation statistics vary across time steps and small or variable batch sizes are common. Hinton, already a central figure in the deep learning revival through backpropagation, Boltzmann machines, and dropout, co-authored the paper that reframed normalization as a per-example rather than per-batch operation — a reframing that would later become indispensable to an architecture (the Transformer) that did not yet exist when the paper was published.

Layer normalization is embedded in essentially every modern large language model: GPT-family decoders, BERT-family encoders, T5, and their descendants all place a layer-norm (or RMSNorm) operation before or after each attention and feed-forward sublayer. Its independence from batch size makes it the natural choice for autoregressive decoding, where inference often proceeds one token — effectively batch size one — at a time, and for training regimes with variable sequence lengths and padding, where naive batch statistics would be contaminated by padding tokens. It is equally at home in recurrent architectures (LSTMs, GRUs) applied to variable-length sequences, and in reinforcement learning settings where "batches" of experience are small, correlated, or nonstationary — exactly the regime where batch normalization's cross-example statistics become unreliable. Beyond Transformers, layer normalization (and its relatives, instance normalization and group normalization) inform a broader lesson in deep learning systems design: the right normalization axis depends on which dimension is a genuine nuisance variable and which carries real signal. Recognizing this tradeoff — and knowing how to diagnose when a network is misbehaving because of a normalization mismatch — is a core skill covered further in Debugging the Training Process, alongside related tools such as Gradient Clipping and careful Weight Initialization, all of which interact closely with how a network's activation statistics evolve during training.

Ask the tutor: How does the choice between Pre-LN and Post-LN placement change gradient flow through a deep Transformer's residual stream?