Weight initialization

The problem

A neural network needs parameter values before its first update. Setting every hidden unit to identical values preserves symmetry: identical units receive identical gradients and remain copies.

Random values break symmetry, but scale matters. Repeated multiplication by weights that are too small shrinks activations and gradients. Weights that are too large can amplify them or push sigmoid and tanh units into saturated regions.

Work through a small example

A neuron sums fan_in inputs. Under the simplifying assumptions that inputs are independent and zero-mean, weights are independent and zero-mean, and inputs and weights are independent of each other, unit input variance and weight variance s^2 give pre-activation variance approximately fan_in*s^2. Choosing s^2=1/fan_in keeps this near one.

ReLU sets roughly half of symmetric inputs to zero. Kaiming initialization compensates with weight variance near 2/fan_in.

The general rule

Xavier initialization balances input and output widths and is commonly paired with tanh-like activations. Kaiming initialization accounts for a rectifier's gain and is commonly paired with ReLU. These are variance-preservation arguments under approximate independence and distribution assumptions.

Initialization gives optimization a reasonable starting signal scale. It does not guarantee convergence, prevent overfitting, or repair an unsuitable architecture.

Implement and inspect

import torch
from torch import nn

torch.manual_seed(4)
x = torch.randn(4096, 100)
layer = nn.Linear(100, 80, bias=False)
nn.init.kaiming_normal_(layer.weight, nonlinearity="relu")

pre = layer(x)
post = torch.relu(pre)
print(tuple(layer.weight.shape))
print(round(x.std().item(), 3))
print(round(pre.std().item(), 3))
print(round(post.std().item(), 3))

The weight shape is (80,100). Expected standard deviations for input, pre-activation, and post-ReLU are approximately 1.00, 1.41, and 0.82 in this seeded CPU example. Small differences across supported environments are normal.

Engineering checks

• Set and record a seed, but test conclusions across several seeds.
• Measure activation and gradient distributions layer by layer.
• Match initializer gain to the actual activation function.
• Do not accidentally reinitialize a loaded pretrained model.

Zero biases are usually safe because random weights already break hidden-unit symmetry. Zeroing every weight is the harmful case.

Go deeper

For ReLU, the Kaiming argument approximately preserves the second moment E[x²], not variance after the activation. ReLU outputs have a nonzero mean, so their standard deviation can be about 0.82 even when the relevant second moment is preserved. Correlations, residual branches, normalization, attention, and finite width also violate the simplest derivation. Deep networks may need architecture-specific scaling.

Dynamical isometry asks whether input-output Jacobian singular values remain near one, a stronger condition than stable activation variance. It helps explain why two initializers with similar variances can produce different optimization behavior.

Practice and recap

Question: A ReLU layer has fan_in=200. What Kaiming standard deviation does the simple rule suggest?

Initialization should break symmetry and keep signals usable. Verify the actual network rather than trusting a formula without measurements.