Kaiming Init Derivation

The Kaiming (He) initialization derivation constructs a variance-preserving weight distribution for ReLU networks, concluding $\text{Var}(w) = 2/d_{\text{in}}$. Full derivation chain: ① Core identity: for a single output neuron $y = \sum_{i=1}^{d} w_i x_i$ ($d = d_{\text{in}}$), with $w_i$ i.i.d., $E[w_i] = 0$ and $w_i$ independent of $x_i$, $\text{Var}(y) = \sum_{i=1}^{d}\text{Var}(w_i x_i) = \sum_{i=1}^{d}\left(E[w_i^2]E[x_i^2] - E[w_i]^2E[x_i]^2\right)$. Since zero-mean $w$, $x$ give $\text{Var}(w)=E[w^2]$, $\text{Var}(x)=E[x^2]$, we get $\text{Var}(y) = d\,E[w^2]\,E[x^2] = d\,\text{Var}(w)\,\text{Var}(x)$. ② Linear case: without an activation (or with a linear one), setting $\text{Var}(y) = \text{Var}(x)$ yields $d\,\text{Var}(w) = 1 \Rightarrow \text{Var}(w) = 1/d$, which is Xavier's forward condition. ③ The role of ReLU: assume the pre-activation $z = w^Tx$ is zero-mean and symmetric about 0 (guaranteed by zero-mean symmetric $w$, $x$), so $P(z<0) = 1/2$; the output $x' = \max(0, z)$ zeroes the negative half, and its mean-square energy is $E[(x')^2] = \int_0^{\infty} z^2 p(z)\,dz = \tfrac{1}{2}\int_{-\infty}^{\infty} z^2 p(z)\,dz = \tfrac{1}{2}\text{Var}(z)$ — ReLU exactly halves the activation energy. ④ Feed into the next layer: $\text{Var}(y') = d\,E[w^2]\,E[(x')^2] = d\,\text{Var}(w)\cdot\tfrac{1}{2}\text{Var}(x)$ (the previous layer already keeps $\text{Var}(z) = \text{Var}(x)$); setting $\text{Var}(y') = \text{Var}(x)$ gives $\tfrac{1}{2}d\,\text{Var}(w) = 1 \Rightarrow \text{Var}(w) = \frac{2}{d}$. Where the factor 2 comes from: ReLU halves the activation energy each layer, so the weight variance must double to preserve variance. ⑤ Backward condition: the gradient $\frac{\partial L}{\partial x} = W^T\frac{\partial L}{\partial y}$ is also a sum of $d_{\text{out}}$ terms, and ReLU's derivative is 1 on the positive half and 0 on the negative half (gradient pass-through is also ½), so backward variance preservation requires $\text{Var}(w) = 2/d_{\text{out}}$. Forward uses fan_in, backward uses fan_out; frameworks such as PyTorch default to fan_in ($d_{\text{in}}$). ⑥ PReLU generalization (negative slope $a$): $\text{Var}(w) = \frac{2}{(1+a^2)d_{\text{in}}}$; $a=0$ recovers ReLU's $2/d_{\text{in}}$.

a top interview derivation question (“derive the factor 2 in Kaiming init”), and the theoretical basis for the default initialization (kaiming_uniform) of all ReLU deep networks such as ResNet.

without the factor-2 compensation, each layer multiplies forward activation variance and backward gradient variance by ½, so after 50 layers signals vanish at $(1/2)^{50} \approx 10^{-15}$ — vanishing gradients; Kaiming keeps every layer's variance at 1, so signals traverse dozens or hundreds of layers stably. It also avoids ReLU saturation (weights too large push inputs into the negative half where gradients are 0 — dead neurons) and keeps the initial output magnitude at $O(1)$, matching the loss scale for stable first updates.

Full derivation: from $\text{Var}(y) = d\,E[w^2]\,E[x^2]$, through ReLU zeroing, derive $\text{Var}(w) = 2/d_{\text{in}}$ (write every step, including the zero-mean and independence assumptions).

Derive the factor 2: use the integral $E[(x')^2] = \tfrac{1}{2}\text{Var}(z)$ to show how ReLU zeroing halves the activation energy; which assumptions does it rely on (zero mean, symmetric distribution, $w$ independent of $x$)?

Why are the forward and backward variance-preserving conditions $2/d_{\text{in}}$ and $2/d_{\text{out}}$ respectively? How does ReLU's ½ gradient pass-through enter the backward derivation? Which does the framework default to (fan_in)?

Kaiming vs Xavier: how does $\text{Var}(w) = 2/d$ reduce to $1/d$ without ReLU? Why do tanh networks use Xavier while ReLU networks must use Kaiming (what goes wrong with $1/d$)?

PReLU generalization: how does $\text{Var}(w) = \frac{2}{(1+a^2)d_{\text{in}}}$ follow from $E[(x')^2] = \tfrac{1+a^2}{2}\text{Var}(z)$; if the assumptions break (asymmetric or nonzero-mean inputs), which step of the derivation fails?