Layer Normalization and the Residual Stream

Layer normalisation normalises across the features of a single example rather than across a batch. The change looks like a substitution of which axis the mean is taken over, and it removes the entire class of problems in which the function being trained depends on how the batch was composed. What it adds is a dependency on how the features are grouped: a normalisation over the channel axis of a feature map, or over the hidden axis of a token, defines a different function from one over all axes at once, and the choice interacts with the residual connections around it.

That interaction is the subject here, because in deep residual and transformer stacks the normalisation’s placement is what determines whether the signal carried by the residual path grows, shrinks, or stays bounded as depth increases.

The statistics without a batch axis

For an input vector $x \in \mathbb{R}^{H}$ belonging to one example (one position, one token, one time step), the layer computes over the $H$ components:

$$\mu = \frac{1}{H}\sum_{i=1}^{H} x_i, \qquad \sigma^2 = \frac{1}{H}\sum_{i=1}^{H}(x_i-\mu)^2, \qquad y_i = \gamma_i \frac{x_i-\mu}{\sqrt{\sigma^2+\epsilon}} + \beta_i$$

Two properties follow from the axis choice. There is no running buffer and no train-eval difference: the same function is applied in both modes, so the mismatch described in the batch normalisation train-eval gap does not arise. And the function is invariant to the batch: the output for one example is unaffected by any other example, so gradient computation per example is exact rather than coupled through batch statistics.

The invariance has a price at the batch level: the layer cannot absorb a per-channel offset that is common across examples, because each example centres itself. If a channel’s activations are systematically offset, the offset is removed from every example independently, and the information that the channel was offset is destroyed rather than corrected. Whether that is desirable depends on whether the offset carries signal.

Budget note on $H$: the normalisation reduces $H$ degrees of freedom to $H-1$, since the mean-subtracted vector is constrained to sum to zero. With $H$ small — say $4$ — the constraint is a substantial fraction of the vector, and the normalised output has rank at most $H-1$. Nothing breaks, but the layer cannot produce an arbitrary output vector of dimension $H$; the gain vector $\gamma$ and bias $\beta$ restore a full-rank affine map only in the sense of adding back fixed directions, not in the sense of recovering the removed one.

The affine parameters and what scale invariance does

The parameters $\gamma, \beta$ are per-feature and are trained. The norm of the pre-normalisation vector is discarded by the division, so the layer’s output depends on the direction of $x$ only. This makes the gradient with respect to $x$ orthogonal to $x$ in a precise sense: $x^\top \nabla_x y = 0$ for the normalised part, because scaling $x$ leaves $y$ unchanged. The learnable gain breaks the invariance, since scaling $x$ now scales the normalised vector, but $\gamma$ is a fixed vector and the invariance to the magnitude of $x$ along directions where $\gamma$ is constant is restored.

The practical consequence for initialization is that the scale of the input to a normalised layer is irrelevant to the layer’s output, and the scale of the layer’s output is set by $|\gamma|$ and the residual stream’s statistics rather than by the preceding weight matrices. This is what makes it possible to train networks whose depth would otherwise require the variance-preserving initialization arguments of variance preservation in weight initialization to hold across dozens of layers.

Placement relative to the residual branch

A residual block computes $x + F(x)$ for some sub-network $F$, and a normalisation layer can sit before $F$, after $F$, or on the branch input. Three arrangements, all of which see use, have distinguishable effects on the signal’s scale in the residual stream.

The original transformer applied normalisation after attention and after the feedforward sublayer and added the result to the stream: $x \leftarrow \mathrm{LN}(x + F(x))$. Every block’s output is normalised, so the stream’s scale is reset at each block. This behaves well at shallow depth and has a well-known cost: the gradient flowing back through the stream passes through the normalisation at every block, and the normalisation’s Jacobian multiplies each block’s contribution. The product of many such Jacobians affects the gradient’s conditioning; the analysis of pre-normalisation versus post-normalisation in transformers shows that post-normalisation requires warmup to remain stable at depth, while pre-normalisation does not.

The alternative places normalisation inside the branch: $x \leftarrow x + F(\mathrm{LN}(x))$. The residual path is then an unmodified identity, so a gradient arriving at the output reaches the input through a sum of identity paths, one per block. The depth of the path does not attenuate it. The stream’s own scale is not reset, however, and grows additively: each block adds its branch’s output, and if those outputs have variance $q$ and are roughly independent, the stream’s variance after $L$ blocks is that of the initial vector plus $L q$. Normalising the branch input keeps the branch’s output at a controlled scale so that $q$ does not grow, but the sum still grows with depth unless the branch outputs become small.

This is the trade: post-normalisation bounds the stream’s scale and costs gradient conditioning; pre-normalisation preserves the gradient path and lets the stream’s scale grow with depth. Empirically the second is more trainable at depth, and its growth is why the final normalisation before the output head exists at all — the head needs a controlled input, and the stream’s own scale at the end of a deep stack is not controlled.

import torch
from torch import nn

class PreNormBlock(nn.Module):
    def __init__(self, d):
        super().__init__()
        self.ln1 = nn.LayerNorm(d)
        self.fc = nn.Sequential(nn.Linear(d, 4 * d), nn.GELU(), nn.Linear(4 * d, d))

    def forward(self, x):
        # normalise only the branch input; the identity path is untouched
        return x + self.fc(self.ln1(x))

x = torch.randn(2, 16, 64)
block = PreNormBlock(64)
for _ in range(48):
    x = block(x)
print(x.shape)

With random weights and a branch whose output has nonzero mean, the stream’s norm after many such blocks is larger than at the start; the growth is additive in variance, not multiplicative, which is why a deep pre-norm stack does not diverge exponentially.

Scale growth and the output head

Additive growth of variance by $L q$ is slow but not free. If $q$ is fixed by the branch’s architecture, the stream’s norm grows as $\sqrt{L}$, which for a stack of a hundred blocks is a factor of ten. The final normalisation removes the magnitude, but the direction of the stream becomes dominated by the accumulated sum, and the relative contribution of any single block’s update to the direction shrinks as $1/\sqrt{L}$. This is the residual stream’s characteristic behaviour: each block writes a small edit into a large shared vector, and the edits compose by addition.

Two mechanisms respond to this. The first is scaling the branch’s output by a constant less than one, so that the accumulated variance is bounded by a geometric rather than an arithmetic sum. The second is to make the branch’s output scale depend on its own normalisation’s gain, which is learned; a gain that shrinks with depth effectively down-weights late blocks’ contributions. Neither restores the $1/\sqrt{L}$ dilution of any single update, which is intrinsic to additive composition.

The gradient clipping discussion connects here: a stack whose stream norm grows is a stack whose loss gradients can be large in absolute terms even when the per-block contribution is small, because the readout at the head amplifies the stream’s norm. Clipping bounds the update, and the interaction with the normalisation’s placement determines which parameters receive the clipped direction.

Layer normalisation is not scale-free in the gradient

A frequent misreading of the invariance is that the layer’s gradient is unaffected by the input’s magnitude. The output is unaffected, and so is the loss, so the derivative of the loss with respect to a scale applied to $x$ before the layer is zero. But the derivative with respect to $x$ inside the layer’s Jacobian contains a factor $1/\sigma$, so a small input variance produces a large sensitivity. If a sublayer’s input variance is tiny — after a residual sum that cancels, or in a branch whose weights have shrunk — the normalised output is a heavily amplified version of the input’s direction, and gradient descent through that layer sees an ill-conditioned map. The $\epsilon$ in the denominator bounds the amplification at $1/\sqrt{\epsilon}$, which is why $\epsilon$ cannot be set arbitrarily small without risking large Jacobian entries in low-variance regimes.