Sources:

• Adjrej Karpathy's video Building makemore Part 4: Becoming a Backprop Ninja.
• Paper: Batch Normalization: Accelerating Deep Network Training by Reducing Internal Covariate Shift
• PyTorch: BatchNorm1d

Batch Normalization

Given samples $\{x_i\}$ over a mini-batch: $\mathcal{B}=\left\{x_{1\dots N}\right\}$, the batch normalization is $$y_i=f(x_i,\mathcal{B}{)}_{\gamma ,\beta }=\gamma {\hat{x}}_i+\beta =\gamma \frac{x_i-\mu }{\sqrt{{\sigma }^2+ϵ}}+\beta .$$ where:

• $y_i$ is the output of the batch normalization layer.
• ${\hat{x}}_i$ is the normalized input $x_i$
• $ϵ$ is a small constant added for avoiding division by 0 error.
• $\gamma$ and $\beta$ are parameters learned during training for each feature, representing the scale and shift to be applied after normalization, respectively.

The detailed process is

$$ $$\begin{array}{rl}& {\mu }_{\mathcal{B}}\leftarrow \frac{1}{m}\sum _{i=1}^m{x}_i\text{ // mini-batch mean }\\ & {\sigma }_{\mathcal{B}}^2\leftarrow \frac{1}{m-1}\sum _{i=1}^m{(x_i-{\mu }_{\mathcal{B}})}^2\text{ // mini-batch variance }\\ & {\hat{x}}_i\leftarrow \frac{x_i-{\mu }_{\mathcal{B}}}{\sqrt{{\sigma }_{\mathcal{B}}^2+ϵ}}\text{ // normalize }\\ & y_i\leftarrow \gamma {\hat{x}}_i+\beta \equiv {\mathrm{BN}}_{\gamma ,\beta }\left(x_i\right)\text{ // scale and shift }\end{array}$$

$$

Derivation of batch norm

We have: $$\frac{\partial L}{\partial x_i}=\sum _{j=1}^N\frac{\partial L}{\partial y_j}\frac{\partial y_j}{\partial x_i}.$$ $\frac{\partial L}{\partial y_j}$ is the upstream gradient, so it's already given here.

Since $y_j=\gamma {\hat{x}}_j+\beta$. $y_j$ is a function of ${\hat{x}}_j$. We have $$\frac{\partial y_j}{\partial x_i}=\frac{\partial y_j}{\partial {\hat{x}}_j}\frac{\partial {\hat{x}}_j}{\partial x_i},$$ and $$\frac{\partial y_j}{\partial {\hat{x}}_j}=\gamma .$$ So $$\frac{\partial y_j}{\partial x_i}=\gamma \frac{\partial {\hat{x}}_j}{\partial x_i}.$$ This means we only need to calculate $\frac{\partial {\hat{x}}_j}{\partial x_i}$.

Recall that ${\hat{x}}_i=\frac{x_i-\mu }{\sqrt{{\sigma }^2+ϵ}}$, we obtain: $$\begin{array}{rl}\frac{\partial {\hat{x}}_j}{\partial x_i}& =\frac{\partial }{\partial x_i}\left(\frac{x_j-\mu }{\sqrt{{\sigma }^2+\epsilon }}\right)\\ & =({\delta }_{ij}-\frac{\partial \mu }{\partial x_i}){({\sigma }^2+\epsilon )}^{-\frac{1}{2}}+(x_j-\mu )(-\frac{1}{2}){({\sigma }^2+\epsilon )}^{-\frac{3}{2}}\frac{\partial {\sigma }^2}{\partial x_i}.\end{array}$$ where ${\delta }_{ij}$ is the Kronecker delta function: $${\delta }_{ij}=\{\begin{array}{ll}0& \text{ if }i\ne j\\ 1& \text{ if }i=j.\end{array}$$

Therefore, we compute the derivation $\frac{\partial \mu }{\partial x_i}$, $\frac{\partial {\sigma }^2}{\partial x_i}$.

For the former: $$ $$\begin{array}{rl}\frac{\partial \mu }{\partial x_i}& =\frac{\partial }{\partial x_i}(\frac{1}{N}(x_1+x_2+\cdots +x_i+\cdots +x_N))\\ & =\frac{1}{N}(0+\cdots +1+\cdots +0)=\frac{1}{N}\end{array}$$

$$

For the latter: $$ $$\begin{array}{rl}\frac{\partial {\sigma }^2}{\partial x_i}& =\frac{\partial }{\partial x_i}\left(\frac{1}{N-1}\sum _{k=1}^N{(x_k-\mu )}^2\right)\\ & =\frac{1}{N-1}\cdot \frac{\partial }{\partial x_i}({(x_1-\mu )}^2+\cdots +{(x_i-\mu )}^2+\cdots )\\ & =\frac{1}{N-1}[2(x_1-\mu )\frac{\partial (-\mu )}{\partial x_i}+2(x_2-\mu )\frac{\partial (-\mu )}{\partial x_i}+\cdots +2(x_i-\mu )\frac{\partial (x_i-\mu )}{\partial x_i}+\cdots ]\\ & =\frac{1}{N-1}[2(x_1-\mu )\cdot \frac{-1}{N}+2(x_2-\mu )\cdot \frac{-1}{N}+\cdots +2(x_i-\mu )\cdot (1-\frac{1}{N})+\cdots ]\\ & =\frac{1}{N-1}\cdot 2[\underset{0}{\underbrace{\sum _{k=1}^N(x_k-\mu }})\cdot \left(\frac{-1}{N}\right)+x_i-\mu ]\\ & =\frac{1}{N-1}\cdot 2\cdot (x_i-\mu )\\ & =\frac{2}{N-1}\cdot (x_i-\mu )\end{array}$$

$$

The transition from 3rd line to 4th line is because $\frac{\partial \mu }{\partial x_i}=\frac{1}{N}$.

So we have: $$ $$\begin{array}{rl}\frac{\partial {\hat{x}}_j}{\partial x_i}& =({\delta }_{ij}-\frac{1}{N}){({\sigma }^2+\epsilon )}^{-\frac{1}{2}}+(x_j-\mu )(-\frac{1}{2}){({\sigma }^2+\epsilon )}^{-\frac{3}{2}}\frac{2}{N-1}(x_i-\mu )\\ & =({\delta }_{ij}-\frac{1}{N}){({\sigma }^2+\epsilon )}^{-\frac{1}{2}}-\frac{1}{N-1}{({\sigma }^2+\epsilon )}^{-\frac{1}{2}}\cdot \frac{x_i-\mu }{\sqrt{{\sigma }^2+\epsilon }}\cdot \frac{x_j-\mu }{\sqrt{{\sigma }^2+\epsilon }}\\ & ={({\sigma }^2+\epsilon )}^{-\frac{1}{2}}[{\delta }_{ij}-\frac{1}{N}-\frac{1}{N-1}({\hat{x}}_i\cdot {\hat{x}}_j)]\\ & =\frac{{({\sigma }^2+\epsilon )}^{-\frac{1}{2}}}{N}[N\cdot {\delta }_{ij}-1-\frac{N}{N-1}\left({\hat{x}}_i{\hat{x}}_j\right)].\end{array}$$

$$

As a result: $$\begin{array}{rl}\frac{\partial L}{\partial x_i}& =\sum _{j=1}^N\frac{\partial L}{\partial y_j}\frac{\partial y_j}{\partial x_i}=\sum _{j=1}^N\frac{\partial L}{\partial y_j}\frac{\partial y_j}{\partial {\hat{x}}_j}\frac{\partial {\hat{x}}_j}{\partial x_i}\\ & =\sum _{j=1}^N\frac{\partial L}{\partial y_j}\cdot \gamma \cdot \frac{{({\sigma }^2+\epsilon )}^{-\frac{1}{2}}}{N}[N\cdot {\delta }_{ij}-1-\frac{N}{N-1}({\hat{x}}_i\cdot {\hat{x}}_j)]\\ & =\frac{\gamma {({\sigma }^2+\epsilon )}^{-\frac{1}{2}}}{N}\sum _{j=1}^N\frac{\partial L}{\partial y_j}[N\cdot {\delta }_{ij}-1-\frac{N}{N-1}({\hat{x}}_i\cdot {\hat{x}}_j)]\\ & =\frac{\gamma {({\sigma }^2+\epsilon )}^{-\frac{1}{2}}}{N}[\left(N\sum _{j=1}^N\frac{\partial L}{\partial y_j}{\delta }_{ij}\right)-\left(\sum _{j=1}^N\frac{\partial L}{\partial y_j}\right)-\left(\frac{N}{N-1}{x}_i\sum _{j=1}^N\frac{\partial L}{\partial y_j}{\hat{x}}_j\right)]\\ & =\frac{\gamma {({\sigma }^2+\epsilon )}^{-\frac{1}{2}}}{N}(N\frac{\partial L}{\partial y_i}-\sum _{j=1}^N\frac{\partial L}{\partial y_j}-\frac{N}{N-1}{\hat{x}}_i\sum _{j=1}^N\frac{\partial L}{\partial y_j}{\hat{x}}_j).\end{array}$$