Initialization and normalization

Two failure modes at step zero

Symmetry. Initialize every weight to the same value and every hidden unit in a layer computes the same function, receives the same gradient (4.4 Backpropagation, fully derived) and stays identical forever: the layer has one effective unit. Weights must be random. Biases can be zero.

Scale. Too small, and activations shrink layer by layer until the output is constant and gradients vanish; too large, and sigmoid/tanh units saturate (ReLU units explode). The question "how large should random weights be?" has an exact answer.

Variance propagation: deriving Xavier and He

Consider one pre-activation $z_j = \sum_{i=1}^{n_{\text{in}}} w_{ji}x_i$ with weights and inputs independent, zero-mean. For independent zero-mean variables, $\mathrm{Var}(w_{ji}x_i)=\mathrm{Var}(w)\,\mathrm{Var}(x)$ (1.6 Probability I: random variables and distributions), and variances of independent terms add:

To keep the forward signal at constant scale we need $\mathrm{Var}(z)=\mathrm{Var}(x)$, i.e. $\mathrm{Var}(w)=1/n_{\text{in}}$ — LeCun initialization. The backward pass multiplies by $\mathbf{W}^\top$ (4.4 Backpropagation, fully derived), so by the same argument gradient variance is preserved when $\mathrm{Var}(w)=1/n_{\text{out}}$. Xavier/Glorot initialization compromises with the harmonic mean,

which is right for tanh/sigmoid units near their linear regime. ReLU zeroes half of every symmetric distribution, so $\mathbb{E}[\mathrm{ReLU}(z)^2]=\tfrac12\mathrm{Var}(z)$ and the signal halves per layer unless the weights compensate. He (Kaiming) initialization doubles the variance:

Math check. A 100-layer ReLU network with $\mathrm{Var}(w)=1/n_{\text{in}}$ shrinks the activation variance by $\tfrac12$ per layer: $2^{-100}\approx 10^{-30}$ by the output. With He initialization the factor is $1$ per layer. The entire difference between "trains" and "does nothing" is a factor of 2 in a variance.

Three practical notes. PyTorch's default nn.Linear init is kaiming_uniform_ with a=√5, which works out to $\mathrm{Var}(w)=1/(3n_{\text{in}})$ — conservative, and a historical quirk; set your initialization explicitly for anything deep. Residual networks add a further rule: scale the last layer of each residual branch down (GPT-2 multiplies the output projections by $1/\sqrt{2L}$ for $L$ blocks; many ResNets zero-initialize the last BatchNorm gain in each block) so that the sum of $L$ branches does not blow up the residual stream. And check your work: with correct initialization the initial loss of a $K$-way classifier should be $\approx\ln K$ (4.8 Training dynamics and debugging).

Batch normalization

BatchNorm (Ioffe & Szegedy, 2015) standardizes each feature over the mini-batch and then lets the network undo it with learnable parameters. For a feature with batch values $x_1,\dots,x_m$:

For convolutions the statistics are per channel, pooled over batch and spatial positions. At inference there is no batch to average over, so BatchNorm uses running estimates of $\mu$ and $\sigma^2$ accumulated during training (momentum 0.1 in PyTorch) — which is what model.eval() switches on. The bias of a layer immediately before BatchNorm is redundant (it is subtracted off) and is normally omitted.

Worked example. Batch values $(1, 2, 3, 6)$: $\mu=3$, $\sigma^2=\tfrac{4+1+0+9}{4}=3.5$, $\sigma\approx1.871$; normalized $(-1.069,\,-0.535,\,0,\,1.604)$; with $\gamma=2,\ \beta=1$: $(-1.14,\,-0.07,\,1.00,\,4.21)$.

Why it helps. The original story — reducing "internal covariate shift" — did not survive scrutiny. The better-supported explanation (Santurkar et al., 2018) is that BatchNorm makes the loss landscape smoother and the gradients more predictable, which allows learning rates 5–10× higher, makes training far less sensitive to initialization, and adds a mild regularizing noise because every example's normalization depends on its batch-mates.

Where it breaks. Small batches (noisy statistics; below 8–16 per device it degrades badly); sequence models with variable lengths and padding; train/test mismatch when the running statistics drift from the batch statistics; fine-tuning with frozen statistics; reinforcement learning with shifting data distributions; and distributed training, where per-device batches need synchronized BatchNorm to see the global statistics. Combining BatchNorm with dropout also interacts badly (variance shift). These failure cases are exactly why transformers do not use it.

LayerNorm, RMSNorm and GroupNorm

LayerNorm (Ba, Kiros & Hinton, 2016) normalizes across the features of a single example instead of across the batch:

No batch dependence, no running statistics, identical behaviour at training and inference, works for any sequence length — which is why it is the normalization inside transformers (7.4 Inside the block: residuals, LayerNorm, FFN and embeddings), applied to each token's $d$-dimensional vector independently.

RMSNorm (Zhang & Sennrich, 2019) drops the mean subtraction and the $\beta$:

It is cheaper (one reduction instead of two) and performs as well, so Llama, Mistral, Gemma, Qwen and DeepSeek all use it. Geometrically it projects each token vector onto a sphere of radius $\sqrt d$ and lets $\boldsymbol\gamma$ rescale per dimension — the norm-division idea of 1.1 Vectors, dot products and norms.

GroupNorm splits channels into groups and normalizes within each group per example: the right choice for vision models trained with tiny batches (detection, segmentation; 5.5 Object detection, 5.6 Segmentation). InstanceNorm is GroupNorm with one channel per group and appears in style transfer.

Pre-norm, post-norm and the residual stream

A transformer block adds a sub-layer's output to its input: $\mathbf{x}+\mathrm{Sub}(\mathbf{x})$. The original transformer applied normalization after the addition (post-norm, $\mathrm{LN}(\mathbf{x}+\mathrm{Sub}(\mathbf{x}))$), which normalizes the residual stream itself; deep post-norm stacks need careful warmup and still train unstably. Modern models use pre-norm, $\mathbf{x}+\mathrm{Sub}(\mathrm{LN}(\mathbf{x}))$: the residual stream is never rescaled, each sub-layer sees a normalized input, and a final norm is applied before the output head (7.4 Inside the block: residuals, LayerNorm, FFN and embeddings).

The reason pre-norm is stable is the residual connection itself. With $\mathbf{h}^{(\ell+1)} = \mathbf{h}^{(\ell)} + F_\ell(\mathbf{h}^{(\ell)})$,

so there is always an identity path along which the gradient reaches early layers undiminished, no matter how many layers sit in between — the Jacobian product of 4.2 Activation functions no longer has to stay near 1 on its own. Residual connections are the third pillar, with initialization and normalization, that makes depth work (5.3 Architectures: LeNet → AlexNet → VGG → Inception → ResNet → EfficientNet → ConvNeXt).

In code

import torch.nn as nn, math
def init_weights(m):
    if isinstance(m, nn.Linear):
        nn.init.kaiming_normal_(m.weight, nonlinearity="relu")   # He: Var = 2/n_in
        nn.init.zeros_(m.bias)
model.apply(init_weights)
# scale residual output projections in a GPT-style model with L blocks
for name, p in model.named_parameters():
    if name.endswith("proj.weight"):
        nn.init.normal_(p, mean=0.0, std=0.02 / math.sqrt(2 * L))

Where this shows up later

• Checking the initial loss and activation histograms is the first debugging step — 4.8 Training dynamics and debugging; regularization interacts with normalization (weight decay on scale-invariant layers) — 4.7 Regularization in deep nets.
• BatchNorm inside ResNets and GroupNorm in detectors — 5.3 Architectures: LeNet → AlexNet → VGG → Inception → ResNet → EfficientNet → ConvNeXt, 5.5 Object detection; LayerNorm/RMSNorm and pre-norm in the transformer block — 7.4 Inside the block: residuals, LayerNorm, FFN and embeddings; RMSNorm in the Llama recipe — 7.6 Decoder-only: GPT and autoregressive generation.
• Normalizing embeddings before cosine similarity — 9.2 Embeddings deep dive; the fixed-scale residual stream that interpretability research reads — 8.11 Interpretability: what is happening inside.