6
$\begingroup$

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:

enter image description here

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.

$\endgroup$
3
$\begingroup$

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.

|improve this answer|||||
$\endgroup$
0
$\begingroup$

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...

Check This Image Out

|improve this answer|||||
$\endgroup$
  • $\begingroup$ 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. $\endgroup$ – SaldaVonSchwartz Apr 9 '18 at 6:52

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service, privacy policy and cookie policy

Not the answer you're looking for? Browse other questions tagged or ask your own question.