ZerO初始化: 仅使用0和1初始化神经网络.

对神经网络进行训练时,需要对神经网络的参数进行初始化。对于深度网络来说,参数的初始化显得尤为重要。糟糕的初始化不仅会使模型效果变差,还有可能使得模型根本训练不动或者不收敛。

随机初始化的关键是设置方差$σ^2$的大小。

1. 恒等初始化

在参数初始化时,如果让各层之间的权重完全相等,并且使得上一层的输入“完整”的传入下一层,则神经网络各层参数之间的方差不会发生变化。恒等初始化是指通过把神经网络的权重层初始化为一个单位矩阵(恒等变换),使得网络层的输出值与输出值相等。

恒等初始化具有动力等距(Dynamical Isometry)性质,使得神经网络具有稳定的信号传播以及梯度下降的行为。然而恒等初始化建立在各层的维度是相等的假设之上,在实际中这种假设有些过强。

当各层的输入输出维度不相等时,可以把参数矩阵(非方阵)初始化为部分单位矩阵 (Partial Identity Matrix),对于行列中“超出”的部分补零即可:

\[\mathbf{I}^* = \begin{cases} [\mathbf{I}, \mathbf{0}], & \mathbf{I} \in \mathbb{R}^{m\times m},\mathbf{0} \in \mathbb{R}^{m\times n-m},m < n \\ [\mathbf{I}, \mathbf{0}]^T, & \mathbf{I} \in \mathbb{R}^{n\times n},\mathbf{0} \in \mathbb{R}^{m-n\times n},m > n \\ \mathbf{I}, & \text{otherwise} \end{cases}\]

然而当使用部分单位矩阵在训练神经网络时,会出现训练衰减 (Training Degeneracy)现象,即无论隐藏层维度$N_h$有多高,$N_h>N_x$部分的输入在激活函数阶段无法生效,导致神经网络的维度仅仅依赖于输入数据的维度$N_x$,从而极大的限制了神经网络的表达能力。

假设\(\mathcal{F}\)是一个$L$层的神经网络,对于$l_1$有$W_1 \in R^{N_h\times N_x}$,而对于$1<l<L$,有$W_l \in R^{N_h\times N_h}$,对于$l=L$有$W_L \in R^{N_y\times N_h}$。假设$N_h>N_x,N_y$,令$z_l(\cdot)$为第$l$层的激活函数,当初始化\(W_1,W_L = I^{\*},W_l=I\)时,对于任意$x \in R^{N_x}$,有:

\[dim(span(z_l(x) | x \in R^{N_x})) \leq N_x\]

如下图所示,使用部分单位矩阵初始化一个3层的神经网络,在MINST上进行训练,红色虚线表示$N_x$的维度为$784$。可以看出无论$N_h$的维度有多大,$W_2$的秩在训练过程在始终低于$N_x$的维度。

2. ZerO初始化

为了避免直接使用部分单位矩阵作为初始权重参数进行训练而出现的训练衰减问题,本文作者提出了使用哈达玛变换(Hadamard Transform)来初始化权重参数,哈达玛变换即使用哈达玛矩阵进行的线性变换,哈达玛矩阵是均由$+1$与$-1$的元素构成,且满足$H_nH_n^T=nI_n$。哈达玛矩阵可以通过递归的构造得到,设$H_0=1$,则有:

\[H_m = \begin{pmatrix} H_{m-1} & H_{m-1}\\ H_{m-1} & -H_{m-1}\\ \end{pmatrix} = \begin{pmatrix} 1 & 1 & 1& 1 & \cdots\\ 1 & -1 & 1& -1 & \cdots\\ 1 & 1 & -1& -1 & \cdots\\ 1 & -1 & -1& 1 & \cdots\\ \vdots & \vdots & \vdots & \vdots & \ddots \end{pmatrix} \in R^{2m\times 2m}\]

在二维平面中,哈达玛变换可以将标准坐标轴旋转45度,论文作者证明了针对部分单位矩阵应用哈达玛变换,即当初始化权重为\(W_1,W_L = HI^{\*},W_l=I\)时,可以有效的规避训练衰减对神经网络训练带来的伤害,将有:

\[dim(span(z_l(x) | x \in R^{N_x})) \geq N_x\]

应用哈达玛变换的权重神奇的打破了训练衰减,本质上来讲,部分单位矩阵的训练衰减主要是源自于“补零”的操作,使得这些位置的输入在激活函数阶段无法生效,从而使得维度被$N_x$限制,而哈达玛变换通过将基向量“旋转”,打破了在传递过程中零元素的对称性,从而解决了训练衰减的问题。

通过哈达玛变换解决训练衰减的问题后,结合恒等初始化,作者构建了ZerO初始化方法,具体算法步骤如下图所示:

ZerO初始化中,由于部分单位矩阵与哈达玛矩阵都是确定性的,训练得到的结果在重复训练的过程中变换程度更低,因此也使得使用ZerO方式训练出的模型更具可复现性。

3. 实验分析

以上文中MNIST上训练的三层网络为例,下图展现了哈达玛变换打破训练衰减的直观图例:

可以看到,使用哈达玛变换作为初始权重的训练过程与使用随机初始化权重的训练过程有很大区别,相较于随机初始化权重一开始就训练一个较为“复杂”(权重秩很大)的网络,ZerO一开始似乎在训练一个更为“简单”的网络,并在不断的学习过程中逐渐使得网络变得“复杂”。即ZerO初始化的网络具有低秩的学习轨迹(Low-Rank Learning Trajectory)。

为了展示这种现象,作者通过定义:

\[\frac{||W||_F^2}{||W||_2^2} = \sum_{i=1}^k \sigma_i^2(W)\sigma_{\max}^2(W)\]

来衡量稳定时的权重矩阵的秩,从而侧面反映网络的复杂度,此处的$\sigma$表示矩阵的奇异值,通过计算ZerO与其他随机初始化方法的稳定秩,可以得到如下图的结果:

上层图表示 ResNet-18CIFAR-10 的训练结果,下层图表示 ResNet-50ImageNet 的训练结果,从左到右分别提取了 ResNet 第二、三、四组残差块的第一层卷积,可以看到,ZerO初始化具有明显的低秩学习的特征。

这种贪心的低秩学习(greedy low-rank learning,GLRL)现象表明梯度下降隐含地偏好于简单的解决方案,即在梯度下降过程中,偏好以权重矩阵的秩递增的顺序在解空间中搜索,在低秩空间无法找到最小值后,才转去高秩空间进行搜索。GLRL现象可以帮助解释基于梯度下降法的神经网络卓越的泛化能力以及经常收敛到低秩的全局或局部最优解之中。