Xavier & Kaiming Initialization

The problem

Before training, every weight is set to a random number, and the size of those random numbers decides whether the network can train at all. A layer $y = xW$ with $n_{\text{in}}$ inputs adds up $n_{\text{in}}$ products, and so it multiplies the variance of its input by

where $\mathrm{Var}(w)$ is the variance of each weight (the average squared distance from 0, for weights with mean 0). Stack many layers and the factor compounds: above 1 the numbers explode, below 1 they vanish (the theory page shows both, measured). Initialization picks $\mathrm{Var}(w)$ so that the factor is about 1.

Example. $n_{\text{in}} = 512$. Weights with variance 1 make the factor 512: after 3 layers the variance has grown $512^3 \approx 134$ million times. Weights with variance $1/512$ make the factor exactly 1.

The weights have shape $(n_{\text{in}}, n_{\text{out}})$ and are used as x @ W. In the code, $n_{\text{in}}$ is fan_in and $n_{\text{out}}$ is fan_out.

Xavier (Glorot) uniform

Glorot and Bengio (2010) want both the forward activations and the backward gradients to keep their scale. The forward pass suggests $\mathrm{Var}(w) = 1/n_{\text{in}}$, the backward pass $1/n_{\text{out}}$, and the compromise is:

The weights are drawn uniformly between $-a$ and $a$. A uniform distribution on $[-a, a]$ has variance $a^2 / 3$, so to get the variance above:

Example. $n_{\text{in}} = n_{\text{out}} = 512$: $a = \sqrt{6/1024} \approx 0.0765$, and the variance is $0.0765^2 / 3 \approx 0.00195 = 2/1024$.

Kaiming (He) normal

ReLU (rectified linear unit) sets about half of its inputs to 0, which halves the variance at every layer. He et al. (2015) make up for it with a factor 2, using only $n_{\text{in}}$:

The weights are drawn from a normal (bell-shaped) distribution with mean 0 and this variance, so the standard deviation is $\sqrt{2 / n_{\text{in}}}$.

Example. $n_{\text{in}} = 512$: standard deviation $\sqrt{2/512} = 0.0625$.

What to implement

• xavier_uniform(fan_in, fan_out, rng): weights uniform on $[-a, a]$, with $a$ as above.
• kaiming_normal(fan_in, fan_out, rng): weights normal with mean 0 and variance $2 / n_{\text{in}}$.

Both return an array of shape (fan_in, fan_out). rng is a np.random.Generator: use rng.uniform and rng.normal.

Requirements

• Shape (fan_in, fan_out).
• Kaiming uses fan_in only.

Hint

rng.normal(mean, scale, size) takes the standard deviation as scale, not the variance: pass the square root.