deriving the gradient of batch normalization

I'm trying to figure out the gradient of batch norm wrt x for backprop, but I get stuck in what I will call 'the triangle of (gradient) death'.

I present to you the triangle of death (in red), in the context of a computation graph for the batch norm equation:

The problem is that applying chain rule to the triangle results in 1 - 1 = 0, which kills the gradient.

My guess is I am messing up one or more of:

• the derivatives.
• the relationships I think the graph implies with respect to the nodes.
• the graph representation I chose for BN as a composition of functions.

but I'm not sure where my error(s) is/are.

Your derivative computation is correct, so I think your understanding of what BN does is slightly off. If you want a more thorough proof that your computation graph is correct, you can backpropagate from $\bar{x} = x-\mu$ using the partial derivatives with respect to each input in the batch, i.e. $\frac{\partial\bar{x}_i}{\partial x_j}$ from $\bar{x}_i = x_i - \mu$. There are two cases to calculate: when $i=j$:

$$\frac{\partial\bar{x}_i}{\partial x_i} = \frac{\partial x_i}{\partial x_i} - \frac{\partial\mu}{\partial x_i} = 1 - \frac{1}{B}$$

And when $i \neq j$:

$$\frac{\partial\bar{x}_i}{\partial x_j} = \frac{\partial x_i}{\partial x_j} - \frac{\partial\mu}{\partial x_j} = -\frac{1}{B}$$

of which there are $B-1$ cases. Therefore:

\begin{align} \frac{\partial\bar{x}}{\partial x} &= \sum_i^B \sum_j^B \frac{\partial\bar{x}_i}{\partial x_j} \\ &= \sum_i^B 1 - \frac{1}{B} - \frac{B-1}{B} \\ &= \sum_i^B 0 \\ \frac{\partial\bar{x}}{\partial x} &= 0 \end{align}

Note also that $\frac{\partial\mu}{x}=1$, as you also correctly pointed out. What this means is that $\frac{\partial\text{BN}}{\partial x}$ is entirely dependent on the demeaned inputs ($\bar{x}=x-\mu$) and not the original inputs themselves; more importantly, backpropagation through a BN layer is invariant to the scale of the inputs. This directly results in some of the observed benefits of using BN, such as a smaller dependence on the initialization scheme and easier/faster training.

Some image credits apply......

It's from a blog which applies the idea of Computational Graphs as explained in CS-231N by Karapathay..

What differs from your graph is that we need to take into account $*,+,/,-$ symbols also..

Will have a careful look at yours also and update here, but till then, this graph works as I have also derived the expressions from this one...

• the graph you posted also has the "triangle of death" which makes me think I am messing up when computing the derivatives. But not sure what you mean about accounting for the operators (+,-,,etc), I think I am accounting for them. It's just that the graph you posted calls a node "" and I call it a variable name, but it should still be the same. – SaldaVonSchwartz Apr 9 '18 at 6:52