📝 Weight Initialization

Description

< What is it? >

Weight initialization is the choice of the values a network's parameters start at, before the first forward pass. Biases are almost always started at zero; the weights cannot be, and the distribution they are drawn from decides whether the very first gradients are usable.

The reason it matters is the same repeated multiplication that drives Vanishing & Exploding Gradients. A signal passing through $L$ layers is multiplied by $L$ weight matrices on the way forward and by $L$ Jacobians on the way back. If each layer shrinks the signal by a factor slightly below 1, the product decays geometrically with depth; if each layer amplifies it slightly, the product blows up. Initialization sets that per-layer factor at step 0.

< The two degenerate cases >

Start	What happens
All zeros (or any constant)	Every unit in a layer computes the same output and receives the same gradient, so they stay identical forever. The layer has the capacity of a single unit — this is the symmetry problem, and it is why weights must be random.
Too large / too small	Activation variance grows or decays layer by layer. With sigmoid / tanh the units saturate and their derivatives go to zero; with ReLU the activations either explode or die out.

Random initialization breaks symmetry; the scale of that randomness is what keeps the signal alive through depth.

Key points

< The variance argument >

For a linear layer $z = Wx$ with $n_{\text{in}}$ inputs, independent zero-mean weights of variance $\operatorname{Var}(W)$ and inputs of variance $\operatorname{Var}(x)$:

$$\operatorname{Var}(z) = n_{\text{in}} \cdot \operatorname{Var}(W) \cdot \operatorname{Var}(x).$$

Holding $\operatorname{Var}(z) \approx \operatorname{Var}(x)$ — so that neither the forward activations nor the backward gradients change scale with depth — requires $\operatorname{Var}(W) \approx 1 / n_{\text{in}}$. Every standard scheme is a version of this fraction, adjusted for what the activation function does to the variance.

< The standard schemes >

Scheme	Variance	Use with	Why
Xavier([ˈzævɪr]) / Glorot	$\dfrac{2}{n_{\text{in}} + n_{\text{out}}}$	tanh, sigmoid, linear	Averages the forward ($n_{\text{in}}$) and backward ($n_{\text{out}}$) constraints, so activations and gradients keep roughly constant variance. Assumes an activation that is near-linear around 0.
He / Kaiming	$\dfrac{2}{n_{\text{in}}}$	ReLU, Leaky ReLU, GELU	ReLU zeroes about half its inputs and so halves the variance; the factor 2 compensates. Using Xavier with ReLU makes activations decay with depth.
LeCun	$\dfrac{1}{n_{\text{in}}}$	SELU, self-normalizing nets	Forward-only constraint; the fixed point SELU's self-normalizing property is derived from.
Orthogonal	— (orthogonal matrix, optionally scaled by a gain)	deep RNNs, very deep stacks	An orthogonal $W$ preserves vector norms exactly, so repeated multiplication neither shrinks nor grows the signal.

Each comes in a normal and a uniform flavor: draw from $\mathcal{N}(0, \sigma^2)$ with $\sigma^2$ from the table, or from $\mathcal{U}(-a, a)$ with $a = \sqrt{3\sigma^2}$ (which gives the same variance).

• fan_in is the number of inputs to the layer, fan_out the number of outputs. For a conv layer, fan_in = in_channels × kernel_height × kernel_width.
• The choice follows the activation function, not the task. The one question worth asking is "does this layer's nonlinearity kill part of the signal?" — if it does (ReLU family), use He.

< In practice >

import torch.nn as nn

layer = nn.Linear(512, 512)

# ReLU-family activations -> He / Kaiming
nn.init.kaiming_normal_(layer.weight, nonlinearity='relu')

# tanh / sigmoid -> Xavier / Glorot
nn.init.xavier_uniform_(layer.weight, gain=nn.init.calculate_gain('tanh'))

nn.init.zeros_(layer.bias)

• Framework defaults are not always the right scheme. PyTorch's nn.Linear and nn.Conv2d default to a Kaiming-uniform variant with a=√5, which is a historical choice rather than a match for your activation; setting the initializer explicitly is cheap insurance for a deep model.
• Biases start at zero. Exceptions: LSTM forget-gate biases are often set to 1 so the cell remembers by default, and a final-layer bias can be set to the log-odds of the base rate to save the model from learning the class prior.
• Residual branches are often started at zero (a zero-init final BatchNorm $\gamma$, or a learnable scalar at zero), so a deep network starts as the identity and gains depth as training proceeds.
• Normalization reduces the sensitivity but does not remove it. BatchNorm/LayerNorm rescale activations after the fact, so a badly scaled init hurts less — but the backward pass at step 0, and networks without normalization, still depend on it.
• Transformers typically use a small fixed standard deviation (e.g. $0.02$) and scale the residual-projection weights by $1/\sqrt{2L}$ so the residual stream's variance does not grow with layer count.

< How to tell it is wrong >

Run one forward and one backward pass on a single batch, and print per-layer statistics:

• Activation standard deviation shrinking steadily with depth → scale too small (or Xavier used with ReLU).
• Activation standard deviation growing with depth, or inf / NaN at step 0 → scale too large.
• tanh / sigmoid outputs clustered at $\pm 1$ / $0$ and $1$ → saturated from the start.
• A large share of ReLU units at exactly 0 across all batches → dead units.
• Early-layer gradient norms orders of magnitude below later ones → the backward signal is decaying.

< Related ideas >

• Vanishing & Exploding Gradients — the failure mode initialization is chosen to avoid.
• Data Normalization — the same variance argument applied to the network's input layer.
• Gradient Descent — what happens once the first gradients arrive.

Implementation

< Notebook >

Open the weight-initialization notebook on GitHub

< Weight-matrix dimensions >

Q: Why is W₁ shaped (2, 3) and W₂ shaped (1, 2)?

layers_dims = [3, 2, 1] describes this network:

$$3\ \text{input features} \quad \longrightarrow \quad 2\ \text{hidden neurons} \quad \longrightarrow \quad 1\ \text{output neuron}.$$

Each row of a weight matrix belongs to one neuron in the current layer; each column corresponds to one input from the previous layer.

Matrix	Connects	Shape
$W_1$	3 input features to 2 hidden neurons	$(2, 3)$
$W_2$	2 hidden activations to 1 output neuron	$(1, 2)$

For the first layer, $x \in \mathbb{R}^{3 \times 1}$ and $b_1 \in \mathbb{R}^{2 \times 1}$:

$$Z_1 = W_1x + b_1 \qquad\Rightarrow\qquad (2 \times 3)(3 \times 1) + (2 \times 1) = (2 \times 1).$$

The result has one value for each hidden neuron. The second layer receives those two activations:

$$Z_2 = W_2A_1 + b_2 \qquad\Rightarrow\qquad (1 \times 2)(2 \times 1) + (1 \times 1) = (1 \times 1).$$

In general:

$$\boxed{ W_l.\text{shape} = (\text{neurons in layer } l,\ \text{neurons in layer } l - 1) }$$

W_l.shape = (layers_dims[l], layers_dims[l - 1])

Crash course

Weight Initialization Explained — Xavier, He & Why It Matters

Reference

• Symmetry Breaking versus Zero Initialization (Paul Mielke)
• Understanding the difficulty of training deep feedforward neural networks — Glorot & Bengio (2010)
• Delving Deep into Rectifiers — He et al. (2015)
• Exact solutions to the nonlinear dynamics of learning in deep linear networks — Saxe et al. (2014)
• Self-Normalizing Neural Networks — Klambauer et al. (2017)
• torch.nn.init — PyTorch documentation