batch norm反向公式推导

输入

$X_i=(x_{i0},x_{i1},...,x_{i(n-1)})$ $i\in [0,m-1]$ batch-size等于m,特征维度n

输出

$Y_i=(y_{i0},y_{i1},...,y_{i(n-1)})$ $i\in [0,m-1]$ 维度和输入$X$一致

前向计算

1. 均值
   $\mu ={\mu }_0,{\mu }_1,...,{\mu }_n$ 其中
   ${\mu }_p=\frac{1}{m}{\sum }_i{x}_{ip}$

2. 方差
   $\sigma ={\sigma }_0,{\sigma }_1,...,{\sigma }_n$ 其中
   ${\sigma }_p=\frac{1}{m}{\sum }_i(x_{ip}-{\mu }_p{)}^2$

3. 中间结果
   ${\overline{x}}_{ip}=\frac{x_{ip}-{\mu }_p}{\sqrt{{\sigma }_p^2+ϵ}}$

4. 结果
   $y_{ip}={\gamma }_p{\overline{x}}_{ip}+{\beta }_p$ 其中
   参数$\gamma ={\gamma }_0,{\gamma }_1,...,{\gamma }_{n-1}$ 和
   $\beta ={\beta }_0,{\beta }_1,...,{\beta }_{n-1}$
   是learnable parameters

反向计算

$$\frac{\partial O}{\partial x_{ij}}=\sum _{kl}\frac{\partial O}{\partial y_{kl}}\frac{\partial y_{kl}}{\partial x_{ij}}=\sum _{kl}\frac{\partial O}{\partial y_{kl}}\frac{\partial y_{kl}}{\partial {\overline{x}}_{ij}}\frac{\partial {\overline{x}}_{ij}}{\partial x_{ij}}=\sum _{kl}\frac{\partial O}{\partial y_{kl}}{\gamma }_l\frac{\partial {\overline{x}}_{ij}}{\partial x_{ij}}(1)$$

$$\frac{\partial {\overline{x}}_{ij}}{\partial x_{ij}}=\frac{\frac{\partial (x_{kl}-{\mu }_l)}{\partial x_{ij}}\sqrt{{\sigma }_l^2+ϵ}-\frac{\partial \sqrt{{\sigma }_l^2+ϵ}}{\partial x_{ij}}(x_{kl}-{\mu }_l)}{{\sigma }_l^2+ϵ}(2)$$

$$\frac{\partial (x_{kl}-{\mu }_l)}{\partial x_{ij}}={\delta }_{ki}{\delta }_{lj}-{\delta }_{lj}\frac{1}{m}(3)$$
其中
$${\delta }_{pq}=\{\begin{array}{l}1p=q\\ 0else\end{array}$$
这个符号可以替代推导过程中的if-else，遇到求和号可以消除
$$\frac{\partial \sqrt{{\sigma }_l^2+ϵ}}{\partial x_{ij}}=\frac{1}{m}\frac{1}{\sqrt{{\sigma }_l^2+ϵ}}{\delta }_{lj}(x_{il}-{\mu }_l)(4)$$
(3)(4)带入(2)得到
$$\frac{\partial {\overline{x}}_{ij}}{\partial x_{ij}}={\delta }_{lj}\frac{({\delta }_{ki}-\frac{1}{m})\sqrt{{\sigma }_l^2+ϵ}-\frac{1}{m\sqrt{{\sigma }_l^2+ϵ}}(x_{kl}-{\mu }_l)(x_{il}-{\mu }_l)}{{\sigma }_l^2+ϵ}$$
上式带入公式(1)得到
$$\frac{\partial O}{\partial x_{ij}}=\frac{{\gamma }_j}{m\sqrt{{\sigma }_j^2+ϵ}({\sigma }_j^2+ϵ)}(({\sigma }_j^2+ϵ)(m\frac{\partial O}{\partial y_{jj}}-\sum _k\frac{\partial O}{\partial y_{kj}})-(x_{ij}-{\mu }_j)(x_{kj}-{\mu }_j)\sum _k\frac{\partial O}{\partial y_{kj}})(done)$$