《从几何视角来理解模型参数的初始化策略》《浅谈Transformer的初始化、参数化与标准化》等文章,我们讨论过模型的初始化方法,大致的思路是:如果一个n×n的方阵用均值为0、方差为1/n的独立同分布初始化,那么近似于一个正交矩阵,使得数据二阶矩(或方差)在传播过程中大致保持不变。

那如果是m×n的非方阵呢?常见的思路(Xavier初始化)是综合考虑前向传播和反向传播,所以使用均值为0、方差为2/(m+n)的独立同分布初始化。但这个平均更多是“拍脑袋”的,本文就来探究一下有没有更好的平均方案。

基础回顾 #

Xavier初始化是考虑如下的全连接层(设输入节点数为m,输出节点数为n
yj=bj+ixiwi,j
其中bj一般初始化为0,wi,j的初始化均值一般也为0,在《浅谈Transformer的初始化、参数化与标准化》中我们已经算得
E[y2j]=iE[x2i]E[w2i,j]=mE[x2i]E[w2i,j]
所以为了保持二阶矩不变,我们将wi,j的初始化方差设为1/m(均值为0时,方差等于二阶矩)。

但这个推导还只是考虑了前向传播,我们还需要使得模型有合理的梯度,那么还要使得模型在反向传播时也保持稳定。假设模型的损失函数为l,根据链式法则我们有
lxi=jlyjyjxi=jlyjwi,j
注意这时是对j求和,求和的维度为n,所以在相同的假设下有
E[(lxi)2]=jE[(lyj)2]E[w2i,j]=nE[(lyj)2]E[w2i,j]
所以要保持反向传播的二阶矩不变,我们将wi,j的初始化方差设为1/n

一个是1/m,一个1/n,当mn时就有冲突,但两个都同样重要,所以Xavier初始化就直接将两个维度平均一下,以2/(m+n)为方差进行初始化。

几何平均 #

现在让我们来考虑两个复合的全连接层(暂时忽略偏置项):
y=xW1W2
其中xRm,W1Rm×n,W2Rn×m,也就是说,输入是m维,变换为n维后再变换回m维,类似的操作比如BERT的FFN层(但FFN层中间多了个激活函数)。

根据前向传播的稳定性,我们应该要用1/m的方差初始化W1、用1/n的方差初始化W2。但是,如果我们要求W1W2必须用同一方差初始化呢?那么很显然,为了保证x,y的方差不变,W1,W2都需要用方差为1/mn的分布来初始化。如果考虑反向传播时,结果是相同的。

这样一来,我们就导出了一个新的维度平均策略:几何平均mn。通过这个维度平均策略,我们可以使得在多层网络复合的时候,如果输入输出维度不变,那么方差就保持不变(不管前向传播还是反向传播)。而如果是代数平均(m+n)/2,假设m<n,那么根据(m+n)2/4mn,在前向/反向传播的时候方差就会缩小了。

二次平均 #

另外一个思考的角度是作为一个双重最小化问题:假设选用的方差为t,在前向传播时我们希望(mt1)2尽可能小,在反向传播时我们则希望(nt1)2尽可能小,所以综合考虑
(mt1)2+(nt1)2
t=(m+n)/(m2+n2)时,上式取到最小值,所以这得到了一个二次分式的平均方案:(m2+n2)/(m+n)

容易证明:
m2+n2m+nm+n2mn
从推导过程上来看,左端的二次平均是希望每一步前向和反向传播的方差尽可能不变,因此可以认为左端是一个局部最优解;而右端的几何平均,则是希望“最初的输入”和“最终的输出”的方差尽量不变,因此可以认为右端某种意义上来说是一个全局最优解;而中间的代数平均,则是介乎全局最优和局部最优之间的一个解。

如此看来,似乎Xavier初始化“拍脑袋”的代数平均也不失为一个“中庸之道”的选择?

文章小结 #

本文简单思考了初始化方法中非方阵的维度平均方案,一直以来,大家似乎对默认的代数平均都没有什么疑问,而笔者从两种不同的角度得出了不同的平均策略的可能性。至于哪种平均策略更好,笔者也没有仔细做实验,有兴趣的读者自行尝试就好。当然,也可能在当前诸多优化策略之下,默认的初始化方案也工作得很好了,也就没有仔细调节的必要性了。

转载到请包括本文地址:https://spaces.ac.cn/archives/8725

更详细的转载事宜请参考:《科学空间FAQ》

如果您还有什么疑惑或建议,欢迎在下方评论区继续讨论。

如果您觉得本文还不错,欢迎分享/打赏本文。打赏并非要从中获得收益,而是希望知道科学空间获得了多少读者的真心关注。当然,如果你无视它,也不会影响你的阅读。再次表示欢迎和感谢!

如果您需要引用本文,请参考:

苏剑林. (Oct. 18, 2021). 《初始化方法中非方阵的维度平均策略思考 》[Blog post]. Retrieved from https://spaces.ac.cn/archives/8725

@online{kexuefm-8725,
        title={初始化方法中非方阵的维度平均策略思考},
        author={苏剑林},
        year={2021},
        month={Oct},
        url={\url{https://spaces.ac.cn/archives/8725}},
}