# how to derive 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.

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

• 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. – SaldaVonSchwartz Apr 9 '18 at 7:07
• 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$. – shimao Apr 9 '18 at 7:10
• "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. – SaldaVonSchwartz Apr 9 '18 at 7:12
• 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}$. – shimao Apr 9 '18 at 7:17
• in your answer then, 1(i) is a 'one hot' vector but \arrow{1} is a vector of all ones? – SaldaVonSchwartz Apr 10 '18 at 23:04