梯度消失与梯度爆炸原理&初始化方法

基础概念

在详细了解神经网络中梯度消失与梯度爆炸的原理之前，先来回顾一下如下基本公式：
$\begin{aligned} &\mathbf{E}(\boldsymbol{X} * \boldsymbol{Y})=\boldsymbol{E}(\boldsymbol{X}) * \boldsymbol{E}(\boldsymbol{Y})\\ &\mathrm{D}(\boldsymbol{X})=\boldsymbol{E}\left(\mathrm{X}^{2}\right)-[\boldsymbol{E}(\boldsymbol{X})]^{2}\\ &\mathbf{D}(\boldsymbol{X}+\boldsymbol{Y})=\boldsymbol{D}(\boldsymbol{X})+\boldsymbol{D}(\boldsymbol{Y}) \end{aligned}$
其中随机变量X与Y相互独立，根据上面三个公式可推导得出下式关系
$\mathrm{D}(\mathrm{X} * \mathrm{Y})=\mathrm{D}(\mathrm{X}) * \mathrm{D}(\mathrm{Y})+\mathrm{D}(\mathrm{X}) *[\boldsymbol{E}(\boldsymbol{Y})]^{2}+\mathrm{D}(\mathrm{Y}) *[\boldsymbol{E}(\boldsymbol{X})]^{2}$
至此，基础知识准备完毕。

神经网络中的梯度消失与爆炸

面对如下结构的神经网络：

先来计算梯度，令$H_1$表示第一个隐藏层，$H_2$表示第二个隐藏层，W代表权重矩阵，那么第二层隐藏层到输出层之间的权重矩阵$W_2$的梯度为：
$\begin{aligned} \Delta \mathrm{W}_{2} &=\frac{\partial L \text { oss }}{\partial \mathrm{W}_{2}}=\frac{\partial L \text { oss }}{\partial o u t} * \frac{\partial o u t}{\partial H_{2}} * \frac{\partial H_{2}}{\partial w_{2}} \\ &=\frac{\partial L o s s}{\partial o u t} * \frac{\partial o u t}{\partial H_{2}} * H_{1} \end{aligned}$
其中倒数第二步可以推得倒数第一步是因为
$\mathrm{H}_{2}=\mathrm{H}_{1} * \mathrm{W}_{2}$
如果对求导部分不明白，可参考矩阵求导的公式。

由此可见，权重矩阵$W_2$的梯度和第一层隐层的输出值有关，那如果第一层隐层的输出值过大或者过小，就引出了梯度消失或者爆炸的初步概念
$\begin{aligned} &\mathrm{H}_{1} \rightarrow 0 \Rightarrow \Delta \mathrm{W}_{2} \rightarrow 0\\ &\mathrm{H}_{1} \rightarrow \infty \Rightarrow \Delta \mathrm{W}_{2} \rightarrow \infty \end{aligned}$
似乎梯度消失与爆炸现象只与前一层的输出值有关，那自然就想到对权重进行随机初始化的时候令方差和均值小一些，以此控制好输出值范围就行了。可是仅仅这样那为什么层数加深的时候，一些不合适的初始化方法依旧会引起梯度消失或者爆炸呢？
现在假设输入层和隐藏层各自的每个节点的值都满足某种概率分布，先来观察第一层隐藏层的第一个节点是如何计算的：
$\mathbf{H}_{\mathbf{1 1}}=\sum_{i=\mathbf{0}}^{\mathbf{n}} \boldsymbol{X}_{\boldsymbol{i}} * \boldsymbol{W}_{\mathbf{1} i}$
根据这个公式，结合前面基础知识中的基本公式，即可求得第一层隐藏层第一个节点，在（符合某种概率分布初始化的）权重计算下的值的方差：
$\begin{aligned} \mathbf{D}\left(\mathbf{H}_{\mathbf{1 1}}\right) &=\sum_{i=0}^{n} D\left(X_{i}\right) * D\left(\boldsymbol{W}_{\mathbf{1 i}}\right) \\ &=\mathbf{n} *(1 * 1) \\ &=\mathbf{n} \end{aligned}$
上式假设了对每层的权重矩阵W采取均值为0，方差为1的正态分布下的初始化（为便于理解这里假设输入也是满足这样的正态分布）。在均值为0方差为1时，基础知识中最后一个公式就变为这样的形式：
$D(X * Y)=D(X) * D(Y)$
所以最后求得方差是n。那么对于第一层的每个节点（每个节点都和第一个节点是一样的情况），其标准差为
$\operatorname{std}\left(\mathrm{H}_{1}\right)=\sqrt{\mathrm{D}\left(\mathrm{H}_{11}\right)}=\sqrt{n}$
这样第一层的输出值就是满足标准差为$\sqrt{n}$的分布了，以此类推，第二层的输出值就满足标准差为n的分布，如果还有更多层，第三层的标准差是第二层标准差的$\sqrt{n}$倍，以此类推，梯度就会增长至超过计算机的计算范围。
而梯度消失也是同样的道理，如果方差过小，每一层新乘的方差都是小于一或着还要小的数，随着层数越来越深，最后梯度就贴近0，也就是说权值根本不会根据误差的大小而有所更新了。

解决方法

那在层数加深的时候有没有什么解决办法呢？首先考虑最简单的方法，想到每一层的输出值的方差如果能维持不变，对后面没有增长或者降低的累积效应就好了，由此考虑某层输出值的方差是否能等于1，即
$\mathbf{D}\left(\mathbf{H}_{\mathbf{1}}\right)=\boldsymbol{n} * \boldsymbol{D}(\boldsymbol{X}) * \boldsymbol{D}(\boldsymbol{W})=\mathbf{1}$
那取
$D(W)=\frac{1}{n} \Rightarrow \operatorname{std}(W)=\sqrt{\frac{1}{n}}$
就是一个简单的解决办法，根据这个简单的考量，诞生了几种初始化方法。

Xavier初始化

概要：面向饱和激活函数例如sigmoid与tanh的初始化方法。

Xavier在其论文Understanding the difficulty of training deep feedforward neural networks中详细探讨了神经网络具有激活函数时该如何进行初始化，主要利用方差一致性原则（也就是让每个网络层输出值的方差尽可能为1），保持数据尺度维持在恰当范围。如下是从文章中得到的公式：
$\begin{aligned} &\boldsymbol{n}_{\boldsymbol{i}} * \boldsymbol{D}(\boldsymbol{W})=\mathbf{1}\\ &\boldsymbol{n}_{\boldsymbol{i}+\boldsymbol{1}} * \boldsymbol{D}(\boldsymbol{W})=\mathbf{1}\\ &\Rightarrow D(W)=\frac{2}{n_{i}+n_{i+1}} \end{aligned}$
也就是让权值的方差乘$n_i$（当前权值前接的输入层的神经元个数）或者$n_{i+1}$（当前权值后接的输出层的神经元个数）时都为1。这是同时考虑了前向传播与反向传播的数据尺度问题。
不妨看看这样初始化情况下，均匀分布权值范围的上限和下限a的情况（对称是因为假设均值为0）：
$\begin{aligned} &\boldsymbol{W} \sim \boldsymbol{U}[-\boldsymbol{a}, \boldsymbol{a}]\\ &D(W)=\frac{(-a-a)^{2}}{12}=\frac{(2 a)^{2}}{12}=\frac{a^{2}}{3}\\ &\frac{2}{n_{i}+n_{i+1}}=\frac{a^{2}}{3} \Rightarrow a=\frac{\sqrt{6}}{\sqrt{n_{i}+n_{i+1}}}\\ &\Rightarrow W \sim U\left[-\frac{\sqrt{6}}{\sqrt{n_{i}+n_{i+1}}}, \frac{\sqrt{6}}{\sqrt{n_{i}+n_{i+1}}}\right] \end{aligned}$
上式中方差求解过程参考均匀分布的方差公式。最后结论得出是令两个关于$D(W)$的等式相等，可解得a。
由此利用a范围内的均匀分布来初始化权重，即可稳定梯度。但是注意，如果是非饱和激活函数，最后数据的尺度还是会越来越大，Xavier初始化则不再适用。

Xavier初始化在pytorch中调用语句就是torch.nn.init.xavier_uniform_()

Kaiming初始化

概要：针对relu及其变种激活函数的初始化方法

同样保持了方差一致性原则，保持数据尺度维持在恰当返回，令方差为1。通过公式推导，得出权值的方差应该为：
$D(W)=\frac{2}{n_i}$
这里的$n_i$还是输入层神经元个数。如果是relu的变种，在负半轴有一定斜率的激活函数，方差则更改为：
$\mathrm{D}(W)=\frac{2}{\left(1+a^{2}\right) * n_{i}}$
这里的a就是负半轴的斜率。采用均值为0，方差为$D(W)$的正态分布初始化权重，即可将输出维持在一定数据尺度上了。

Kaiming初始化在pytorch中的调用语句是torch.nn.init.kaiming_normal_()

其他初始化方法

除上述方法外，还有Xavier均匀分布、Xavier正态分布、kaiming均匀分布、kaiming正态分布、均匀分布、正态分布、常数分布、正交矩阵初始化、单位矩阵初始化、稀疏矩阵初始化。无论使用哪种初始化方法，其核心都是尽力使得输出值方差为1。
在pytorch中，也有相应的函数可以计算激活函数的方差变化尺度，就是torch.nn.init.calculate_gain()函数。只要给函数输入激活函数的名称以及参数，即可计算得出方差变化的尺度。这个尺度本质上就是输入数据的方差除以经过激活函数之后输出数据的方差，是两者的一个比值。