Derive: Training Deep Networks Without Breaking Them

Unrolling gives $$v_t=g_t+\beta g_{t-1}+{\beta }^2{g}_{t-2}+\cdots +{\beta }^{t-1}{g}_1.$$ With $g_1=1+S$ and $g_k=1$ for $k\ge 2$, split the sum: $$v_t=\underset{\text{baseline}}{\underbrace{\sum _{k=0}^{t-1}{\beta }^k}}+\underset{\text{spike}}{\underbrace{S{\beta }^{t-1}}}.$$ The baseline is a finite geometric series, $(1-{\beta }^t)/(1-\beta )$. The spike appears once, at step 1, and every later step multiplies its contribution by $\beta$, so its contribution at step $t$ is $S{\beta }^{t-1}$.

For $S=5$ and $\beta =0.9$:

• At step 5, the baseline is $(1-0.9^5)/0.1=(1-0.59049)/0.1=4.0951$, and the spike contributes $5\times 0.9^4=5\times 0.6561=3.2805$. The total is 7.3756, and the excess over the no-spike run is 3.2805.
• At step 10, the baseline is $(1-0.9^{10})/0.1=(1-0.348678)/0.1=6.5132$, and the spike contributes $5\times 0.9^9=5\times 0.387420=1.9371$. The total is 8.4503, and the excess is 1.9371.

The excess falls by a factor of 0.9 per step, so it halves in about $\log (0.5)/\log (0.9)=6.58$ steps. After 20 steps it is $5\times 0.9^{19}=0.675$, and after 40 steps it is $5\times 0.9^{39}=0.082$. Section 20.5 gives momentum its memory; section 20.6 says a loss spike's cause is not visible in the loss alone. Put together, the optimizer carries the spike for many steps after the loss looks healthy again. A curve that recovers therefore does not show that the cause has gone, and the gradient norm, which section 20.6 logs before changing a hyperparameter, is the measurement that shows whether the state is still disturbed.

Write the layer as $$y_i=\sum _{j=1}^{n_{\text{in}}}{w}_{ij}{x}_j.$$ The terms are independent and have mean zero, so their variances add: $$\mathrm{Var}(y_i)=\sum _{j=1}^{n_{\text{in}}}\mathrm{Var}(w_{ij}{x}_j)=n_{\text{in}}{\sigma }^2{v}_x.$$ The forward factor is $a=n_{\text{in}}{\sigma }^2$. For the backward pass, the gradient at input $j$ is $$\frac{\partial L}{\partial x_j}=\sum _{i=1}^{n_{\text{out}}}{w}_{ij}\frac{\partial L}{\partial y_i},$$ so the same variance rule gives $\mathrm{Var}(\partial L/\partial x_j)=n_{\text{out}}{\sigma }^2{v}_g$. The backward factor is $b=n_{\text{out}}{\sigma }^2$.

Setting $a=1$ asks for ${\sigma }^2=1/n_{\text{in}}$. Setting $b=1$ asks for ${\sigma }^2=1/n_{\text{out}}$. Both hold only when $n_{\text{in}}=n_{\text{out}}$. If the fan-in and fan-out differ, any single choice leaves one direction shrinking and the other growing; the ratio between the two factors is $b/a=n_{\text{out}}/n_{\text{in}}$, which the choice of ${\sigma }^2$ cannot change.

Xavier chooses ${\sigma }^2=2/(n_{\text{in}}+n_{\text{out}})$. Then $$a=\frac{2n_{\text{in}}}{n_{\text{in}}+n_{\text{out}}},b=\frac{2n_{\text{out}}}{n_{\text{in}}+n_{\text{out}}}.$$ Their product is $$ab=\frac{4n_{\text{in}}{n}_{\text{out}}}{(n_{\text{in}}+n_{\text{out}}{)}^2}\le 1,$$ because $(n_{\text{in}}-n_{\text{out}}{)}^2\ge 0$ rearranges to $4n_{\text{in}}{n}_{\text{out}}\le (n_{\text{in}}+n_{\text{out}}{)}^2$. Equality holds exactly when $n_{\text{in}}=n_{\text{out}}$. Xavier also makes $a+b=2$, so it balances the two factors by their sum; their product is the quantity that pays for an unequal pair.

For $n_{\text{in}}=100$ and $n_{\text{out}}=900$, $a=200/1000=0.2$ and $b=1800/1000=1.8$. The product is $0.2\times 1.8=0.36$. A stack of $L$ such layers multiplies the forward variance by $0.2^L$ and the backward gradient variance by $1.8^L$, and a full forward-backward pass by $0.36^L$. Even the best single constant cannot make both directions neutral; it can only choose where the imbalance lands. Section 20.1 supplies the product picture, and this derivation gives the exact size of the compromise that section 20.2 describes in words.

(a) The block factor is $1+c$, so its magnitude is $|1+c|$. The inequality $|1+c|<1$ is equivalent to $-1<1+c<1$, which gives $-2<c<0$. At $c=0$ the factor is 1, and at $c=-2$ it is $-1$, whose magnitude is 1. For $c>0$ the factor exceeds 1. For $c<-2$ the factor is below $-1$, so its magnitude exceeds 1. The sign matters separately: at $c=-2$ the product alternates in sign but keeps its size, while immediately inside the interval it shrinks and alternates.

(b) With $c=-1.5$ the factor is $1-1.5=-0.5$. After 12 blocks the product is $$(-0.5{)}^{12}=\frac{1}{2^{12}}=\frac{1}{4096}=0.000244141,$$ which is 0.000244 to six decimal places. The sign is positive because 12 is even. The gradient reaches the input, but it is 4096 times smaller than the gradient at the loss.

(c) The skip guarantees a route, not a size. A branch with slope between $-2$ and $0$ cancels part of the skip and leaves a factor below 1 in magnitude, so a deep stack still vanishes; at $c=-1$ the factor is exactly 0 and the gradient dies at that block. A branch with slope below $-2$ makes the stack explode with alternating sign. The identity case is $c=0$: the branch adds nothing, the block passes its input through, and every factor is 1. Section 20.4's table uses positive branch slopes, where the factor exceeds 1 and the risk is growth rather than decay. This exercise marks the boundary on the other side and shows that the sign of the branch slope decides which failure the skip produces.