4.1 权值初始化
最后更新于
最后更新于
本章代码：https://github.com/zhangxiann/PyTorch_Practice/blob/master/lesson4/grad_vanish_explod.py
在搭建好网络模型之后，一个重要的步骤就是对网络模型中的权值进行初始化。适当的权值初始化可以加快模型的收敛，而不恰当的权值初始化可能引发梯度消失或者梯度爆炸，最终导致模型无法收敛。下面分 3 部分介绍。第一部分介绍不恰当的权值初始化是如何引发梯度消失与梯度爆炸的，第二部分介绍常用的 Xavier 方法与 Kaiming 方法，第三部分介绍 PyTorch 中的 10 种初始化方法。
考虑一个 3 层的全连接网络。
$H_{1}=X \times W_{1}$，$H_{2}=H_{1} \times W_{2}$，$Out=H_{2} \times W_{3}$
其中第 2 层的权重梯度如下：
$\begin{aligned} \Delta \mathrm{W}{2} &=\frac{\partial \mathrm{Loss}}{\partial \mathrm{W}{2}}=\frac{\partial \mathrm{Loss}}{\partial \mathrm{out}} \frac{\partial \mathrm{out}}{\partial \mathrm{H}_{2}} \frac{\partial \mathrm{H}{2}}{\partial \mathrm{w}{2}} \ &=\frac{\partial \mathrm{Loss}}{\partial \mathrm{out}} \frac{\partial \mathrm{out}}{\partial \mathrm{H}_{2}} \mathrm{H}_{1} \end{aligned}$
所以$\Delta \mathrm{W}{2}$依赖于前一层的输出$H{1}$。如果$H_{1}$ 趋近于零，那么$\Delta \mathrm{W}{2}$也接近于 0，造成梯度消失。如果$H{1}$ 趋近于无穷大，那么$\Delta \mathrm{W}_{2}$也接近于无穷大，造成梯度爆炸。要避免梯度爆炸或者梯度消失，就要严格控制网络层输出的数值范围。
下面构建 100 层全连接网络，先不适用非线性激活函数，每层的权重初始化为服从$N(0,1)$的正态分布，输出数据使用随机初始化的数据。
输出为：
也就是数据太大(梯度爆炸)或者太小(梯度消失)了。接下来我们在forward()函数中判断每一次前向传播的输出的标准差是否为 nan，如果是 nan 则停止前向传播。
输出如下：
可以看到每一层的标准差是越来越大的，并在在 31 层时超出了数据可以表示的范围。
下面推导为什么网络层输出的标准差越来越大。
首先给出 3 个公式：
$E(X \times Y)=E(X) \times E(Y)$：两个相互独立的随机变量的乘积的期望等于它们的期望的乘积。
$D(X)=E(X^{2}) - [E(X)]^{2}$：一个随机变量的方差等于它的平方的期望减去期望的平方
$D(X+Y)=D(X)+D(Y)$：两个相互独立的随机变量之和的方差等于它们的方差的和。
可以推导出两个随机变量的乘积的方差如下：
$D(X \times Y)=E[(XY)^{2}] - [E(XY)]^{2}=D(X) \times D(Y) + D(X) \times [E(Y)]^{2} + D(Y) \times [E(X)]^{2}$
如果$E(X)=0$，$E(Y)=0$，那么$D(X \times Y)=D(X) \times D(Y)$
我们以输入层第一个神经元为例：
$\mathrm{H}{11}=\sum{i=0}^{n} X_{i} \times W_{1 i}$
其中输入 X 和权值 W 都是服从$N(0,1)$的正态分布，所以这个神经元的方差为：
$\begin{aligned} \mathbf{D}\left(\mathrm{H}{11}\right) &=\sum{i=0}^{n} \boldsymbol{D}\left(X_{i}\right) * \boldsymbol{D}\left(W_{1 i}\right) \ &=n (1 1) \ &=n \end{aligned}$
标准差为：$\operatorname{std}\left(\mathrm{H}{11}\right)=\sqrt{\mathbf{D}\left(\mathrm{H}{11}\right)}=\sqrt{n}$，所以每经过一个网络层，方差就会扩大 n 倍，标准差就会扩大$\sqrt{n}$倍，n 为每层神经元个数，直到超出数值表示范围。对比上面的代码可以看到，每层神经元个数为 256，输出数据的标准差为 1，所以第一个网络层输出的标准差为 16 左右，第二个网络层输出的标准差为 256 左右，以此类推，直到 31 层超出数据表示范围。可以把每层神经元个数改为 400，那么每层标准差扩大 20 倍左右。从$D(\mathrm{H}{11})=\sum{i=0}^{n} D(X_{i}) \times D(W_{1 i})$，可以看出，每一层网络输出的方差与神经元个数、输入数据的方差、权值方差有关，其中比较好改变的是权值的方差$D(W)$，所以$D(W)= \frac{1}{n}$，标准差为$std(W)=\sqrt\frac{1}{n}$。因此修改权值初始化代码为nn.init.normal_(m.weight.data, std=np.sqrt(1/self.neural_num)),结果如下：
修改之后，没有出现梯度消失或者梯度爆炸的情况，每层神经元输出的方差均在 1 左右。通过恰当的权值初始化，可以保持权值在更新过程中维持在一定范围之内，不过过大，也不会过小。
上述是没有使用非线性变换的实验结果，如果在forward()中添加非线性变换tanh，每一层的输出方差还是会越来越小，会导致梯度消失。因此出现了 Xavier 初始化方法与 Kaiming 初始化方法。
Xavier 是 2010 年提出的，针对有非线性激活函数时的权值初始化方法，目标是保持数据的方差维持在 1 左右，主要针对饱和激活函数如 sigmoid 和 tanh 等。同时考虑前向传播和反向传播，需要满足两个等式：$\boldsymbol{n}{\boldsymbol{i}} * \boldsymbol{D}(\boldsymbol{W})=\mathbf{1}$和$\boldsymbol{n}{\boldsymbol{i+1}} * \boldsymbol{D}(\boldsymbol{W})=\mathbf{1}$，可得：$D(W)=\frac{2}{n_{i}+n_{i+1}}$。为了使 Xavier 方法初始化的权值服从均匀分布，假设$W$服从均匀分布$U[-a, 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}$，解得：$\boldsymbol{a}=\frac{\sqrt{6}}{\sqrt{n_{i}+n_{i+1}}}$，所以$W$服从分布$U\left[-\frac{\sqrt{6}}{\sqrt{n_{i}+n_{i+1}}}, \frac{\sqrt{6}}{\sqrt{n_{i}+n_{i+1}}}\right]$
所以初始化方法改为：
并且每一层的激活函数都使用 tanh，输出如下：
可以看到每层输出的方差都维持在 0.6 左右。
PyTorch 也提供了 Xavier 初始化方法，可以直接调用：
上面的初始化方法都使用了tanh_gain = nn.init.calculate_gain('tanh')。
x = torch.randn(10000) out = torch.tanh(x)
gain = x.std() / out.std() print('gain:{}'.format(gain))
tanh_gain = nn.init.calculate_gain('tanh') print('tanh_gain in PyTorch:', tanh_gain)
gain:1.5982500314712524 tanh_gain in PyTorch: 1.6666666666666667
```
结果表示，原有数据分布的方差经过 tanh 之后，标准差会变小 1.6倍左右。
虽然 Xavier 方法提出了针对饱和激活函数的权值初始化方法，但是 AlexNet 出现后，大量网络开始使用非饱和的激活函数如 ReLU 等，这时 Xavier 方法不再适用。2015 年针对 ReLU 及其变种等激活函数提出了 Kaiming 初始化方法。
针对 ReLU，方差应该满足：$\mathrm{D}(W)=\frac{2}{n_{i}}$；针对 ReLu 的变种，方差应该满足：$\mathrm{D}(W)=\frac{2}{n_{i}}$，a 表示负半轴的斜率，如 PReLU 方法，标准差满足$\operatorname{std}(W)=\sqrt{\frac{2}{\left(1+a^{2}\right) * n_{i}}}$。代码如下：`nn.init.normal_(m.weight.data, std=np.sqrt(2 / self.neural_num))，或者使用 PyTorch 提供的初始化方法：nn.init.kaiming_normal_(m.weight.data)`，同时把激活函数改为 ReLU。
PyTorch 中提供了 10 中初始化方法
Xavier 均匀分布
Xavier 正态分布
Kaiming 均匀分布
Kaiming 正态分布
均匀分布
正态分布
常数分布
正交矩阵初始化
单位矩阵初始化
稀疏矩阵初始化
每种初始化方法都有它自己使用的场景，原则是保持每一层输出的方差不能太大，也不能太小。
参考资料
如果你觉得这篇文章对你有帮助，不妨点个赞，让我有更多动力写出好文章。
我的文章会首发在公众号上，欢迎扫码关注我的公众号张贤同学。
