前向传播

假设待归一化的$m$维向量为$x$，均值和标准差分别是$\mu{(x)}$和$\sigma{(x)}$，LayerNorm的参数是$w$和$b$，那么层归一化后的输出为：
$$
y = w \odot \frac{x - \mu}{\sqrt{\sigma^2 + \epsilon}} + b
$$

这里的极小量$\epsilon$是为了防止标准差为0导致溢出。为了简单起见，我们不加上这一项，原公式也就变成了：
$$
y = w \odot \frac{x - \mu}{\sigma} + b
$$

反向传播

假设损失函数$\mathcal{L}$对输出$y$的梯度是$\frac{\partial{\mathcal{L}}}{\partial{y}}$，那么这里需要求三个梯度：$\frac{\partial{\mathcal{L}}}{\partial{w}}$、$\frac{\partial{\mathcal{L}}}{\partial{b}}$和$\frac{\partial{\mathcal{L}}}{\partial{x}}$。

令$\hat x = \frac{x - \mu}{\sigma}$，那么原公式就变成了：
$$
y = w \odot \hat x + b
$$

两个参数的梯度非常简单：
$$
\begin{aligned}
\frac{\partial{\mathcal{L}}}{\partial{w_i}} &= \frac{\partial{\mathcal{L}}}{\partial{y_i}} \cdot \frac{\partial{\mathcal{y_i}}}{\partial{w_i}} = \frac{\partial{\mathcal{L}}}{\partial{y_i}} \cdot \hat x_i \\
\frac{\partial{\mathcal{L}}}{\partial{b_i}} &= \frac{\partial{\mathcal{L}}}{\partial{y_i}} \cdot \frac{\partial{\mathcal{y_i}}}{\partial{b_i}} = \frac{\partial{\mathcal{L}}}{\partial{y_i}}
\end{aligned}
$$

对输入$x$的梯度等于：
$$
\frac{\partial{\mathcal{L}}}{\partial{x_i}} = \frac{1}{\sigma} \cdot \left[\frac{\partial{\mathcal{L}}}{\partial{y_i}} \cdot w_i - \frac{1}{m}\cdot \left( \sum_j {\frac{\partial{\mathcal{L}}}{\partial{y_j}} \cdot w_j} + \hat x_i \cdot \sum_j {\frac{\partial{\mathcal{L}}}{\partial{y_j}} \cdot w_j \cdot \hat x_j}\right)\right]
$$

推导过程

对输入$x$的梯度可以写成：
$$
\begin{aligned}
\frac{\partial{\mathcal{L}}}{\partial{x_i}} &= \sum_j {\frac{\partial{\mathcal{L}}}{\partial{y_j}} \cdot \frac{\partial{\mathcal{y_j}}}{\partial{\hat x_j}}} \cdot \frac{\partial{\mathcal{\hat x_j}}}{\partial{x_i}} \\
&= \sum_j {\frac{\partial{\mathcal{L}}}{\partial{y_j}} \cdot w_j \cdot \frac{\partial{\mathcal{\hat x_j}}}{\partial{x_i}}}
\end{aligned}
$$

这里只需要计算最后一项就行了：
$$
\begin{aligned}
\frac{\partial{\mathcal{\hat x_j}}}{\partial{x_i}} &= \frac{\partial}{\partial{x_i}}\left(\frac{x_j - \mu}{\sigma}\right) \\
&= \frac{\partial{(x_j - \mu)}}{\partial x_i} \cdot \sigma^{-1} + (x_j - \mu) \cdot \frac{\partial\sigma}{\partial x_i} \cdot (-\sigma^{-2}) \\
&= \left(\delta_{ij} - \frac{\partial{\mu}}{\partial x_i}\right) \cdot \sigma^{-1} - \sigma^{-2} \cdot (x_j - \mu) \cdot \frac{\partial\sigma}{\partial x_i}
\end{aligned}
$$

其中$\delta_{ij}$只有当$i=j$的时候才会等于1，否则都等于0。这里只需要求出均值和标准差对$x_i$的梯度就行了。直接给出结论，证明比较简单，放在了文末：
$$
\begin{aligned}
\frac{\partial \mu}{\partial x_i} &= \frac{1}{m} \\
\frac{\partial \sigma}{\partial x_i} &= \frac{1}{m} \cdot \sigma^{-1} \cdot (x_i - \mu)
\end{aligned}
$$

代入可以得到：
$$
\begin{aligned}
\frac{\partial{\mathcal{\hat x_j}}}{\partial{x_i}} &= \left(\delta_{ij} - \frac{\partial{\mu}}{\partial x_i}\right) \cdot \sigma^{-1} - \sigma^{-2} \cdot (x_j - \mu) \cdot \frac{\partial\sigma}{\partial x_i} \\
&= \left(\delta_{ij} - \frac{1}{m}\right) \cdot \sigma^{-1} - \sigma^{-2} \cdot (x_j - \mu) \cdot \frac{1}{m} \cdot \sigma^{-1} \cdot (x_i - \mu) \\
&= \sigma^{-1} \cdot \delta_{ij} - \frac{1}{m} \cdot \sigma^{-1} - \frac{1}{m} \cdot \sigma^{-3} \cdot (x_i - \mu) \cdot (x_j - \mu) \\
&= \sigma^{-1} \cdot \delta_{ij} - \frac{1}{m} \cdot \sigma^{-1} - \frac{1}{m} \cdot \sigma^{-1} \cdot {\hat x_i} \cdot {\hat x_j}
\end{aligned}
$$

最后带入梯度$\frac{\partial{\mathcal{L}}}{\partial{x_i}}$中可以得到：
$$
\begin{aligned}
\frac{\partial{\mathcal{L}}}{\partial{x_i}} &= \sum_j {\frac{\partial{\mathcal{L}}}{\partial{y_j}} \cdot w_j \cdot \frac{\partial{\mathcal{\hat x_j}}}{\partial{x_i}}} \\
&= \sum_j {\frac{\partial{\mathcal{L}}}{\partial{y_j}} \cdot w_j \cdot \left(\sigma^{-1} \cdot \delta_{ij} - \frac{1}{m} \cdot \sigma^{-1} - \frac{1}{m} \cdot \sigma^{-1} \cdot {\hat x_i} \cdot {\hat x_j}\right)} \\
&= \frac{1}{\sigma} \cdot \frac{\partial{\mathcal{L}}}{\partial{y_i}} \cdot w_i - \frac{1}{m \sigma} \cdot \sum_j {\frac{\partial{\mathcal{L}}}{\partial{y_j}} \cdot w_j} - \frac{\hat x_i}{m \sigma} \cdot \sum_j {\frac{\partial{\mathcal{L}}}{\partial{y_j}} \cdot w_j \cdot \hat x_j} \\
&= \frac{1}{\sigma} \cdot \left[\frac{\partial{\mathcal{L}}}{\partial{y_i}} \cdot w_i - \frac{1}{m}\cdot \left( \sum_j {\frac{\partial{\mathcal{L}}}{\partial{y_j}} \cdot w_j} + \hat x_i \cdot \sum_j {\frac{\partial{\mathcal{L}}}{\partial{y_j}} \cdot w_j \cdot \hat x_j}\right)\right]
\end{aligned}
$$

均值和标准差的梯度

均值的梯度为：
$$
\begin{aligned}
\frac{\partial \mu}{\partial x_i} &= \frac{\partial}{\partial x_i} \left(\frac{1}{m} \cdot \sum_j{x_j}\right) \\
&= \frac{1}{m}
\end{aligned}
$$

标准差的计算公式可以写成$\sigma = \left[\mu(x^2) - \mu^2(x)\right]^{\frac{1}{2}}$，所以梯度为：
$$
\begin{aligned}
\frac{\partial \sigma}{\partial x_i} &= \frac{\partial}{\partial x_i} \left[\mu(x^2) - \mu^2(x)\right]^{\frac{1}{2}} \\
&= \frac{1}{2} \cdot \left[\mu(x^2) - \mu^2(x)\right]^{-\frac{1}{2}} \cdot \left(\frac{2}{m} \cdot x_i - \frac{2}{m} \cdot \mu \right) \\
&= \frac{1}{m} \cdot \sigma^{-1} \cdot (x_i - \mu)
\end{aligned}
$$