Weight initialization: Xavier/Glorot, He/Kaiming

The problem, from first principles

Vanishing and exploding gradients: why a deep, badly-initialized net fails to train, later in this track, names the failure mode directly: stack enough layers together, and if each layer's output is even slightly too large or too small relative to its input, that small multiplicative factor compounds across every layer, and by the time the signal reaches the last layer (on the forward pass) or the first layer (on the backward pass), it has either exploded to a huge number or vanished to essentially zero. Where you START the weights, before training even begins, is the single biggest lever for avoiding this: initialize every layer so that the VARIANCE of its output roughly matches the variance of its input, and that compounding effect never gets a foothold in the first place.

The exact right initialization scale depends on how many inputs a neuron has (fan_in, since summing more terms naturally increases variance) and, for some schemes, how many outputs (fan_out) and what nonlinearity comes right after the layer. Xavier/Glorot initialization (Glorot & Bengio, 2010) derived the right scale assuming a roughly linear activation (or tanh, which is close to linear near zero) and BOTH the forward AND backward pass need to preserve variance, hence it depends on both fan_in and fan_out. He/Kaiming initialization (He et al., 2015) derived a DIFFERENT scale specifically for ReLU, because ReLU zeroes out roughly half of its input (everything negative), which itself halves the output variance relative to a linear activation, so Kaiming compensates with extra variance to make up for that loss.

From theory to code

Implement four small functions that each compute a single scalar: xavier_uniform_bound, xavier_normal_std, kaiming_uniform_bound, and kaiming_normal_std. Each one returns either a uniform-distribution bound a (weights sampled from [-a, a]) or a normal-distribution standard deviation, given the layer's fan-in (and, for Xavier, fan-out).

Constraints

• xavier_uniform_bound(fan_in, fan_out) and xavier_normal_std(fan_in, fan_out) must use BOTH fan_in and fan_out.
• kaiming_uniform_bound(fan_in, gain) and kaiming_normal_std(fan_in, gain) must use ONLY fan_in, not fan_out, and must accept a gain parameter (default sqrt(2), the standard ReLU gain).
• All four must return a strictly positive float for any positive fan_in/fan_out.
• A uniform distribution on [-a, a] has variance a^2 / 3; use this relationship to derive each uniform bound from its corresponding variance target.

Hints