5. 前向/反向传播——Layer Normalization

参考资料

cs231n Course Materials: Backprop
Derivatives, Backpropagation, and Vectorization
cs231n Lecture 4:Neural Networks and Backpropagation
cs231n Assignment 2
笔记: Batch Normalization及其反向传播

5. Layer Normalization

其实Layer Normalization基本就是把Batch Normalization对输入X的第一个维度N做的事情作用在了第二个维度D上(不过可学习参数$\gamma$，$\beta$仍然是D维)，所以实现的话其实只需要把输入转置一下，然后再适当的位置转置回去就差不多了(当然，因为是对第二维D求均值和方差，测试阶段也可以单独算，就不需要再使用滑动平均了)。

不过，这里还是推导一下。

前向传播

"""Forward pass for layer normalization.

    Input:
    - X: Data of shape (N, D)
    - gamma: Scale parameter of shape (D,)
    - beta: Shift paremeter of shape (D,)
    - ln_param: Dictionary with the following keys:
        - eps: Constant for numeric stability

    Returns a tuple of:
    - Y: of shape (N, D)
    - cache: A tuple of values needed in the backward pass
"""

$${\mu }_i=\frac{1}{D}\sum _{j=1}^D{X}_{i,j}\text{\hspace{1em}(5.1)}$$
$${\sigma }_i=\sqrt{\frac{1}{D}\sum _{j=1}^D{\left(X_{i,j}-{\mu }_i\right)}^2+ϵ}\text{\hspace{1em}(5.2)}$$
$${\hat{X}}_{i,j}=\frac{X_{i,j}-{\mu }_i}{{\sigma }_i}\text{\hspace{1em}(5.3)}$$
$$Y_{i,j}={\gamma }_j{\hat{X}}_{i,j}+{\beta }_j\text{\hspace{1em}(5.4)}$$

反向传播

由式(5.4)得
$$\frac{\partial L}{\partial {\beta }_j}=\sum _i\frac{\partial L}{\partial Y_{i,j}}\frac{\partial Y_{i,j}}{\partial {\beta }_j}=\sum _i\frac{\partial L}{\partial Y_{i,j}}\text{\hspace{1em}(5.5)}$$
$$\frac{\partial L}{\partial {\gamma }_j}=\sum _i\frac{\partial L}{\partial Y_{i,j}}\frac{\partial Y_{i,j}}{\partial {\gamma }_j}=\sum _i\frac{\partial L}{\partial Y_{i,j}}{\hat{X}}_{i,j}\text{\hspace{1em}(5.6)}$$
$$\frac{\partial L}{\partial {\hat{X}}_{i,j}}=\frac{\partial L}{\partial Y_{i,j}}\frac{\partial Y_{i,j}}{\partial {\hat{X}}_{i,j}}=\frac{\partial L}{\partial Y_{i,j}}{\gamma }_j\text{\hspace{1em}(5.7)}$$

---

由式(5.3)可得
$$\begin{array}{rl}\frac{\partial L}{\partial {\hat{\sigma }}_i}& =\sum _{j=1}^D\frac{\partial L}{\partial {\hat{X}}_{i,j}}\frac{\partial {\hat{X}}_{i,j}}{\partial {\hat{\sigma }}_i}\\ & =-\sum _{j=1}^D\frac{\partial L}{\partial {\hat{X}}_{i,j}}\frac{X_{i,j}-{\mu }_i}{{\left(\widehat{{\sigma }_i}\right)}^2}\\ & =-\frac{1}{{\sigma }_i}\sum _{k=1}^D\frac{\partial L}{\partial {\hat{X}}_{i,k}}{\hat{X}}_{i,k}\end{array}\text{\hspace{1em}(5.8)}$$
$$\begin{array}{rl}\frac{\partial L}{\partial {\mu }_i}& =\frac{\partial L}{\partial {\hat{\sigma }}_i}\frac{\partial {\hat{\sigma }}_i}{\partial {\mu }_i}+\sum _{j=1}^D\frac{\partial L}{\partial {\hat{X}}_{i,j}}\frac{\partial {\hat{X}}_{i,j}}{\partial {\mu }_i}\\ & =\frac{\partial L}{\partial {\hat{\sigma }}_i}\frac{\frac{1}{D}{\sum }_{j=1}^D-2\left(X_{i,j}-{\mu }_i\right)}{2\sqrt{\frac{1}{D}{\sum }_{j=1}^D{\left(X_{i,j}-{\mu }_i\right)}^2}}+\sum _{j=1}^D\frac{\partial L}{\partial {\hat{X}}_{i,j}}\frac{\partial {\hat{X}}_{i,j}}{\partial {\mu }_i}\\ & =\frac{\partial L}{\partial {\hat{\sigma }}_i}\frac{\frac{1}{D}{\sum }_{j=1}^D\left({\mu }_i-X_{i,j}\right)}{\sqrt{\frac{1}{D}{\sum }_{j=1}^D{\left(X_{i,j}-{\mu }_i\right)}^2}}+\sum _{j=1}^D\frac{\partial L}{\partial {\hat{X}}_{i,j}}\frac{\partial {\hat{X}}_{i,j}}{\partial {\mu }_i}\\ & =\frac{\partial L}{\partial {\hat{\sigma }}_i}\frac{0}{\sqrt{\frac{1}{D}{\sum }_{j=1}^D{\left(X_{i,j}-{\mu }_i\right)}^2}}+\sum _{j=1}^D\frac{\partial L}{\partial {\hat{X}}_{i,j}}\frac{\partial {\hat{X}}_{i,j}}{\partial {\mu }_i}\\ & =\sum _{j=1}^D\frac{\partial L}{\partial {\hat{X}}_{i,j}}\frac{\partial {\hat{X}}_{i,j}}{\partial {\mu }_i}\\ & =\sum _{j=1}^D\frac{\partial L}{\partial {\hat{X}}_{i,j}}\left(-\frac{1}{{\hat{\sigma }}_i}\right)\\ & =\sum _{k=1}^D\frac{\partial L}{\partial {\hat{X}}_{i,k}}\left(-\frac{1}{{\hat{\sigma }}_i}\right)\end{array}\text{\hspace{1em}(5.9)}$$

---

$$\begin{array}{rl}\frac{\partial L}{\partial X_{i,j}}& =\frac{\partial L}{\partial {\hat{X}}_{i,j}}\frac{\partial {\hat{X}}_{i,j}}{\partial X_{i,j}}+\frac{\partial L}{\partial \widehat{{\sigma }_i}}\frac{\partial \widehat{{\sigma }_i}}{\partial X_{i,j}}+\frac{\partial L}{\partial {\mu }_i}\frac{\partial {\mu }_i}{\partial X_{i,j}}\\ & =\frac{\partial L}{\partial {\hat{X}}_{i,j}}\frac{1}{{\hat{\sigma }}_i}+\frac{\partial L}{\partial \widehat{{\sigma }_i}}\frac{\frac{2}{D}\left(X_{i,j}-{\mu }_i\right)}{2\sqrt{\frac{1}{D}{\sum }_{j=1}^D{\left(X_{i,j}-{\mu }_i\right)}^2}}+\frac{\partial L}{\partial {\mu }_i}\frac{1}{D}\\ & =\frac{\partial L}{\partial {\hat{X}}_{i,j}}\frac{1}{{\hat{\sigma }}_i}+\frac{\partial L}{\partial \widehat{{\sigma }_i}}\frac{\frac{2}{D}\left(X_{i,j}-{\mu }_i\right)}{2{\hat{\sigma }}_i}+\frac{\partial L}{\partial {\mu }_i}\frac{1}{D}\\ & =\frac{\partial L}{\partial {\hat{X}}_{i,j}}\frac{1}{{\hat{\sigma }}_i}+\frac{\partial L}{\partial \widehat{{\sigma }_j}}\frac{1}{D}{\hat{X}}_{i,j}+\frac{\partial L}{\partial {\mu }_i}\frac{1}{D}\\ & =\frac{\partial L}{\partial {\hat{X}}_{i,j}}\frac{1}{{\hat{\sigma }}_i}-\left(\frac{1}{{\hat{\sigma }}_i}\sum _{k=1}^D\frac{\partial L}{\partial {\hat{X}}_{i,k}}{\hat{X}}_{i,k}\right)\frac{1}{D}{\hat{X}}_{i,j}-\sum _{k=1}^D\frac{\partial L}{\partial {\hat{X}}_{i,k}}\left(\frac{1}{{\hat{\sigma }}_i}\right)\frac{1}{D}\\ & =\frac{1}{D}\frac{1}{{\hat{\sigma }}_i}\left(D\frac{\partial L}{\partial {\hat{X}}_{i,j}}-{\hat{X}}_{i,j}\sum _{k=1}^D\frac{\partial L}{\partial {\hat{X}}_{i,k}}{\hat{X}}_{i,k}-\sum _{k=1}^D\frac{\partial L}{\partial {\hat{X}}_{i,k}}\right)\end{array}\text{\hspace{1em}(5.10)}$$