初始化方法中非方阵的维度平均策略思考

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

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

基础回顾 #

Xavier初始化是考虑如下的全连接层（设输入节点数为$m$，输出节点数为$n$）
\begin{equation} y_j = b_j + \sum_i x_i w_{i,j}\end{equation}
其中$b_j$一般初始化为0，$w_{i,j}$的初始化均值一般也为0，在《浅谈Transformer的初始化、参数化与标准化》中我们已经算得
\begin{equation}
\mathbb{E}[y_j^2] = \sum_{i} \mathbb{E}[x_i^2] \mathbb{E}[w_{i,j}^2]= m\mathbb{E}[x_i^2]\mathbb{E}[w_{i,j}^2]\end{equation}
所以为了保持二阶矩不变，我们将$w_{i,j}$的初始化方差设为$1/m$（均值为0时，方差等于二阶矩）。

但这个推导还只是考虑了前向传播，我们还需要使得模型有合理的梯度，那么还要使得模型在反向传播时也保持稳定。假设模型的损失函数为$l$，根据链式法则我们有
\begin{equation}\frac{\partial l}{\partial x_i} = \sum_j \frac{\partial l}{\partial y_j} \frac{\partial y_j}{\partial x_i}=\sum_j \frac{\partial l}{\partial y_j} w_{i,j}\end{equation}
注意这时是对$j$求和，求和的维度为$n$，所以在相同的假设下有
\begin{equation}
\mathbb{E}\left[\left(\frac{\partial l}{\partial x_i}\right)^2\right] = \sum_{j} \mathbb{E}\left[\left(\frac{\partial l}{\partial y_j}\right)^2\right] \mathbb{E}[w_{i,j}^2]= n \mathbb{E}\left[\left(\frac{\partial l}{\partial y_j}\right)^2\right]\mathbb{E}[w_{i,j}^2]\end{equation}
所以要保持反向传播的二阶矩不变，我们将$w_{i,j}$的初始化方差设为$1/n$。

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

几何平均 #

现在让我们来考虑两个复合的全连接层（暂时忽略偏置项）：
\begin{equation} y = xW_1 W_2
\end{equation}
其中$x\in\mathbb{R}^m,W_1\in\mathbb{R}^{m\times n},W_2\in\mathbb{R}^{n\times m}$，也就是说，输入是$m$维，变换为$n$维后再变换回$m$维，类似的操作比如BERT的FFN层（但FFN层中间多了个激活函数）。

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

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

二次平均 #

另外一个思考的角度是作为一个双重最小化问题：假设选用的方差为$t$，在前向传播时我们希望$(mt-1)^2$尽可能小，在反向传播时我们则希望$(nt-1)^2$尽可能小，所以综合考虑
\begin{equation}(mt-1)^2 + (nt-1)^2
\end{equation}
当$t=(m+n)/(m^2+n^2)$时，上式取到最小值，所以这得到了一个二次分式的平均方案：$(m^2+n^2)/(m+n)$。

容易证明：
\begin{equation}\frac{m^2+n^2}{m+n} \geq \frac{m+n}{2}\geq \sqrt{mn}\end{equation}
从推导过程上来看，左端的二次平均是希望每一步前向和反向传播的方差尽可能不变，因此可以认为左端是一个局部最优解；而右端的几何平均，则是希望“最初的输入”和“最终的输出”的方差尽量不变，因此可以认为右端某种意义上来说是一个全局最优解；而中间的代数平均，则是介乎全局最优和局部最优之间的一个解。

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

文章小结 #

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

转载到请包括本文地址：https://kexue.fm/archives/8725

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