Normalization and residual connections

Problem

Stack many layers and two things go wrong. First, the size of the activations drifts: each layer multiplies by weights, so the numbers shrink toward zero or blow up, and every layer must keep re-adapting to the shifting scale of the one below. Second, gradients travel back through every layer by the chain rule, a long product of factors, and a product of many numbers below 1 vanishes while one above 1 explodes (see backpropagation). Networks beyond a handful of layers simply fail to train.

Intuition

Residual connection: instead of a block computing y = F(x), let it compute y = x + F(x). The block now only learns a correction to its input, and the identity path gives the gradient a direct route: the derivative of x + F(x) with respect to x is 1 + F′(x), so even when F′ is tiny the gradient still flows. Stack blocks and you get a residual stream: one running vector per position that every block reads from and adds to.

Normalization: before (or after) a block, rescale each vector to a standard size: subtract a mean and divide by a standard deviation, then let two learned parameters (a scale γ and shift β) restore whatever range the next layer wants. The variants differ in which numbers share a mean and variance: all features of one example (LayerNorm), one feature across the whole batch (BatchNorm), or a group of channels in one example (GroupNorm).

In a transformer the "example" is one token position, and LayerNorm normalizes the E features of that position alone, so it never looks at other tokens or other sequences in the batch.

Mechanics

The residual stream

y = x + F(x)    shapes: x, F(x), y all [N, L, E]

F must return the same shape as x so the two can be added (when widths differ, a projection on the skip path fixes it). In a transformer block F is either attention or the feed-forward network; each adds its output into the stream. Reading the model this way, the stream is the model's working memory and each block writes a small update into it. Dropout is applied to F(x) before the addition, so a dropped block output leaves the identity path intact.

LayerNorm

y = (x − E[x]) / √(Var[x] + ε) · γ + β    statistics over the last dimension(s), x: [.., E], γ, β: [E]

The mean and (biased) variance are computed over the E features of each position independently, so each row ends up with zero mean and unit variance, then γ (initialized to 1) and β (initialized to 0) rescale it. ε (default 1e-5) keeps the division finite. Because nothing depends on other examples, LayerNorm behaves identically in training and evaluation, works with any batch size (even 1) and any sequence length, which is why transformers use it. It also makes the output invariant to scaling or shifting its input: multiply a row by 10 or add 100 to every entry and the result is unchanged.

RMSNorm

y = x / √(mean(x²) + ε) · γ

RMSNorm drops the mean subtraction (and β): it only divides by the root mean square. It is a little cheaper and, empirically, as good for many language models, so recent open-weight models commonly use it. Unlike LayerNorm it does not remove a constant offset, so the output keeps the direction of x exactly. If you leave eps=None PyTorch uses the machine epsilon of the computation type.

BatchNorm and the train/eval split

BatchNorm normalizes each feature (channel) over the batch (and spatial positions, in images): statistics are per column in the diagram above. In training it uses the current mini-batch's mean and variance and also updates two stored buffers by an exponential moving average:

running ← (1 − momentum) · running + momentum · batch statistic    momentum = 0.1 by default

(the running variance uses the unbiased batch variance, dividing by n − 1). In evaluation mode it does not look at the batch at all: it normalizes with the stored running mean and variance, so one example gives the same answer whatever else is in the batch. This is why model.eval() matters, why BatchNorm struggles with tiny batches (noisy statistics), and why it is awkward for variable-length sequences. The module's state is weight (γ), bias (β), and buffers running_mean, running_var, num_batches_tracked; the buffers are saved in the state dict but are not trained by gradient descent.

GroupNorm and others

GroupNorm splits the C channels of one example into G groups and normalizes each group on its own, so it needs no batch and no running statistics. With G = 1 it is a LayerNorm over all channels; with G = C every channel is normalized alone (instance norm). It requires num_channels % num_groups == 0. It is the usual choice for convolutional models with small batches.

Pre-norm versus post-norm

Where the normalization sits relative to the residual addition changes training behavior:

In TransformerEncoderLayer the switch is the norm_first argument (default False, i.e. post-norm). The code's own comment points to Fig. 1 of the "On Layer Normalization in the Transformer Architecture" paper for the difference.

Green is the skip path. In pre-norm it carries x from the input to the output without passing through any normalization, which is what keeps gradients healthy at depth.

The feed-forward block and dropout

The other half of a transformer block is a position-wise two-layer network applied to each position separately: linear2(dropout(activation(linear1(x)))), expanding E to dim_feedforward (commonly 4E) and back, with ReLU by default (GELU or SwiGLU in many newer models). Dropout zeroes each activation with probability p during training and rescales the rest by 1/(1 − p); in evaluation mode it does nothing. In TransformerEncoderLayer it is applied inside the feed-forward block, to the attention output (dropout1) and to the feed-forward output (dropout2), each before its residual addition.

Worked example

The input bars are drawn at the same scale as the outputs, so the first chart is simply tall.

Implementation

import torch
import torch.nn as nn

x = torch.randn(8, 16, 64)                        # [batch, positions, features]

ln = nn.LayerNorm(64)                             # weight [64] = 1, bias [64] = 0
y = ln(x)                                         # per position: mean 0, variance 1
rms = nn.RMSNorm(64)                              # weight only
bn = nn.BatchNorm1d(64)                           # expects [N, C] or [N, C, L]
y = bn(x.transpose(1, 2))                         # statistics per feature over batch and positions
bn.running_mean, bn.running_var, bn.num_batches_tracked
bn.eval()                                         # now uses the stored statistics
gn = nn.GroupNorm(num_groups=8, num_channels=64)  # 64 % 8 == 0

# a transformer sublayer, both styles
class Block(nn.Module):
    def __init__(self, dim, pre_norm=True, p=0.1):
        super().__init__()
        self.norm1, self.norm2 = nn.LayerNorm(dim), nn.LayerNorm(dim)
        self.attn = nn.MultiheadAttention(dim, 4, batch_first=True)
        self.ff = nn.Sequential(nn.Linear(dim, 4 * dim), nn.GELU(), nn.Linear(4 * dim, dim))
        self.drop = nn.Dropout(p)
        self.pre_norm = pre_norm

    def forward(self, x):
        if self.pre_norm:                          # x = x + F(norm(x))
            h = self.norm1(x)
            x = x + self.drop(self.attn(h, h, h, need_weights=False)[0])
            return x + self.drop(self.ff(self.norm2(x)))
        x = self.norm1(x + self.drop(self.attn(x, x, x, need_weights=False)[0]))
        return self.norm2(x + self.drop(self.ff(x)))

layer = nn.TransformerEncoderLayer(64, 4, batch_first=True, norm_first=True)   # the built-in

LayerNorm and a residual block in plain Python:

import math

def layer_norm(row, gamma, beta, eps=1e-5):
    mean = sum(row) / len(row)
    var = sum((v - mean) ** 2 for v in row) / len(row)         # biased variance
    return [(v - mean) / math.sqrt(var + eps) * g + b for v, g, b in zip(row, gamma, beta)]

def rms_norm(row, gamma, eps=1e-5):
    scale = 1 / math.sqrt(sum(v * v for v in row) / len(row) + eps)
    return [v * scale * g for v, g in zip(row, gamma)]

def block(x, sublayer, norm, pre_norm):
    if pre_norm:
        return [a + b for a, b in zip(x, sublayer(norm(x)))]
    return norm([a + b for a, b in zip(x, sublayer(x))])

Tradeoffs

• BatchNorm needs a decent batch. With a handful of examples the batch statistics are noisy; with variable-length sequences the padding corrupts them. It works well for image convolutions at moderate batch sizes and badly for most language models.
• Train/eval mismatch. Forgetting model.eval() leaves BatchNorm on batch statistics and Dropout on; predictions then vary with the batch and with chance. LayerNorm and RMSNorm have no such mode switch.
• Pre-norm trains more easily, post-norm can end slightly better. Pre-norm's clean skip path stabilises deep stacks, but the stream grows and later layers contribute relatively less; many large models therefore use pre-norm plus a final norm.
• RMSNorm saves a little. No mean, no bias: slightly fewer operations and parameters (4096 versus 8192 for E = 4096), with similar quality in practice, but it cannot remove a shared offset.
• Pitfalls. nn.LayerNorm(64) normalizes only the last dimension, so giving it an image tensor [N, C, H, W] with 64 normalizes over W. Using BatchNorm1d on a [N, L, E] tensor without transposing normalizes the wrong axis (it expects the feature dimension second). Epsilon matters in half precision. Frozen BatchNorm layers during fine-tuning need .eval() on those layers specifically.
• When not to use it. A shallow network with a well-scaled input often needs neither a residual path nor normalization; add them when depth or instability calls for it.

Check yourself

Connections

References

• He, Zhang, Ren and Sun. Deep Residual Learning for Image Recognition. CVPR, 2016 (arXiv 1512.03385).
• Ioffe and Szegedy. Batch Normalization: Accelerating Deep Network Training by Reducing Internal Covariate Shift. ICML, 2015 (arXiv 1502.03167).
• Ba, Kiros and Hinton. Layer Normalization. 2016 (arXiv 1607.06450).
• Zhang and Sennrich. Root Mean Square Layer Normalization. NeurIPS, 2019 (arXiv 1910.07467).
• Wu and He. Group Normalization. ECCV, 2018 (arXiv 1803.08494).
• Xiong et al. On Layer Normalization in the Transformer Architecture. ICML, 2020 (arXiv 2002.04745).
• Srivastava et al. Dropout: A Simple Way to Prevent Neural Networks from Overfitting. JMLR, 2014.
• PyTorch v2.15.0: normalization.py, batchnorm.py and TransformerEncoderLayer.