I struggle to understand how batch normalization (BN) enables larger learning rates during gradient descent according to the original paper. I am aware that some of the explanations given in the latter have been debunked, but I would like to understand the logic behind them anyway.
The central claim is that BN has this effect on the learning rate because it prevents exploding gradients. I find the intuition behind this best explained in a video by Ian Goodfellow, where he uses the "simplest possible network" for illustration:
$\hat{y} = abcde$
so, a network that consists of 5 one-unit layers (where $a/b/c/d/e$ are the respective weights of the units), and which does not introduce non-linearity through activation functions. Obviously, during forward propagation, the value of $a$ will determine the statistics of the activation at $d$, as Goodfellow explains. Similarly, during backpropagation, the value of $d$ will influence the gradient of $a$ since the derivative w.r.t. $a$ is
$\frac{\delta \hat{y}}{\delta a} = bcde$
so far so good. Adding normalization steps before/after each layer prevents this interaction between layers and keeps the gradients from exploding (due to the normalized value range). This way, gradient descent can make large modifications to parameters, without having to adjust to the propagated effect of said modifications in later iterations, causing more linear progress and less oscillations. Am I correct so far? Now, I have been trying to apply the same logic to the network shown in the below picture (taken from this article):
Here, the partial derivative of the cost function w.r.t. the weight $w_{1}$, is given by:
$\frac{\delta J}{\delta w_{1}} = \frac{\delta J}{\delta \hat{y}} \frac{\delta \hat{y}}{\delta z_{2}} \frac{\delta z_{2}}{\delta a_{1}} \frac{\delta a_{1}}{\delta z_{1}} \frac{\delta z_{1}}{\delta w_{1}}$
if the $ReLU$ is used as the activation function and considering that $z_{i} = w_{i}a_{i-1} + b_{i}$, this becomes (leaving out $\frac{\delta J}{\delta \hat{y}}$ for simplicity, and assuming that $z_{i} > 0$):
$\begin{align} \frac{\delta J}{\delta w_{1}} &= \frac{\delta J}{\delta \hat{y}} \cdot ReLU'(z_{2}) \cdot w_{2} \cdot ReLU'(z_{1}) \cdot x_{1}\\ &= \frac{\delta J}{\delta \hat{y}} \cdot 1 \cdot w_{2} \cdot 1 \cdot x_{1} \end{align}$
My problem is that in the original paper, BN is applied before the activation, so $BN(w_{i}a_{i-1} + b)$, i.e. $BN(z_{i})$. However, $ReLU'(z_{i})$ is always 1 or 0. And if a different activation is used, such as the sigmoid, then $\sigma'(z_{i})$ is always $\leq 1$. Point being, that I'm struggling to imagine how normalizing $z_{i}$ can make such a big difference, since the value range of $g(z_{i})$ is anyway very restricted for any activation function $g$. In the explanation by Goodfellow, the normalized values go into the multiplication unmodified, so it makes more sense to normalize them.
PS: I have asked a similar question about exploding gradients before ... so I guess the idea just confuses me.
