参考资料
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
"""
μ
i
=
1
D
∑
j
=
1
D
X
i
,
j
(5.1)
\mu_{i}=\frac{1}{D}\sum_{j=1}^{D}{X_{i,j}}\tag{5.1}
μi=D1j=1∑DXi,j(5.1)
σ
i
=
1
D
∑
j
=
1
D
(
X
i
,
j
−
μ
i
)
2
+
ϵ
(5.2)
\sigma_i=\sqrt{\frac{1}{D}\sum_{j=1}^{D}{\left(X_{i,j}-\mu_i\right)^2}+\epsilon}\tag{5.2}
σi=D1j=1∑D(Xi,j−μi)2+ϵ(5.2)
X
^
i
,
j
=
X
i
,
j
−
μ
i
σ
i
(5.3)
\hat{X}_{i,j}=\frac{X_{i,j}-\mu_i}{\sigma_{i}}\tag{5.3}
X^i,j=σiXi,j−μi(5.3)
Y
i
,
j
=
γ
j
X
^
i
,
j
+
β
j
(5.4)
Y_{i,j}=\gamma_j\hat{X}_{i,j}+\beta_j\tag{5.4}
Yi,j=γjX^i,j+βj(5.4)
反向传播
由式(5.4)得
∂
L
∂
β
j
=
∑
i
∂
L
∂
Y
i
,
j
∂
Y
i
,
j
∂
β
j
=
∑
i
∂
L
∂
Y
i
,
j
(5.5)
\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}}}}\tag{5.5}
∂βj∂L=i∑∂Yi,j∂L∂βj∂Yi,j=i∑∂Yi,j∂L(5.5)
∂
L
∂
γ
j
=
∑
i
∂
L
∂
Y
i
,
j
∂
Y
i
,
j
∂
γ
j
=
∑
i
∂
L
∂
Y
i
,
j
X
^
i
,
j
(5.6)
\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}}\tag{5.6}
∂γj∂L=i∑∂Yi,j∂L∂γj∂Yi,j=i∑∂Yi,j∂LX^i,j(5.6)
∂
L
∂
X
^
i
,
j
=
∂
L
∂
Y
i
,
j
∂
Y
i
,
j
∂
X
^
i
,
j
=
∂
L
∂
Y
i
,
j
γ
j
(5.7)
\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\tag{5.7}
∂X^i,j∂L=∂Yi,j∂L∂X^i,j∂Yi,j=∂Yi,j∂Lγj(5.7)
由式(5.3)可得
∂
L
∂
σ
^
i
=
∑
j
=
1
D
∂
L
∂
X
^
i
,
j
∂
X
^
i
,
j
∂
σ
^
i
=
−
∑
j
=
1
D
∂
L
∂
X
^
i
,
j
X
i
,
j
−
μ
i
(
σ
i
^
)
2
=
−
1
σ
i
∑
k
=
1
D
∂
L
∂
X
^
i
,
k
X
^
i
,
k
(5.8)
\begin{aligned}\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(\hat{\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{aligned}\tag{5.8}
∂σ^i∂L=j=1∑D∂X^i,j∂L∂σ^i∂X^i,j=−j=1∑D∂X^i,j∂L(σi^)2Xi,j−μi=−σi1k=1∑D∂X^i,k∂LX^i,k(5.8)
∂
L
∂
μ
i
=
∂
L
∂
σ
^
i
∂
σ
^
i
∂
μ
i
+
∑
j
=
1
D
∂
L
∂
X
^
i
,
j
∂
X
^
i
,
j
∂
μ
i
=
∂
L
∂
σ
^
i
1
D
∑
j
=
1
D
−
2
(
X
i
,
j
−
μ
i
)
2
1
D
∑
j
=
1
D
(
X
i
,
j
−
μ
i
)
2
+
∑
j
=
1
D
∂
L
∂
X
^
i
,
j
∂
X
^
i
,
j
∂
μ
i
=
∂
L
∂
σ
^
i
1
D
∑
j
=
1
D
(
μ
i
−
X
i
,
j
)
1
D
∑
j
=
1
D
(
X
i
,
j
−
μ
i
)
2
+
∑
j
=
1
D
∂
L
∂
X
^
i
,
j
∂
X
^
i
,
j
∂
μ
i
=
∂
L
∂
σ
^
i
0
1
D
∑
j
=
1
D
(
X
i
,
j
−
μ
i
)
2
+
∑
j
=
1
D
∂
L
∂
X
^
i
,
j
∂
X
^
i
,
j
∂
μ
i
=
∑
j
=
1
D
∂
L
∂
X
^
i
,
j
∂
X
^
i
,
j
∂
μ
i
=
∑
j
=
1
D
∂
L
∂
X
^
i
,
j
(
−
1
σ
^
i
)
=
∑
k
=
1
D
∂
L
∂
X
^
i
,
k
(
−
1
σ
^
i
)
(5.9)
\begin{aligned}\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{aligned} \tag{5.9}
∂μi∂L=∂σ^i∂L∂μi∂σ^i+j=1∑D∂X^i,j∂L∂μi∂X^i,j=∂σ^i∂L2D1∑j=1D(Xi,j−μi)2D1∑j=1D−2(Xi,j−μi)+j=1∑D∂X^i,j∂L∂μi∂X^i,j=∂σ^i∂LD1∑j=1D(Xi,j−μi)2D1∑j=1D(μi−Xi,j)+j=1∑D∂X^i,j∂L∂μi∂X^i,j=∂σ^i∂LD1∑j=1D(Xi,j−μi)20+j=1∑D∂X^i,j∂L∂μi∂X^i,j=j=1∑D∂X^i,j∂L∂μi∂X^i,j=j=1∑D∂X^i,j∂L(−σ^i1)=k=1∑D∂X^i,k∂L(−σ^i1)(5.9)
∂ L ∂ X i , j = ∂ L ∂ X ^ i , j ∂ X ^ i , j ∂ X i , j + ∂ L ∂ σ i ^ ∂ σ i ^ ∂ X i , j + ∂ L ∂ μ i ∂ μ i ∂ X i , j = ∂ L ∂ X ^ i , j 1 σ ^ i + ∂ L ∂ σ i ^ 2 D ( X i , j − μ i ) 2 1 D ∑ j = 1 D ( X i , j − μ i ) 2 + ∂ L ∂ μ i 1 D = ∂ L ∂ X ^ i , j 1 σ ^ i + ∂ L ∂ σ i ^ 2 D ( X i , j − μ i ) 2 σ ^ i + ∂ L ∂ μ i 1 D = ∂ L ∂ X ^ i , j 1 σ ^ i + ∂ L ∂ σ j ^ 1 D X ^ i , j + ∂ L ∂ μ i 1 D = ∂ L ∂ X ^ i , j 1 σ ^ i − ( 1 σ ^ i ∑ k = 1 D ∂ L ∂ X ^ i , k X ^ i , k ) 1 D X ^ i , j − ∑ k = 1 D ∂ L ∂ X ^ i , k ( 1 σ ^ i ) 1 D = 1 D 1 σ ^ i ( D ∂ L ∂ X ^ i , j − X ^ i , j ∑ k = 1 D ∂ L ∂ X ^ i , k X ^ i , k − ∑ k = 1 D ∂ L ∂ X ^ i , k ) (5.10) \begin{aligned}\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{\hat{\sigma_i}}}\frac{\partial{\hat{\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{\hat{\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{\hat{\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{\hat{\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{aligned}\tag{5.10} ∂Xi,j∂L=∂X^i,j∂L∂Xi,j∂X^i,j+∂σi^∂L∂Xi,j∂σi^+∂μi∂L∂Xi,j∂μi=∂X^i,j∂Lσ^i1+∂σi^∂L2D1∑j=1D(Xi,j−μi)2D2(Xi,j−μi)+∂μi∂LD1=∂X^i,j∂Lσ^i1+∂σi^∂L2σ^iD2(Xi,j−μi)+∂μi∂LD1=∂X^i,j∂Lσ^i1+∂σj^∂LD1X^i,j+∂μi∂LD1=∂X^i,j∂Lσ^i1−(σ^i1k=1∑D∂X^i,k∂LX^i,k)D1X^i,j−k=1∑D∂X^i,k∂L(σ^i1)D1=D1σ^i1(D∂X^i,j∂L−X^i,jk=1∑D∂X^i,k∂LX^i,k−k=1∑D∂X^i,k∂L)(5.10)