Xavier and He initialisation: the variance by hand

Worked solutions

Problem 1

1. $\mathbb{E}[wx] = \mathbb{E}[w]\,\mathbb{E}[x] = 0$.Independence factorises the expectation of the product, and $\mathbb{E}[w] = 0$ kills it whatever $x$'s mean is.
2. $\operatorname{Var}(wx) = \mathbb{E}[w^2x^2] - 0^2 = \mathbb{E}[w^2]\,\mathbb{E}[x^2]$.The definition of variance with step 1's mean; $w^2$ and $x^2$ are functions of independent variables, so they are independent too.
3. $\mathbb{E}[w^2] = \operatorname{Var}(w) + (\mathbb{E}[w])^2 = \sigma_w^2$.The definition of variance rearranged, with a zero mean.
4. $\mathbb{E}[x^2] = \operatorname{Var}(x) + (\mathbb{E}[x])^2 = 4 + 1 = 5$.The same identity for $x$, whose mean is not zero, so its square stays.
5. $\operatorname{Var}(wx) = \sigma_w^2\,\mathbb{E}[x^2] = 0.25 \times 5 = 1.25$, while $\operatorname{Var}(w)\operatorname{Var}(x) = 1$The product of the variances misses $\sigma_w^2(\mathbb{E}[x])^2$: because $w$'s sign is random, even a constant input contributes spread to $wx$. The two agree only when $\mathbb{E}[x] = 0$, which a ReLU output never satisfies (Problem 3).

Problem 2

1. $\mathbb{E}[y_i] = \sum_j \mathbb{E}[W_{ij}]\,\mathbb{E}[x_j] = 0$.Linearity of expectation, then independence of $W$ from $x$, then the zero mean of the weights.
2. $\mathbb{E}[y_i^2] = \sum_j\sum_k \mathbb{E}[W_{ij}W_{ik}\,x_jx_k]$.Square the sum and use linearity; there are $n_{\text{in}}^2$ terms.
3. For $j \neq k$: $\mathbb{E}[W_{ij}W_{ik}\,x_jx_k] = \mathbb{E}[W_{ij}]\,\mathbb{E}[W_{ik}]\,\mathbb{E}[x_jx_k] = 0$.$W_{ij}$, $W_{ik}$ and the input are mutually independent. The cross terms vanish because of the weights' zero means, so correlated inputs do no harm.
4. For $j = k$: $\mathbb{E}[W_{ij}^2\,x_j^2] = \sigma_w^2\,q$.Problem 1, step 2, with $\mathbb{E}[W_{ij}^2] = \sigma_w^2$.
5. $\operatorname{Var}(y_i) = n_{\text{in}}\,\sigma_w^2\,q$, so $\sigma_w^2 = 1/n_{\text{in}}$ keeps $\mathbb{E}[y_i^2] = q$Only the $n_{\text{in}}$ diagonal terms survive, and with a zero mean the second moment is the variance. This is LeCun's initialisation. It preserves the scale of a linear network, and of a tanh network near $0$, where $\tanh t \approx t$. For $n_{\text{in}} = 256$ it is a standard deviation of $1/16$.

Problem 3

1. $\operatorname{relu}(z)^2 = z^2$ when $z > 0$ and $0$ when $z \le 0$.$\operatorname{relu}$ passes positive values and zeroes the rest.
2. $\mathbb{E}[z^2;\, z > 0] = \mathbb{E}[z^2;\, z < 0]$, where $\mathbb{E}[u;\, A]$ is the expectation of $u$ times the indicator of the event $A$.$-z$ has the same distribution as $z$, and $z \mapsto -z$ swaps the two events while leaving $z^2$ unchanged.
3. $\mathbb{E}[z^2;\, z > 0] + \mathbb{E}[z^2;\, z < 0] = \mathbb{E}[z^2] = \sigma^2$.The two events cover everything except $z = 0$, where $z^2 = 0$; and $\mathbb{E}[z] = 0$, so the second moment is the variance.
4. $\mathbb{E}[\operatorname{relu}(z)^2] = \sigma^2/2$.Steps 1 to 3: the left side is one of two equal halves of $\sigma^2$.
5. For $z \sim \mathcal{N}(0, \sigma^2)$: $\mathbb{E}[\operatorname{relu}(z)] = \dfrac{1}{\sigma\sqrt{2\pi}}\displaystyle\int_0^\infty t\,e^{-t^2/(2\sigma^2)}\,dt = \dfrac{\sigma^2}{\sigma\sqrt{2\pi}} = \dfrac{\sigma}{\sqrt{2\pi}}$.The integrand is the normal density times $t$ on $t > 0$, and $-\sigma^2 e^{-t^2/(2\sigma^2)}$ is an antiderivative of $t\,e^{-t^2/(2\sigma^2)}$.
6. $\mathbb{E}[\operatorname{relu}(z)^2] = \sigma^2/2$; for normal $z$, $\mathbb{E}[\operatorname{relu}(z)] = \sigma/\sqrt{2\pi}$ and $\operatorname{Var}(\operatorname{relu}(z)) = \sigma^2\big(\tfrac12 - \tfrac1{2\pi}\big) \approx 0.341\,\sigma^2$The variance is step 4 minus step 5 squared. The next layer needs the second moment, $\sigma^2/2$, not this variance (Problem 2), and the $\tfrac12$ is what He's factor $2$ undoes (Problem 4).

Problem 4

1. $\operatorname{Var}(y_i^{(l)}) = n_{l-1}\,\sigma_l^2\,\mathbb{E}\big[(h_j^{(l-1)})^2\big]$.Problem 2 with $x = h^{(l-1)}$, which is computed from earlier layers only and so is independent of $W^{(l)}$.
2. $\mathbb{E}\big[(h_j^{(l-1)})^2\big] = \tfrac12\operatorname{Var}(y_j^{(l-1)})$.Problem 3 applied to $z = y_j^{(l-1)}$, which is symmetric by hypothesis.
3. $\operatorname{Var}(y_i^{(l)}) = \tfrac12\,n_{l-1}\,\sigma_l^2\,\operatorname{Var}(y_j^{(l-1)})$.Substitute step 2 into step 1.
4. Negating row $i$ of $W^{(l)}$ negates $y_i^{(l)}$ and leaves the joint distribution of the weights and $h^{(l-1)}$ unchanged, so $y_i^{(l)}$ is symmetric.The weights are symmetric and independent of everything below them. The first layer's output is symmetric for the same reason, whatever the input, so the hypothesis holds at every layer from $l = 2$ on.
5. $\operatorname{Var}(y^{(l)}) = \tfrac12\,n_{l-1}\sigma_l^2\operatorname{Var}(y^{(l-1)})$; He initialisation $\sigma_l^2 = 2/n_{l-1}$ makes the factor exactly $1$The standard deviation is $\sqrt{2/n_{\text{in}}}$, which is $1/16$ for $n_{\text{in}} = 512$. PyTorch's kaiming_normal_ with nonlinearity='relu' uses this, as $\text{gain}/\sqrt{\text{fan\_in}}$ with gain $\sqrt2$; its default mode is 'fan_in'.

Problem 5

1. $\operatorname{Var}(y^{(l)}) = \tfrac c2\operatorname{Var}(y^{(l-1)})$ for $l \ge 2$.Problem 4, step 3, with $n_{l-1}\sigma_l^2 = c$.
2. $\operatorname{Var}(y^{(L)}) = \big(\tfrac c2\big)^{L-1}\operatorname{Var}(y^{(1)})$.Apply step 1 for $l = 2, \dots, L$: that is $L - 1$ factors. The first layer reads the raw input, not a ReLU output, so it is not one of them.
3. For $c = 1$, $L = 21$: $2^{-20} = 1/1048576 \approx 9.5 \times 10^{-7}$.$2^{10} = 1024$, so $2^{20} = 1024^2$.
4. For $c = 2.2$: $1.1^{20} \approx 6.7$.The same formula with $c/2 = 1.1$.
5. $\operatorname{Var}(y^{(L)})/\operatorname{Var}(y^{(1)}) = (c/2)^{L-1}$; for $c = 1$, $L = 21$ it is $2^{-20} \approx 9.5\times10^{-7}$, a standard deviation $1024$ times smallerThe standard deviation scales by the square root, $2^{-10}$. A constant factor per layer becomes exponential in depth, so only $c = 2$ neither vanishes nor explodes; $10\%$ too much variance per layer already multiplies it by $6.7$ over twenty layers.

Problem 6

1. $dx_j = \sum_{i=1}^{n_{\text{out}}} W_{ij}\,dy_i$.Entry $j$ of $W^\top dy$ combines column $j$ of $W$ with $dy$: one term per output.
2. $\mathbb{E}[dx_j] = 0$ and $\mathbb{E}[dx_j^2] = n_{\text{out}}\,\sigma_w^2\,\mathbb{E}[dy_i^2]$.Problem 2's argument with the roles swapped: the cross terms $\mathbb{E}[W_{ij}W_{kj}\,dy_i\,dy_k]$ vanish for $i \ne k$ because $W_{ij}$ and $W_{kj}$ are independent with mean $0$, and the sum now runs over $n_{\text{out}}$ terms.
3. The forward pass keeps its scale when $n_{\text{in}}\sigma_w^2 = 1$ (Problem 2), the backward pass when $n_{\text{out}}\sigma_w^2 = 1$.Both hold only if $n_{\text{in}} = n_{\text{out}}$.
4. $\sigma_w^2 = \dfrac{2}{n_{\text{in}} + n_{\text{out}}}$ gives $\tfrac12\big(n_{\text{in}}\sigma_w^2 + n_{\text{out}}\sigma_w^2\big) = 1$.Glorot and Bengio's compromise: the variance is the reciprocal of the average fan, so the two factors average to $1$.
5. $\operatorname{Var}(dx_j) = n_{\text{out}}\,\sigma_w^2\,\mathbb{E}[dy_i^2]$; for $784 \to 256$, Xavier gives $\sigma_w^2 = 2/1040 = 1/520$, a forward factor $784/520 \approx 1.51$ and a backward factor $256/520 \approx 0.49$Neither direction is preserved exactly; each is off by the ratio of one fan to the average. The derivation assumes linear units, so on a ReLU network both factors are halved again (Problems 5 and 7).

Problem 7

1. $\operatorname{relu}'(y_i^{(l)})$ is $1$ with probability $\tfrac12$ and $0$ otherwise.$y_i^{(l)}$ is symmetric about $0$ (Problem 4, step 4) and continuous, so it is positive half the time.
2. $\mathbb{E}[(dy_i^{(l)})^2] = \mathbb{E}[\operatorname{relu}'(y_i^{(l)})^2]\,\mathbb{E}[(dh_i^{(l)})^2] = \tfrac12\,\mathbb{E}[(dh_i^{(l)})^2]$.The two factors are treated as independent, and $0^2 = 0$, $1^2 = 1$, so the square of the mask is the mask.
3. $\mathbb{E}[(dh_j^{(l-1)})^2] = n_l\,\sigma_l^2\,\mathbb{E}[(dy_i^{(l)})^2]$.Problem 6, step 2: $W^{(l)}$ has $n_l$ rows, so the backward sum has $n_l$ terms.
4. $\mathbb{E}[(dh_j^{(l-1)})^2] = n_l\cdot\dfrac{2}{n_{l-1}}\cdot\tfrac12\,\mathbb{E}[(dh_i^{(l)})^2] = \dfrac{n_l}{n_{l-1}}\,\mathbb{E}[(dh_i^{(l)})^2]$.Steps 2 and 3 with $\sigma_l^2 = 2/n_{l-1}$: the ReLU's $\tfrac12$ cancels the $2$.
5. From $dh^{(4)}$ to $dh^{(1)}$: $\dfrac{n_4}{n_3}\cdot\dfrac{n_3}{n_2}\cdot\dfrac{n_2}{n_1} = \dfrac{n_4}{n_1} = \dfrac{128}{1024} = \dfrac18$.Step 4 for $l = 4, 3, 2$; the product telescopes.
6. $\mathbb{E}[(dh^{(l-1)})^2] = (n_l/n_{l-1})\,\mathbb{E}[(dh^{(l)})^2]$; from layer $4$ back to layer $1$ the factor is $n_4/n_1 = 1/8$; $\sigma_l^2 = 2/n_l$ (fan-out) makes each factor $1$With fan-in initialisation the backward pass is off only by the ratio of two widths, not by an exponential in depth, which is why the fan-in mode is a safe default. PyTorch exposes the other choice as mode='fan_out'.

Problem 8

1. Each output $y_{o,r,c}$ is a sum over $i = 1, \dots, 64$ and $a, b = 0, 1, 2$: $64 \times 3 \times 3 = 576$ products.Fan-in counts the products summed into one output, and a convolution sums over every input channel and every kernel position.
2. $\operatorname{Var}(y_{o,r,c}) = 576\,\sigma_w^2\,\mathbb{E}[x^2]$.Problem 2 with $n_{\text{in}} = 576$: the $576$ weights in one output's sum are distinct and independent. The same weights are reused at other positions, which correlates different outputs but does not change the variance of one.
3. An input $x_{i,r,c}$ away from the border appears in $128 \times 3 \times 3 = 1152$ outputs, each through a different weight.It sits at each of the $9$ kernel offsets of some window, for each of the $128$ output channels; the backward sum of Problem 6 has that many terms.
4. fan-in $= 576$, fan-out $= 1152$; He $\sigma_w^2 = 2/576 = 1/288$ (standard deviation $\approx 0.0589$); Xavier $\sigma_w^2 = 2/(576 + 1152) = 1/864$Both are channels times the receptive-field size $k_hk_w$, which is how PyTorch computes the fans of a convolution weight.

Problem 9

1. $f_\alpha(z)^2 = z^2$ for $z > 0$ and $\alpha^2z^2$ for $z < 0$.The leaky ReLU multiplies negative inputs by $\alpha$, so their squares by $\alpha^2$.
2. $\mathbb{E}[f_\alpha(z)^2] = \tfrac12\sigma_z^2 + \alpha^2\cdot\tfrac12\sigma_z^2 = \tfrac12(1 + \alpha^2)\,\sigma_z^2$.Problem 3, steps 2 and 3: each sign carries half of $\sigma_z^2$.
3. $\operatorname{Var}(y^{(l)}) = \tfrac12(1 + \alpha^2)\,n_{\text{in}}\,\sigma_w^2\operatorname{Var}(y^{(l-1)})$, so $\sigma_w^2 = \dfrac{2}{(1 + \alpha^2)\,n_{\text{in}}}$ and $g = \sqrt{2/(1 + \alpha^2)}$.Problem 4, step 3, with step 2 in place of the ReLU's $\tfrac12$. $\alpha = 0$ gives He's $\sqrt2$ and $\alpha = 1$, a linear unit, gives LeCun's $1$.
4. $\operatorname{Var}\big(U(-b, b)\big) = \displaystyle\int_{-b}^{b}\frac{t^2}{2b}\,dt = \frac{b^2}{3}$.The density is $1/(2b)$ on the interval and the mean is $0$; $\int_{-b}^b t^2\,dt = 2b^3/3$.
5. $b^2/3 = g^2/n_{\text{in}}$ gives $b = g\sqrt{3/n_{\text{in}}}$.Match the uniform's variance to the normal one's.
6. Slope $\sqrt5$ in step 3's gain: $g^2 = 2/(1 + 5) = 1/3$, so $b = \sqrt{1/3}\,\sqrt{3/n_{\text{in}}} = 1/\sqrt{n_{\text{in}}}$.kaiming_uniform_ computes its bound as gain times $\sqrt{3/\text{fan\_in}}$, with the leaky-ReLU gain of step 3 evaluated at its a; a slope above $1$ is not a sensible activation, only a way to set the scale.
7. $\mathbb{E}[f_\alpha(z)^2] = \tfrac12(1+\alpha^2)\sigma_z^2$, $g = \sqrt{2/(1+\alpha^2)}$, uniform bound $b = g\sqrt{3/n_{\text{in}}}$; with slope $\sqrt5$ the bound is $1/\sqrt{n_{\text{in}}}$ and the variance $1/(3n_{\text{in}})$That variance is a sixth of He's $2/n_{\text{in}}$, so a plain stack of default nn.Linear layers with ReLUs shrinks the variance by $1/6$ per layer (Problem 5 with $c = 1/3$). The same step 5 gives He uniform $\sqrt{6/n_{\text{in}}}$ and Xavier uniform $\sqrt{6/(n_{\text{in}} + n_{\text{out}})}$.

Problem 10

1. $q_l = \mathbb{E}[(h_i + \beta F_i)^2] = q_{l-1} + 2\beta\,\mathbb{E}[h_iF_i] + \beta^2\,\mathbb{E}[F_i^2]$, with $h = h^{(l-1)}$ and $F = F_l(h)$.Expand the square and use linearity of expectation.
2. $q_l = (1 + \beta^2)\,q_{l-1}$.The cross term is $0$ and $\mathbb{E}[F_i^2] = q_{l-1}$, by the hypotheses on the branch.
3. $q_L = (1 + \beta^2)^L\,q_0$.Apply step 2 for each of the $L$ blocks.
4. $\beta = 1$, $L = 10$: $2^{10} = 1024$.The skip and the branch each carry the full second moment, and they add.
5. $\beta = 1/\sqrt L$: $(1 + 1/L)^L$, which for $L = 10$ is $1.1^{10} \approx 2.594$.Step 3 with $\beta^2 = 1/L$.
6. $L\log(1 + 1/L) < L\cdot\tfrac1L = 1$, so $(1 + 1/L)^L < e$.$\log(1 + t) < t$ for $t > 0$, because $\log$ is concave and $t$ is its tangent line at $t = 0$.
7. $q_L/q_0 = (1 + \beta^2)^L$: $1024$ for $\beta = 1$, $L = 10$, and $(1 + 1/L)^L \approx 2.594 < e$ for $\beta = 1/\sqrt L$, $L = 10$Uncorrelated signals add their second moments, so a stack of variance-preserving branches still grows exponentially with depth. Scaling the branches by $1/\sqrt L$ caps the growth at $e$ however deep the stack; some ResNet recipes instead start each branch's last layer at zero, which is $\beta = 0$ at initialisation.