bn

Backward pass for Batch Normalization

The computational graph for batch norm can be simplified to:

bitmap

$\mathbf{x}$ $\mu$ $n$ $\sigma$ $n$ times to compute normalized x-values. We need to add up all these contributions backwards.

So we need to compute:

\begin{matrix} (1) & \frac{\partial \frac{x - μ}{σ}}{\partial x} = \underset{standard deviation}{\underset{⏟}{\frac{\partial \frac{x - μ}{σ}}{\partial σ} \frac{\partial σ}{\partial x}}} + \underset{inputs}{\underset{⏟}{\frac{\partial \frac{x - μ}{σ}}{\partial x}}} + \underset{mean}{\underset{⏟}{(\frac{\partial \frac{x - μ}{σ}}{\partial μ} + \frac{\partial σ}{\partial μ}) \frac{\partial μ}{\partial x}}} \end{matrix}

$\partial L / \partial\mathbf{x}$ ultimately so we need to use the upstream gradient as per the Chain Rule.

\frac{\partial L}{\partial x} = \frac{\partial L}{\partial \frac{x - μ}{σ}} \frac{\partial \frac{x - μ}{σ}}{\partial x}

$\ref{norm_grad}$ .

\begin{aligned} (2) & \begin{aligned} \frac{\partial \frac{x - μ}{σ}}{\partial σ} & = - \sum_{i}^{n} \frac{x_{i} - μ}{σ^{2}} \end{aligned} \\ (3) & \begin{aligned} \frac{\partial σ}{\partial x} & = \frac{1}{2 n σ} \frac{\sum_{i}^{n} {(x_{i} - μ)}^{2}}{\partial x} \\ = \frac{1}{2 n σ} \frac{\sum_{i}^{n} x_{i}^{2} + \sum_{i}^{n} μ^{2} - 2 \sum_{i}^{n} x_{i} μ}{\partial x} \\ = \frac{1}{2 n σ} \cdot 2 (x - μ) \\ = \frac{x - μ}{n σ} \end{aligned} \\ (4) & \begin{array}{r} \frac{\partial \frac{x - μ}{σ}}{\partial x} = {[\begin{array}{c} \frac{1}{σ} & \dots & \frac{1}{σ} \end{array}]}^{⊺} \end{array} \end{aligned}

$\partial \sigma / \partial \mu$ $0$ (saving us from applying the chain rule many times!), so we only need 2 more derivatives

\begin{aligned} (5) & \begin{array}{r} \frac{\partial \frac{x - μ}{σ}}{\partial μ} = - \sum_{i}^{n} \frac{1}{σ} \end{array} \\ (6) & \begin{array}{r} \frac{\partial μ}{\partial x} = {[\begin{array}{c} \frac{1}{n} & \dots & \frac{1}{n} \end{array}]}^{⊺} \end{array} \end{aligned}

$(2)\ldots(6)$ $\ref{norm_grad}$ , multiply by the upstream gradient and factorize/simplify.

Fool's way (T.T)

Say we expand the normalized vector like so:

\hat{x} = [\begin{matrix} \frac{x_{0} - \sum_{m}^{n} \frac{x_{m}}{n}}{\sqrt{\sum_{k}^{n} \frac{{(x_{k} - \sum_{m}^{n} \frac{x_{m}}{n})}^{2}}{n}}} \\ ⋮ \\ \frac{x_{n} - \sum_{m}^{n} \frac{x_{m}}{n}}{\sqrt{\sum_{k}^{n} \frac{{(x_{k} - \sum_{m}^{n} \frac{x_{m}}{n})}^{2}}{n}}} \end{matrix}]

Then we can compute the gradient as

\begin{matrix} (7) & \frac{\partial L}{\partial x_{i}} = \sum_{j}^{n} \frac{\partial {\hat{x}}_{j}}{\partial x_{i}} \frac{\partial L}{\partial {\hat{x}}_{j}} \end{matrix}

We're going to use the quotient rule to compute the derivative. The derivative for the standard deviation is particularly bashy, so I'll include it here. The rest should be easy.

\begin{aligned} \frac{\partial σ}{\partial x_{i}} & = \frac{1}{2 σ n} \frac{\sum_{k}^{n} {(x_{k} - \sum_{m}^{n} \frac{x_{m}}{n})}^{2}}{\partial x_{i}} \\ = \frac{1}{2 σ n} (\frac{\sum_{k}^{n} {x_{k}}^{2}}{\partial x_{i}} + \frac{\sum_{k}^{n} {(\sum_{m}^{n} \frac{x_{m}}{n})}^{2}}{\partial x_{i}} - \frac{2 \sum_{k}^{n} x_{k} \sum_{m}^{n} \frac{x_{m}}{n}}{\partial x_{i}}) \\ = \frac{1}{2 σ n} (2 x_{i} + \frac{1}{n} \frac{{(\sum_{m}^{n} x_{m})}^{2}}{\partial x_{i}} - \frac{2}{n} \frac{\sum_{k}^{n} x_{k} \sum_{m}^{n} x_{m}}{\partial x_{i}}) \\ = \frac{1}{2 σ n} (2 x_{i} + \frac{1}{n} \frac{{(\sum_{m}^{n} x_{m})}^{2}}{\partial x_{i}} - \frac{2}{n} (\frac{x_{i} (x_{0} + \dots + x_{i} + \dots + x_{m})}{\partial x_{i}} + \frac{\sum_{m}^{n} x_{m} \sum_{k \neq i}^{n} x_{k}}{\partial x_{i}})) \\ = \frac{1}{2 σ n} (2 x_{i} + \frac{1}{n} \frac{{(\sum_{m}^{n} x_{m})}^{2}}{\partial x_{i}} - \frac{2}{n} (\frac{{x_{i}}^{2} + \sum_{k \neq i}^{n} x_{i} x_{k}}{\partial x_{i}} + \sum_{k \neq i}^{n} x_{k})) \\ = \frac{1}{2 σ n} (2 x_{i} + \frac{1}{n} \frac{{(\sum_{m}^{n} x_{m})}^{2}}{\partial x_{i}} - \frac{2}{n} (2 x_{i} + \sum_{k \neq i}^{n} x_{k} + \sum_{k \neq i}^{n} x_{k})) \\ = \frac{1}{2 σ n} (2 x_{i} + \frac{1}{n} \frac{{(\sum_{m}^{n} x_{m})}^{2}}{\partial x_{i}} - \frac{4 \sum_{k}^{n} x_{k}}{n}) \\ = \frac{1}{2 σ n} (2 x_{i} + \frac{2 \sum_{k}^{n} x_{k}}{n} - \frac{4 \sum_{k}^{n} x_{k}}{n}) \\ = \frac{x_{i} - \frac{\sum_{k}^{n} x_{k}}{n}}{σ n} \\ = \frac{x_{i} - μ}{σ n} \end{aligned}

Completing the quotient rule, we end up with 2 cases

\begin{array}{r} \frac{\partial {\hat{x}}_{j}}{\partial x_{i}} = {\begin{cases} \frac{σ (\frac{- 1}{n})}{σ^{2}} - \frac{(x_{j} - μ) (\frac{x_{i} - μ}{σ n})}{σ^{2}} & if i \neq j, \\ \frac{σ (1 - \frac{1}{n})}{σ^{2}} - \frac{(x_{i} - μ) (\frac{x_{i} - μ}{σ n})}{σ^{2}} & if i = j \end{cases} \end{array}

$\ref{eq:7}$ .

\begin{aligned} \begin{aligned} LHS & = (\sum_{j \neq i}^{n} - \frac{1}{n σ} \frac{\partial L}{\partial {\hat{x}}_{j}}) + (\frac{n - 1}{n σ} \frac{\partial L}{\partial {\hat{x}}_{i}}) \\ = (- \frac{1}{n σ} \sum_{j}^{n} \frac{\partial L}{\partial {\hat{x}}_{j}}) + (\frac{1}{n σ} \frac{\partial L}{\partial {\hat{x}}_{i}} + \frac{n - 1}{n σ} \frac{\partial L}{\partial {\hat{x}}_{i}}) \\ = (- \frac{1}{n σ} \sum_{j}^{n} \frac{\partial L}{\partial {\hat{x}}_{j}}) + (\frac{1}{σ} \frac{\partial L}{\partial {\hat{x}}_{i}}) \end{aligned} \\ \begin{aligned} RHS & = - \sum_{j}^{n} \frac{(x_{j} - μ) (x_{i} - μ)}{n σ^{3}} \frac{\partial L}{\partial {\hat{x}}_{j}} \\ = - \frac{{\hat{x}}_{i}}{n σ} \sum_{j}^{n} {\hat{x}}_{j} \frac{\partial L}{\partial {\hat{x}}_{j}} \end{aligned} \end{aligned}

$1/n\sigma$ $\mathbf{x}$ .

Another way

alt

$\mathbf{p}$ $\mathbf{x}$ .