I'm following the derivative calculation of Batch Norm paper:
Something doesn't seem right. In the 3rd equation shouldn't we lose the 2nd term as the sum is equal to 0 ($\mu_B$ is the mean of the $x_i$ over the batch)?
And worse, in the 4th equation, for the same consideration, once we will sum the loss over the entire batch, we will lose the middle term. And then the 1st and 3rd term will cancel each other to give out 0.
This obviously can't be right, but I don't seem to spot my mistake(s).
Ok, so after checking also with the original authors - in the 3rd equation ($\partial \mathcal l/\partial \mu_\mathcal B$) the 2nd term does indeed simplify to 0.
The 4th equation ($\partial \mathcal l/\partial x_i$) however is not summed by itself but is multiplied as part of an outer product with the inputs to the layer - so when computing the gradient of the entire batch, we will have a vector of size $(N,)$ per neuron, or $(N,k)$ matrix per layer if we have $k$ neurons. To compute the downstream $\partial\mathcal l/\partial W $ we will do an outer product $a\cdot \partial\mathcal l/\partial x$ (where $a$ are the activations from the last layer / input to the current layer, and $\cdot$ is an outer product) [this is just one (simple) way to present it, we could also use tensors instead]. It's true that if we sum the matrix across the rows / batch it will sum to 0, but this is not what we do for the gradients.
Also note that in the 4th equation, $\partial \mathcal l/\partial \sigma^2_\mathcal B, \partial \mathcal l/\partial \mu_\mathcal B$ are of size $(,k)$ and need to be broadcast to the size of the matrix $(N,k)$ when computing the gradient for the entire batch.
Update: note that the 4th equation can also be simplified more if we plug in the $\partial \mathcal l/\partial \sigma^2_\mathcal B, \partial \mathcal l/\partial \mu_\mathcal B$ derivatives, and also replace some terms with $\hat x$:
$$ \frac{\partial \mathcal L}{\partial x} = \frac{1}{n \sqrt {\sigma^2+\epsilon}}[n\frac{\partial \mathcal L}{\partial \hat x} - 1^T\frac{\partial \mathcal L}{\partial \hat x}-\hat x(1^T\frac{\partial \mathcal L}{\partial \hat x}\hat x)]\\ $$
You can check the full derivation on this YouTube video I made here.
