[深度学习] 使用Xavier进行权重初始化的原理以及caffe实现

0.背景说明

这几天开始在看caffe源码, 看到官方代码中有mnist的examples. 其中使用的网络模型是lenet. 在layer name:conv1中使用xavier对该层layer进行权重初始化, 因此本文探究一下xavier初始化方法的原理作用以及caffe代码实现方法.

1.Xavier介绍以及推导

1.0.符号表达式定义以及采用的模型

为方便公式推导，使用一个简单的三层神经网络；为方便描述，规定以下符号表达式定义：

参数名	含义
$N^l$	第$l$层单元数量
$A^l$	第$l$层激活值向量，维度为$[N^l\times 1]$
$a_j^l$	$A^l$中的元素，代表第$l$层第$j$个单元的激活值
$W^l$	第$l$层的权重矩阵，维度为$[N^l,N^{l+1}]$
$w_{jk}^l$	$W^l$中的元素，代表第$l-1$层的第$j$个单元连接到第$l$层第$k$个单元的权重
$B^l$	第$l$层偏置项向量,维度为$[N^l]$
$b_k^l$	$B^l$中的元素，代表第$l$层第$k$个单元的偏置项
$Z^l$	第$l$层激活函数的加权输入向量，即$Z^l=W^l\times A^{l-1}+B^l$
$z_j^l$	$Z^l$中的元素，代表第$l$层第$j$个单元的加权输入
$C$	尝试优化的损失函数,本文使用平方损失函数$\frac{1}{2}(Y-O{)}^2$作为损失函数
$\sigma$	激活函数,因此 ${\sigma }^l=\sigma (Z^l)$，其中该函数应用到输入向量的每一个元素。
$X$	神经网络输入向量
${\delta }^l$	${\delta }^l=\frac{\partial C}{\partial Z^l}$为损失函数对第 L 层加权输入向量的梯度，同样也成为误差方向。
${\delta }_k^l$	${\delta }^l=\frac{\partial C}{\partial z_k^l}$为损失函数对第$l$层第$k$个单元的加权输入的梯度
$M$	网络总层数

网络模型图：

1.1.前向传播

假设第$l$层网络的输入为第$l-1$层的激活值向量$A^{l-1}$，向量维度为$[N^{l-1}\times 1]$
以第$l=2$为例：

输入	加权计算	激活	输出
$a_1^1,a_2^1$	$z_1^2=w_{11}^2\ast a_1^1+w_{21}^2\ast a_2^1+b_1^2$	$a_1^2=\sigma (z_1^2)$	$a_1^2$
$a_1^1,a_2^1$	$z_2^2=w_{12}^2\ast a_1^1+w_{22}^2\ast a_2^1+b_2^2$	$a_2^2=\sigma (z_2^2)$	$a_2^2$
$a_1^1,a_2^1$	$z_3^2=w_{13}^2\ast a_1^1+w_{23}^2\ast a_2^1+b_3^2$	$a_3^2=\sigma (z_3^2)$	$a_3^2$

可以得出$l$层输出为:
$$A^l=\sum _{k=1}^{N^l}\sum _{j=1}^{N^{l-1}}\sigma (w_{jk}^l\ast a_j^{l-1}+b_k^l)\text{\hspace{1em}(1)}$$
考虑计算第$l$层的第$k$个神经元的加权输入$z_k^l$的方差:
$$\begin{array}{rl}D(z_k^l)& =D(\sum _{j=1}^{N^{l-1}}{w}_{jk}^l{a}_j^{l-1}+b_k^l)\end{array}\text{\hspace{1em}(2)}$$
将偏置项初始化为$0$, 则:
$$\begin{array}{rl}D(z_k^l)& =D(\sum _{j=1}^{N^{l-1}}{w}_{jk}^l{a}_j^{l-1})\end{array}\text{\hspace{1em}(3)}$$
由

“独立变量和的方差等于独立变量的方差的和”

得到:
$$\begin{array}{rl}D(z_k^l)& =\sum _{j=1}^{N^{l-1}}D(w_{jk}^l{a}_j^{l-1})\end{array}\text{\hspace{1em}(4)}$$
假设权重参数和$l-1$层激活值相互独立, 由接下来推导两个独立变量乘积项的方差,推导如下:
由方差定义可得
$$\begin{array}{rl}D(xw)& =E((xw-E(xw){)}^2)\\ & =E(x^2{w}^2-2xwE(xw)+(E(xw){)}^2)\\ & =E(x^2{w}^2)-2(E(kw){)}^2+(E(xw){)}^2\\ & =E(x^2{w}^2)-(E(xw){)}^2\end{array}\text{\hspace{1em}(5-1)}$$
其中$D()$代表方差, $E()$代表期望值,在这里也可以等价于平均值
假设$x,w$相互独立,有
$$\begin{array}{r}E(x^2{w}^2)=E(x^2)E(w^2)\end{array}\text{\hspace{1em}(5-2)}$$
$$\begin{array}{r}E(xw)=E(x)E(w)\end{array}\text{\hspace{1em}(5-3)}$$
结合式$(5-2)(5-3)$有
$$\begin{array}{r}D(xw)=E(x^2)E(w^2)-(E(xw){)}^2\end{array}\text{\hspace{1em}(5-3)}$$
同样的由方差定义得
$$\begin{array}{rl}D(x)& =E((x-E(x){)}^2)\\ & =E(x^2-2xE(x)+E(x){)}^2\\ & =E(x^2)-2(E(x){)}^2+E(x){)}^2\\ & =E(x^2)-E(x{)}^2\end{array}\text{\hspace{1em}(5-4)}$$
可得
$$\begin{array}{r}E(x^2)=D(x)+E(x{)}^2\end{array}\text{\hspace{1em}(5-5)}$$
同样的
$$\begin{array}{r}E(w^2)=D(w)+E(w{)}^2\end{array}\text{\hspace{1em}(5-6)}$$
代入式(5-3)得
$$\begin{array}{rl}D(xw)=& D(x)D(w)+E(x{)}^{2}D(w)+\\ & D(x)E(w{)}^2+E(x{)}^{2}E(w{)}^2-(E(xw){)}^2\\ =& D(x)D(w)+E(x{)}^{2}D(w)+D(x)E(w{)}^2\end{array}\text{\hspace{1em}(5-7)}$$
推导结束, 代入到式(4),且假设初始化权重和激活值服从均值为$0$的分布,得到:
$$\begin{array}{rl}D(z_k^l)& =\sum _{j=1}^{N^{l-1}}D(w_{jk}^l)D(a_j^{l-1})\\ & =\sum _{j=1}^{N^{l-1}}D(w_k^l)D(A^{l-1})\\ & =N^{l-1}D(w_k^l)D(A^{l-1})\end{array}\text{\hspace{1em}(6)}$$
考虑
$$\begin{array}{rl}{N}^{l}D(Z^l)& =\sum _{k=1}^{N^l}D(z_k^l)\\ & =\sum _{k=1}^{N^l}{N}^{l-1}D(w_k^l)D(A^{l-1})\\ & =\sum _{k=1}^{N^l}{N}^{l-1}D(W^l)D(A^{l-1})\\ & =N^l{N}^{l-1}D(W^l)D(A^{l-1})\end{array}\text{\hspace{1em}(7)}$$
可得$l$层加权输入的方差$D(Z^l)$:
$$\begin{array}{rl}D(Z^l)& =N^{l-1}D(W^l)D(A^{l-1})\end{array}\text{\hspace{1em}(8)}$$
关于激活函数$\sigma$我们做如下假设：

• 激活函数对称；
• $\sigma \prime (0)=1$；
• 初始时，状态值落在激活函数的线性区域：$\sigma \prime (sik)\approx 1$。

三个关于激活函数的假设，称为$Glorot$激活函数假设。基于这三个假设， 我们可以认为激活函数在初始化时候近似于恒等变换，即第$l$层加权输入和激活函数输出近似相等$z^l=a^l$.
此时，有：
$$\begin{array}{rl}D(Z^l)& =N^{l-1}D(W^l)D(A^{l-1})\\ & =N^{l-1}D(W^l)D(Z^{l-1})\\ & =N^{l-1}D(W^l)N^{l-2}D(W^{l-1})D(Z^{l-2})\\ & =D(X)\prod _{i=1}^{l-1}{N}^{i}D(W^{i+1})\end{array}\text{\hspace{1em}(9)}$$
我们希望每一层的加权输入方差$D(Z^l)$都相等, 只需要确保每一层$N^{l-1}D(W^l)$等于$1$，那么只需要满足：
$$D(W^l)=\frac{1}{N^{l-1}}\text{\hspace{1em}(10)}$$
这里$N^{l-1}$代表$l$层的输入端数量。

1.2.反向传播

模型图：
同样的，以第$1$层为例：
考虑第一层第一个单元的加权输入对于损失函数的梯度：
$$\begin{array}{rl}{\delta }_1^1& =\frac{\partial C}{\partial z_1^1}\\ & =\frac{\partial C}{\partial a_1^1}\ast \frac{\partial a_1^1}{\partial z_1^1}\\ & =(\frac{\partial C}{\partial z_1^2}\ast \frac{\partial z_1^2}{\partial a_1^1}+\frac{\partial C}{\partial z_2^2}\ast \frac{\partial z_2^2}{\partial a_1^1}+\frac{\partial C}{\partial z_3^2}\ast \frac{\partial z_3^2}{\partial a_1^1})\ast \sigma \prime (z_1^1)\\ & =(\frac{\partial C}{\partial z_1^2}\ast w_{11}^2+\frac{\partial C}{\partial z_2^2}\ast w_{12}^2+\frac{\partial C}{\partial z_3^2}\ast w_{13}^2)\ast \sigma \prime (z_1^1)\\ & =({\delta }_1^2\ast w_{11}^2+{\delta }_2^2\ast w_{12}^2+{\delta }_3^2\ast w_{13}^2)\ast \sigma \prime (z_1^1)\end{array}\text{\hspace{1em}(11)}$$
一般，第$l$层第$j$个单元对于损失函数的梯度可以表示为：
$$\begin{array}{rl}D({\delta }_j^l)& =D((\sum _{k=1}^{N^{l+1}}{w}_{jk}^{l+1}{\delta }_k^{l+1})\ast \sigma \prime (z_j^l))\end{array}$$
由上一节对激活函数的三个假设可知， 加权输入$z^L$的均值为$0$，且在初始化阶段激活函数的导数$\sigma \prime (z_j)$近似为 1:
$$\begin{array}{rl}& =D(\sum _{k=1}^{N^{l+1}}{w}_{jk}^{l+1}{\delta }_k^{l+1})\\ & =\sum _{k=1}^{N^{l+1}}D(w_{jk}^{l+1}{\delta }_k^{l+1})\\ & =\sum _{k=1}^{N^{l+1}}D(w_{jk}^{l+1})D({\delta }_k^{l+1})\\ & =N^{l+1}D(w_j^{l+1})D({\delta }^{l+1})\end{array}\text{\hspace{1em}(12)}$$
考虑
$$\begin{array}{rl}{N}^{l}D({\delta }^l)& =\sum _{j=0}^{N^l}D({\delta }_j^l)\\ & =\sum _{j=0}^{N^l}{N}^{l+1}D(w_j^{l+1})D({\delta }^{l+1})\\ & =N^l{N}^{l+1}D(W^{l+1})D({\delta }^{l+1})\end{array}\text{\hspace{1em}(13)}$$
得到第$l$层加权输入的梯度$D({\delta }^l)$为：
$$\begin{array}{rl}D({\delta }^l)& =N^{l+1}D(W^{l+1})D({\delta }^{l+1})\\ & =N^{l+1}D(W^{l+1})N^{l+2}D(W^{l+2})D({\delta }^{l+2})\\ & =D({\delta }^M)\prod _{i=1}^{M-1}{N}^{i}D(W^i)\end{array}\text{\hspace{1em}(14)}$$
其中${\delta }^M$是最后一层加权输入对于损失函数的梯度。
因此为了确保在反向传播中所有层的加权输入对于损失函数的梯度$D({\delta }^l)$保持不变，需要满足的条件是
$$D(W^l)=\frac{1}{N^l}\text{\hspace{1em}(15)}$$
这里$N^l$代表$l$层的输出端数量。

1.3.前向传播以及方向传播的权衡

通常来说一个网络层的的输入端数量和输出端数量不保持一致，$Glorot$和$Bengio$建议使用输入端和输出端数量的均值来作为每一层的权重方差。即就是:
$$D(W^l)=\frac{2}{N^l+N^{l-1}}\text{\hspace{1em}(16)}$$
得到的权重应该满足的方差， 可以计算该方差对应的均匀分布:
由均匀分布方差公式,$W\sim U(a,b)$对应的方差为$\frac{(b-a{)}^2}{12}$,可得方差对应的均匀分布应为：
$$W\sim U(-\frac{\sqrt{6}}{\sqrt{N^l+N^{l-1}}},\frac{\sqrt{6}}{\sqrt{N^l+N^{l-1}}})\text{\hspace{1em}(17)}$$

2.caffe代码实现

/**
 * @brief Fills a Blob with values @f$ x \sim U(-a, +a) @f$ where @f$ a @f$ is
 *        set inversely proportional to number of incoming nodes, outgoing
 *        nodes, or their average.
 *
 * A Filler based on the paper [Bengio and Glorot 2010]: Understanding
 * the difficulty of training deep feedforward neuralnetworks.
 *
 * It fills the incoming matrix by randomly sampling uniform data from [-scale,
 * scale] where scale = sqrt(3 / n) where n is the fan_in, fan_out, or their
 * average, depending on the variance_norm option. You should make sure the
 * input blob has shape (num, a, b, c) where a * b * c = fan_in and num * b * c
 * = fan_out. Note that this is currently not the case for inner product layers.
 *
 * TODO(dox): make notation in above comment consistent with rest & use LaTeX.
 */
template <typename Dtype>
class XavierFiller : public Filler<Dtype> {
 public:
  explicit XavierFiller(const FillerParameter& param)
      : Filler<Dtype>(param) {}
  virtual void Fill(Blob<Dtype>* blob) {
    /* count参数是权重矩阵的总个数， 权重矩阵维度[输出端个数, 输入端个数] */
    CHECK(blob->count());
    /* fan_in是输入端个数 */
    int fan_in = blob->count() / blob->shape(0);
    // Compatibility with ND blobs
    int fan_out = blob->num_axes() > 1 ?
                  blob->count() / blob->shape(1) : 
                  blob->count(); /* 当输入端个数为1时, count()参数即就是输出端个数 */
    /** 默认情况下, 只考虑输入端个数;
      * 当配置了FillerParameter_VarianceNorm_AVERAGE参数, 考虑输入和输出;
      * 当配置了FillerParameter_VarianceNorm_FAN_OUT参数, 只考虑输出;
      */ 
    Dtype n = fan_in;  // default to fan_in
    if (this->filler_param_.variance_norm() ==
        FillerParameter_VarianceNorm_AVERAGE) {
      n = (fan_in + fan_out) / Dtype(2);
    } else if (this->filler_param_.variance_norm() ==
        FillerParameter_VarianceNorm_FAN_OUT) {
      n = fan_out;
    }
    Dtype scale = sqrt(Dtype(3) / n);
    /* 生成均匀分布 */
    caffe_rng_uniform<Dtype>(blob->count(), -scale, scale,
        blob->mutable_cpu_data());
    CHECK_EQ(this->filler_param_.sparse(), -1)
         << "Sparsity not supported by this Filler.";
  }
};

具体的含义已在代码注释中说明了， 代码路径在/caffe/include/caffe/filler.hpp

参考

链接: Xavier Initialization.