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.