Why take the gradient of the moments (mean and variance) when using Batch Normalization in a Neural Network?

Question

When doing Batch Normalization (BN) it makes sense to me to treat the BN transform as a layer that we need to do back propagatiod and thus have derivatives to update its parameters (for each layer) the scale $\gamma^{(k)}$ and shift $\beta^{(k)}$.

However, what does not make sense to me is why we would need the derivates of the moments (mean $\mu$ and variance $\sigma$) using gradient descent. This is clearly reflected in the paper (but not explained clearly) when they even have derivatives with respect to these two quantities on page 4:

$$ \frac{\partial l}{\partial \sigma^2_{\mathcal{B}}} = \sum^{m}_{i=1} \frac{\partial \mathcal{l}}{\partial \hat x_i} \cdot (x_i - \mu_{\mathcal{B}} ) \cdot \frac{-1}{2} ( \sigma^2_{\mathcal{B}} + \epsilon )^{-\frac{3}{2}}$$

$$ \frac{\partial l}{\partial \mu_{\mathcal{B}}} = \left(\sum^{m}_{i=1} \frac{\partial \mathcal{l}}{\partial \hat x_i} \cdot \frac{-1}{\sqrt{\sigma_{\mathcal{B}}^2 + \epsilon }} \right) + \frac{\partial l}{\partial \sigma^2_{\mathcal{B}}} \cdot \frac{ \sum^m_{i=1} -2 (x_i - \mu_{\mathcal{B}} }{m} ) $$

The thing that is specifically confusing me is that I thought the means and standard deviations were constant during a specific epoch of training (so their derivatives should be zero. I thought that because the original paper said:

we make the second simplification: since we use mini-batches in stochastic gradient training, each mini-batch produces estimates of the mean and variance of each activation.

Furthermore, in their pseudocode they even compute the (moments) mean and variance according to the current batch:

which further confuses me why there would even be any derivatives with respect to such quantities.

Furthermore, they seem to be the population mean and variance during inference, which makes me further suspect that the moments (mean and variance) should be not be variables. Are they parameters, variables or constants? Someone knows?

---

2: Ioffe S. and Szegedy C. (2015),
"Batch Normalization: Accelerating Deep Network Training by Reducing Internal Covariate Shift",
Proceedings of the 32nd International Conference on Machine Learning, Lille, France, 2015.
Journal of Machine Learning Research: W&CP volume 37

Answer 1

Derivatives of the moments are used for backpropagation.

Along the two derivatives of the moments on page 4, it gives the derivatives with respect to the input, which makes use of the derivative of the moments, $$\frac{\partial l}{\partial x}=\frac{\partial l}{\partial \hat{x}}\cdot\frac{1}{\sqrt{\sigma^2+\epsilon}}+\frac{\partial l}{\partial \sigma^2}\cdot\frac{2(x-\mu)}{m}+\frac{\partial l}{\partial \mu}\cdot\frac{1}{m},$$ which will be used for computing the derivatives of the parameters in previous layers by the chain rule.

IMO, the moments are not treated as parameters nor constants, they can be thought of as some intermediate results of computing the output of the layer.

Answer 2

There's no reason for the moments to be thought of as constants. The quote

we make the second simplification: since we use mini-batches in stochastic gradient training, each mini-batch produces estimates of the mean and variance of each activation.

doesn't imply that they are constant, it just makes them scalars.

Another way to look at it is that there are constant only for a mini-batch (but so are the input and the output!). Furthermore, the output (and therefore the loss) is dependant on the change from mini batch to mini batch, and thus the derivative is logically non-zero.

As an aside, these quantities are only used to calculate the gradient towards inputs, as mentioned by @dontloo .

As for your last question, they're variables, but completely tied to the input variables and the network's parameters. (Not only the input variables due to batch normalization being appliable (and mostly interesting) in-network)

Answer 3

There is a reason, why you should propagate gradient through mean and variance. I will try to show you a simple example. Let's just start with the mean. Imagine, that you have two samples (with one attribute) with values: $x_1$ and $x_2$.

Now mean is $m = \frac{x_1 + x_2}{2}$ and output after subtracting the mean would be: $y_1 = x_1 - m$ and $y_2 = x_2 - m$.

Now consider the backward pass and imagine, that both of those gradients $\frac{\partial L}{\partial y_1}$ and $\frac{\partial L}{\partial y_2}$ are equal to one. If you do not pass gradient through $m$, then also $\frac{\partial L}{\partial x_1}$ and $\frac{\partial L}{\partial x_2}$ would be equal to one. And you will increase both inputs. But this would have zero effect in the next iteration since you are subtracting the mean! It would actually be detrimental since lower layers would get some confusing signal.

Now if you pass gradients through mean, then $\frac{\partial L}{\partial m} = -\left( \frac{\partial L}{\partial y_1} + \frac{\partial L}{\partial y_2} \right)$. Which is two in our example.

And $\frac{\partial L}{\partial x_1} = \frac{\partial L}{\partial y_1} - \frac{1}{2} \frac{\partial L}{\partial m}$. Which is zero in our example.

Thus we can say, that passing gradient through mean prevents changing mean of the inputs, which is meaningless.

And we can say same thing about passing gradients through variance, but demonstration would be much more complicated.