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.


From the way you have written it, $x_i$ are scalars, and $x$ is a vector representing the input batch. $\mu$ is a scalar.

$\frac{\partial}{\partial x} x_i = \mathbf{1}(i)$, where $\mathbf{1}(i)$ is an indicator vector with the $i$th component set to 1. Therefore, $\frac{\partial \mu}{\partial x}$ is actually $\frac{1}{B} \vec 1$.

Also, it's not correct to write $\frac{\partial \overline x}{\partial x} = 1$, in the bottom of the triangle, since that is what we are trying to calculate!

$\frac{\partial \overline x}{\partial \mu} = -\vec 1^T$, not -1, so $\frac{\partial \overline x}{\partial \mu} \frac{\partial \mu}{\partial x} = -\frac{1}{B} \mathbb{I}$, giving us $\frac{\partial (x-\mu)}{\partial x} = I -\frac{1}{B} \mathbb{I}$, where I use $\mathbb{I}$ to denote the matrix of all 1's.

In other words, $\frac{\partial \overline x_i}{\partial x_i} = \frac{B-1}{B}$, and $\frac{\partial \overline x_i}{\partial x_j} = -\frac{1}{B}$ for $i \neq j$.

  • $\begingroup$ I have a few questions: "$\frac{\partial \overline x}{\partial \mu} = -\vec 1^T$, not 1," I have it as -1 already in the graph, or do you mean something else? Also, in $I -\frac{1}{B}$, is I the identity matrix? So far seems like my mistakes are then because of treating the individual x_i as the whole x vector. $\endgroup$ – SaldaVonSchwartz Apr 9 '18 at 7:07
  • $\begingroup$ I use $I$ for the identity matrix, $1$ for the scalar, $\vec 1$ for a vector of 1's, and $\mathbb{I}$ for a matrix of 1's. Yes, it's important to differentiate $x_i$ and the vector $x$. $\endgroup$ – shimao Apr 9 '18 at 7:10
  • $\begingroup$ "Also, it's not correct to write $\frac{\partial \overline x}{\partial x} = 1$, in the bottom of the triangle, since that is what we are trying to calculate!" I agree, yeah it felt kind of wrong when I wrote it, but what would I write instead? cause as far as the computation graph I DO compute the gradient from x bar PLUS the gradient from mu, which in turn is computed from the gradient of x bar with respect to mu. $\endgroup$ – SaldaVonSchwartz Apr 9 '18 at 7:12
  • $\begingroup$ Since $\overline x = x - \mu$, I would write $\frac{\partial \overline x}{\partial x} = \frac{\partial x}{\partial x} - \frac{\partial \mu}{\partial x} = I - \frac{\partial \mu}{\partial x}$. So in a sense, the gradient flowing across that bottom edge is just $I$. I'm not aware of any notation for what to call that edge, but certainly not $\frac{\partial \overline x}{\partial x}$. $\endgroup$ – shimao Apr 9 '18 at 7:17
  • $\begingroup$ in your answer then, 1(i) is a 'one hot' vector but \arrow{1} is a vector of all ones? $\endgroup$ – SaldaVonSchwartz Apr 10 '18 at 23:04

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.