Weight initialization

Problem

A deep network multiplies its input by one weight matrix after another, with an activation function in between. If each layer makes its output a little larger than its input, say by a factor of 1.5, twenty layers make it 1.5²⁰ ≈ 3300 times larger. If each makes it a little smaller, the signal fades to nothing. Backpropagation multiplies by the same matrices on the way back, so the gradients suffer the same fate: the early layers get gradients that are vanishingly small or enormous, and training stalls or diverges before it starts.

We cannot set the weights to zero either. If all weights in a layer start equal, every unit computes the same function and gets the same gradient, so they stay equal forever. The weights must be random to break the symmetry, and their scale must keep the signal's size roughly constant from layer to layer.

Intuition

A unit adds up n inputs, each multiplied by a random weight. Sums of n independent random terms have variance proportional to n, so a layer with 1000 inputs produces outputs about √1000 ≈ 32 times larger than a layer with one input, for the same weight scale. To cancel that, scale the weights by 1/√n. Then correct for what the activation does: ReLU zeroes half its inputs, which halves the average squared value, so multiply the variance back up by 2. That is Kaiming initialization: weight variance 2/n.

The whole page in one line: keep the per-layer factor at 1.

Mechanics

Forward pass

Take one pre-activation z = Σ_i=1..n w_i h_i, with the weights independent of the inputs, zero-mean and of variance Var(w). Each term has mean zero and, using the rules from expectation and variance, variance Var(w) · E[h²]. The terms are uncorrelated, so

Var(z) = n · Var(w) · E[h²]

It is E[h²], not Var(h), because ReLU outputs are not zero-mean. Now z is symmetric around zero, and ReLU keeps the positive half, so E[relu(z)²] = Var(z)/2. Chaining the two,

E[h′²] = (n · Var(w) / 2) · E[h²]     ⇒     after L layers: E[h_L²] = Π_l (n_l Var(w_l) / 2) · E[h₀²]

A product of L factors is stable only if each factor is 1, which gives He et al.'s condition Var(w) = 2/n_in. For an activation that is close to the identity near zero, like tanh, the ReLU factor ½ is replaced by about 1, which gives Var(w) = 1/n_in.

Backward pass

The gradient with respect to a layer's input is W^⊤ times the gradient at its output, times the activation's slope. The same argument now sums over the n_out units the input feeds, and ReLU's slope is 1 for half of them and 0 for the rest:

Var(∂L/∂h) = (n_out · Var(w) / 2) · Var(∂L/∂h′)

So the backward pass wants Var(w) = 2/n_out. When a layer is square the two agree; when it is not, you pick one (mode="fan_in" or "fan_out" in PyTorch). Glorot and Bengio, deriving for activations in their linear regime, split the difference by using the average of the two fans, Var(w) = 2/(n_in + n_out).

The named schemes

A uniform distribution on [−a, a] has variance a²/3, so each scheme comes in a normal and a uniform flavour with bound a = √3 · std. PyTorch writes every scheme as a gain divided by the square root of a fan:

• Xavier / Glorot: std = gain · √(2 / (fan_in + fan_out)); uniform bound gain · √(6 / (fan_in + fan_out)).
• Kaiming / He: std = gain / √fan, with fan = fan_in by default; uniform bound gain · √(3 / fan).
• Gains from calculate_gain: 1 for linear, convolution and sigmoid; 5/3 for tanh; √2 for ReLU; √(2 / (1 + slope²)) for leaky ReLU (slope 0.01 by default); 3/4 for SELU, a value the source says was found empirically.
• Fans: for a weight of shape (out, in, k₁, k₂, …), fan_in = in · k₁k₂… and fan_out = out · k₁k₂…, so a 3×3 convolution with 64 input channels has fan_in 576. Note that a Linear weight is stored as (out_features, in_features).

What PyTorch does by default

nn.Linear.reset_parameters calls kaiming_uniform_(weight, a=math.sqrt(5)). With the default nonlinearity="leaky_relu" and slope √5, the gain is √(2 / (1 + 5)) = √(1/3), so

bound = √3 · gain / √fan_in = √3 · √(1/3) / √fan_in = 1 / √fan_in     (fan_in = 64: 0.125)

The source comment says exactly this, and the class docstring documents the weights as U(−√k, √k) with k = 1/in_features. The bias uses the same bound. The convolution layers (_ConvNd) do the same. The weight variance is 1/(3 fan_in), so in a plain ReLU stack the per-layer factor is fan_in · 1/(3 fan_in) / 2 = 1/6: the default is not Kaiming for ReLU, and the signal shrinks layer by layer. Modern networks mostly get away with it because normalization layers and residual connections re-establish the scale at every block. nn.Embedding uses normal_, N(0, 1).

Orthogonal, truncated and residual scaling

• Orthogonal (orthogonal_, Saxe et al.): draw a Gaussian matrix, take its QR decomposition, fix the signs of Q's columns by the signs of R's diagonal (so Q is uniformly distributed), and multiply by gain. Every singular value is exactly 1, so a deep linear network preserves the norm of every input exactly, not just on average. It is common for recurrent weight matrices, which are applied many times.
• Truncated normal (trunc_normal_(mean=0, std=1, a=−2, b=2)): redraws values outside [a, b], to avoid rare huge weights. The cutoffs are absolute values, not multiples of std; with std=0.02 the defaults truncate nothing in practice.
• Residual scaling: in a residual network x + f(x), each block adds variance, so the stream grows with depth. GPT-2 scaled the weights of the layers feeding the residual stream by 1/√N, with N the number of residual layers, so their sum stays about the same size.

Worked example

Each scheme is a straight line on a log scale with slope log(factor)/2; only a factor of 1 is flat.

Implementation

import math
import torch.nn as nn

def init_weights(module):
    if isinstance(module, (nn.Linear, nn.Conv2d)):
        nn.init.kaiming_normal_(module.weight, mode="fan_in", nonlinearity="relu")
        if module.bias is not None:
            nn.init.zeros_(module.bias)

model = nn.Sequential(*[m for _ in range(20) for m in (nn.Linear(512, 512), nn.ReLU())])
model.apply(init_weights)            # visits every submodule

# tanh layers: Xavier with the tanh gain
nn.init.xavier_uniform_(tanh_layer.weight, gain=nn.init.calculate_gain("tanh"))

# recurrent weights: orthogonal
nn.init.orthogonal_(rnn.weight_hh_l0)

# GPT-2-style: scale the projections that write into the residual stream
n_residual_layers = 2 * len(blocks)   # attention and MLP each add to the stream
for block in blocks:
    for proj in (block.attn_out, block.mlp_out):
        proj.weight.data /= math.sqrt(n_residual_layers)

The functions end in _ because they fill the tensor in place, under torch.no_grad(), and return it. They take an optional generator for reproducible draws. When you load pretrained weights, initialization only matters for the layers you add.

Tradeoffs

• Match the activation. Kaiming for ReLU-family, Xavier (with the right gain) for tanh and linear layers. The derivations assume zero-mean, independent weights and inputs; they are a starting point, not a guarantee.
• fan_in or fan_out. fan_in keeps activations stable, fan_out keeps gradients stable. For square layers they agree. Otherwise, choosing one leaves the other pass off by Π n_l+1/n_l, which telescopes to the ratio of the last width to the first, so neither choice compounds with depth.
• Normalization makes it less critical. With LayerNorm or BatchNorm after each layer the activation scale is reset every block, which is why the PyTorch default rarely causes trouble in modern models. Plain deep stacks without normalization are where initialization decides whether training works at all.
• Average, not exact. The variance rule holds on average over random weights. Narrow layers fluctuate, as the worked example shows; orthogonal initialization removes that fluctuation for linear maps.
• Scale still interacts with the optimizer. The optimizer and learning rate decide how far weights move from the starting point; a good initialization only gives them a well-scaled place to start.

Connections

References

• Glorot and Bengio. Understanding the difficulty of training deep feedforward neural networks. AISTATS 2010.
• He, Zhang, Ren and Sun. Delving Deep into Rectifiers: Surpassing Human-Level Performance on ImageNet Classification. ICCV 2015. arXiv:1502.01852
• Saxe, McClelland and Ganguli. Exact solutions to the nonlinear dynamics of learning in deep linear neural networks. 2013. arXiv:1312.6120
• Radford et al. Language Models are Unsupervised Multitask Learners (GPT-2), 2019, section 2.3 (residual scaling).
• PyTorch v2.15.0: torch/nn/init.py, torch/nn/modules/linear.py.