使用自然梯度的神经网络并行训练【译】

本篇是对论文 Parallel training of DNNs with Natural Gradient and Parameter Averaging前几个部分的翻译。

摘要

我们描述了一种在kaidi语音识别工具中使用的一种面向多核,多GPU机器的神经网络训练框架。为了尽可能与硬件无关,我们需要一种方法来使用多台机器,而不会产生过多的网络流量。这里使用的方法是周期性的平均网络参数(每一或者两分钟),然后将平均之后的参数重新分布到机器上进行进一步的训练。每台机器都只能看见自己的训练数据。本身来说,这种方法并不能很好的工作。但是,我们有另外一种方法,就是正确且有效的实现了针对随机梯度下降的自然梯度,它可以允许我们的周期性平均方法工作的很好,并且极大的提高了在单个机器上的SGD收敛程度。

介绍

神经网络的并行训练一般是数据并行和模型并行的结合。一般的数据并行的方法是每个minibatch交换一下模型参数。这里我们介绍的并行训练框架使用了不同的数据并行方式:我们在不同的机器上有多个执行SGD的进程,每隔几分钟平均一下模型参数并且将他们重新分布到各自的机器上。这对于大规模的语音识别系统训练是非常有效的,但是在我们的例子中,只有和我们实现的有效的自然梯度随机梯度下降结合起来工作的才比较好。考虑到DNN的非收敛性,在这篇文章中我们不解释为什么参数平均可以工作的很好,或者为什么NG-SGD作用很大。这篇文章的要点在于描述我们的方法,并且从经验上说明他们工作的很好。这篇文章的重要性在于,我们展现了当增加GPU数量时,可以收获线性的加速效果,在没有平凡的数据传输的情况下。

在第二部分我们描述了问题的设置,就是DNN在语音识别中的应用-虽然我们的点子比这更加宽泛。在第三部分我们介绍了并行训练方法。在第四部分我们描述了在自然梯度方法背后更加一般的点子,虽然大部分技术细节都被放在了附录部分。在这篇文章中我们不给出任何证明,但是我们在第五部分讨论一些我们认为可以但是证明不了的东西。第六部分有带有/不带有自然梯度和并行的SGD收敛性的实验。我们在第七部分做总结。

NG-SGD有两个版本,一个是简单的版本,还有一个在线版本。技术细节分别在附录A和B。附录C有关于我们DNN实现的背景知识。

问题设置

在训练语音识别的DNN模型时,直接的问题就是将向量$\mathbf{x} \in \mathbf{R}^D$分类为离散的label($y \in \mathcal{Y}$)。维度$D$常常数以百计。$\mathbf{x}$是从声学信号中提取的短时谱特性,$\mathcal{Y}$是CD-phone聚类之后绑定的HMM状态,正常情况下,5000左右是典型的。每一对$(\mathbf{x}, y)$对应于音频数据的一帧。帧长和帧移一般是25ms和100/s,$\mathbf{x}$含有几个相邻上下文的谱信息。我们最终不是只对这个$y$感兴趣,而是所有$y$的log概率$\log p(y|\mathbf{x})$,因为我们要将它作为在维特比搜索算法中的权值,用来生成最大概率的词序列。训练的目标函数是在所有训练集上的帧,给出$\mathbf{x}$时$y$出现的概率:$\sum_{i}\log(y_i|\mathbf{x}_i)$。由于我们要最小化这个目标,我们使用SGD这个方法有点轻微的用词不当,它是梯度上升。监督标记$y$由训练语料的抄本生成的HMM模型进行维特比对齐得到。

带参数平均的SGD

参数平均概览

在我们的训练中参数平均十分简单。我们有$N$个机器(比如$N = 4$),每一个机器独自的在训练数据的随机子集上运行,我们允许它们的参数逐渐的发散。在每个机器训练完一定数量$K$的数据时(典型的$K = 400000$)。我们平均所有进程的参数,之后将结果重新分发(实际中,我们通过在GridEngine,或者在我们使用的任何管理系统中生成新的作业)。重复这个过程直到完成在全部数据上的若干轮训练,比如说10轮。

我们将训练的外部迭代定义为每个作业处理完$K$批数据的时间。那么每一轮外部迭代的次数和训练数据量和$K$有关。

我们发现给并行的SGD流程设置一个有效的学习率是非常有用的,因为学习率$\eta_t$被每个独立的作业使用,需要除以作业数$N$。当我们为了获得线性的加速而增加作业量$N$时,我们需要按比例的增大学习率以保证有效的学习率不变。这个概念是我们进行参数平均时,任何一个SGD作业的参数更新被稀释为$N$倍造成的。

我们将其设置为参数平均而不是将各个作业的参数变化相加的原因是稳定性考虑。试想这样一种情形,在参数空间中某些方向的Hessian矩阵非常大(我们的学习率也足够大),以至于SGD在执行完$K$个样本时,随机梯度下降已经到达了平衡点。假如有四个作业,那么处理完$K$个样本之后将参数变化量相加,得到的参数不再接近于平衡态,而是在相反的方向,相对于起始点三倍远的地方。显然这会导致模型发散。

我们SGD实现的其他方面

在这里我们提供一些关于SGD实现上的其他方面的一些细节的特性,即用来避免模型发散的学习率调度策略和强制最大参数变化量。

还有一些其他的和主题并不直接相关的我们放在附录C中,即基于CPU和GPU的SGD(C.1),数据随机化(C.2),一般模型平均(C.4),混合分量,子空间(C.5),输入数据的正则化(C.6),参数初始化(C.7),序列训练(C.8)以及使用i-vector的在线解码。

学习率调度

(Senior et al., 2013)中指出,在训练DNN的声学模型时,学习率的指数衰减效果很好,我们也独立的发现了这一现象。通常在训练阶段,我们使用指数衰减将学习率调小10倍。除非特殊说明,我们这里提到的实验,学习率初始设置为0.01,到0.001停止。我们提前说明训练的论数,典型的在4到20之间(如果数据更多,轮数就少一些)。

需要说明的是我们这里的实验都没有独立的针对SGD和NG-SGD调整学习率(我们只是使用了以前在NG-SGD上表现比较好的值),我们在过去也做过拓展性的实验,在小数据集上,学习率是独立调整的。最终在所有的情况下,我们发现NG-SGD是有帮助的。现在在大规模数据集上不方便做这些重复了。

最大参数变化

在深度学习上SGD训练阶段,常见的问题是参数会突然变的很大,目标函数也会变成负无穷。这被称为参数发散。常见的方法是减小学习率重新训练,但是这十分的不方便。为了避免这种情况,我们更改了SGD程序,在每个minbatch上强制一个最大的参数变化量。这样的限制在训练的早期比较活跃,尤其在趋向于输出层的层中。我们在附录C.3中做了进一步说明。

针对SGD的自然梯度

这节描述我们对SGD的自然梯度修正,我们通过一个对称正定矩阵,即对Fisher矩阵的逆的估计来缩放梯度。

从技术上讲,自然梯度是指在黎曼参数表面,沿着常规参数空间中的弯曲路径走一步。它极难计算,但是,前面的工作(Yang & Amari, 1998; Roux et al., 2007)已经使用了“自然梯度”来描述像我们这样的方法,即使用一个Fisher矩阵的估计作为学习率矩阵,因此我们跟随着他们,称呼我们的方法为“自然梯度”。

我们可以在SGD中将常量的学习率替换为一个矩阵

在SGD中,学习率常常被假设为一个常量$\eta_t$,伴随着时间减少,更新等式如下:

其中$\mathbf{g}_t$是目标函数在时刻$t$的梯度(从一个训练样本或者minibatch获得)。然而,将这个标量用一个对称正定矩阵替换是可能的,我们可以改写上式为:

其中$\mathbf{E}_t$是学习率的矩阵分量,为了证明方便,我们没有把$\eta_t$合并到$\mathbf{E}_t$中去。$\mathbf{E}_t$随机是可被接受的:如果我们用可以提前获知的正常数限制$\mathbf{E}_t$的特征值的上下限,那么在给出$\boldsymbol{\theta}$时的$\mathbf{g}_t$和$\mathbf{E}_t$是独立采样的,我们可以证明在某些条件被满足的情况下的收敛性,就像使用一个标量的学习率一样(Bottou, 1998, Sec. 4.2.2)。

总的来说,学习率矩阵不是一个当前处理的样本数据的函数,否则它可能会阻止收敛到一个局部最优解。举一个例子,对于特定类型的训练数据,较小的矩阵将通过对该数据进行加权来明确地偏差学习。

Fisher矩阵的逆是一个合适学习率矩阵

在和自然梯度这个点子相关的统计学习理论中,有理由可以说明为什么我们可以将$\mathbf{E}_t$设置为Fisher矩阵的逆。比如(Murata & Amari, 1999)和(Roux et al., 2007)。Fisher矩阵是在我们要学习一个分布时最直接的定义,而不是分类问题,比如我们现在处理的这个。设想$x$,这个我们要对它的分布建模的变量,它可能是离散的,连续的,$f(x;\boldsymbol{\theta})$是给出$\boldsymbol{\theta}$时$x$的可能性。Fisher信息矩阵$\mathcal{I}(\boldsymbol{\theta})$被定义为log概率对参数偏导的二阶矩。

以上偏导在信息论中称为“score”。在某些情况下,Fisher矩阵和Hessian是相同的。很明显为什么Hessian的逆可以称为一个很好的梯度下降方向。这些条件十分严格,包括正确的模型,$\boldsymbol{\theta}$是基于可以描述正确的数据分布的值的。但是即使这些条件不适用,Fisher矩阵在某种意义上也是和Hessian相同的,也就是说,它在参数变换的情况下也随之变换,因此它的逆依旧是学习率矩阵的一个很好的选择。

将Fisher信息矩阵的概念推广到预测问题$p(y;x,\boldsymbol{\theta})$上也是非常容易的。假设我们已经知道了$x$的分布$q(x)$,那么$p(y, x; \boldsymbol{\theta}) = q(x)p(y;x, \boldsymbol{\theta})$。不难看出“score”就等于$\frac{\partial}{\partial \boldsymbol{\theta}} \log f(x;y, \boldsymbol{\theta})$,因为$q(x)$不依赖$\boldsymbol{\theta}$,所以没有关于$q(x)$的表达式了。在计算Fisher矩阵时计算的期望是在$x$和$y$联合分布上的期望。这个结论在(Roux et al., 2007, Section 3)中同样存在。

更一般的来说,对于任意的目标函数,不一定得是log概率或者log似然,我们都可以计算出一个类似于Fisher矩阵的东西,在变量变化时,它的变换方式和Hessian矩阵类似,比如,它的逆也可以是一个学习率矩阵的合理选择。

在实际中我们需要估计Fisher矩阵

对于大规模的问题,比如参数量达到百万的语音识别问题,即使一次Fisher矩阵的求逆也是不切实际的因为它的时间复杂度是$O(n^3)$的。但是处理它的分解形式还是可以的。之前也有文献在这方面。在(Roux et al., 2007)中,Fisher矩阵被分解成多个对角阵,其中每个阵使用一个低秩的矩阵估计。这个对角阵的点子在(Bastian et al., 2011)也被研究了,每一个阵对应一个权值矩阵,我们的方法也是用的类似的思想。在未发表的手稿(Yang & Amari, 1997)中(有些材料说1998年发表了),作者尝试去证明在一些假设下,对于单个隐层的神经网络的Fisher矩阵是一个Kronecker积的形式。虽然我们考虑的网络形式比他们更加通用,Kronecker积同样出现在我们Fisher矩阵的分解中。

需要说明的是,不进行分解也是可以使用NG的,只是每一轮的时间代价会剧增,可以看一下(Pascanu & Bengio, 2013)中的例子,他们使用阶段牛顿法近似Fisher矩阵的逆。

Fisher矩阵的分解

我们的分解形式是这样的:对于一个有$I$个权值矩阵的神经网络,我们把Fisher矩阵划分为$I$对角阵,每一个对应一个权值矩阵。考虑Fisher矩阵第$i$个对角块,它对应的权值矩阵是$\boldsymbol{W}_i$,在这里我们假设没有单独的偏置存在。那么第$i$块Fisher矩阵可以视为两个对称正定矩阵$\boldsymbol{A}_i$,$\boldsymbol{B}_i$的Kronecker积,其中$\boldsymbol{A}_i$的维度是$\boldsymbol{W}_i$的行数,$\boldsymbol{A}_i$的维度是$\boldsymbol{W}_i$的列数。我们进一步将$\boldsymbol{A}_i$和$\boldsymbol{B}_i$分解为一个低秩的对称矩阵和单位矩阵的倍数的和,那么Fisher矩阵$\boldsymbol{F}$可以被写为:

其中$\boldsymbol{A}_i$,$\boldsymbol{B}_i$以 $\lambda\boldsymbol{I} + \boldsymbol{X}\boldsymbol{X}^T$的形式分解。$\boldsymbol{A}_i$,$\boldsymbol{B}_i$在Kronecker积中的顺序取决于矩阵的存储形式,行主序的还是列主序。实际中我们不显式的处理Kronecker积或者向量化的权值矩阵,所以这个选择不重要。不难证明如果Fisher矩阵可以被这么分解,那么它的逆也可以以同样的形式分解。

我们如何估计Fisher矩阵

我们有两种方法估计Fisher矩阵的分解形式

  • 简单方法:我们从一个minibatch里面拿一些其他的数据来估计Fisher矩阵。这样做十分高效,细节在附录A中。
  • 在线方法:我们从之前所有的minibatch中估计Fisher矩阵。时间较远的minibatch使用一个遗忘因子来减少权重。细节在附录B中。

我们通常使用在线方法,因为在GPU上它跑的很快,模型学习的也很快。这可能是因为它对Fisher矩阵的估计是比较干净的。我们描述简单方法是因为它更加好懂,可以帮助我们启发在线方法。

向量操作

虽然我们描述Fisher矩阵是一个Kronecker积的形式,但是在我们没有显式的实现它。

假设我们当前一次训练一批/个数据。那么SGD按照如下方法更新第$i$个权值矩阵:

$\boldsymbol{x}_{it}$是目标函数的偏导,$\boldsymbol{y}_{it}$是权值作用的输入。这个公式最初存在BP算法中。

在我们的自然梯度下降算法中,这被改写成:

其中$\boldsymbol{A}_{it}$和$\boldsymbol{B}_{it}$是Fisher矩阵的因子。很容易表明,这相当于将参数阶乘以由A和B量形成的Fisher矩阵的倒数。

minibatch的操作

如果不是一次训练一批数据,而是minibatch个数据。那么上一节中的向量$\boldsymbol{x}_{it}$和$\boldsymbol{y}_{it}$就相应的变成了矩阵$\boldsymbol{X}_{it}, \boldsymbol{Y}_{it}$,那么更新公式写为:

注意,我们没有像有些作者一样,将梯度除以minibatch的大小,这样可以更加容易的独立的调minibatch的大小和学习率。NG的更新公式为:

其中带横杠的参数表示修正过的$\boldsymbol{X},\boldsymbol{Y}$。在在线版本中可以写为:

在简单方法中,由于$\boldsymbol{A},\boldsymbol{B}$是从minibatch中的其他元素估计的,所以我们不能这么写(每行单独的乘法),附录A中有高效的实现。

从编程的角度说,我们可以这么描述NG的接口:

  • 简单方法:给出一个minibatch$\boldsymbol{X}_{it}$,每一行是minibatch的一个元素,我们通过每个样本来估计Fisher矩阵,并乘以他们的逆,返回修正过的$\bar{\boldsymbol{X}}_{it}$
  • 在线方法:给出一个minibatch$\boldsymbol{X}_{it}$和前一次的Fisher矩阵因子估计,计算$\bar{\boldsymbol{X}}_{it} = \boldsymbol{X}_{it} \boldsymbol{A}_{i(t-1)}^{-1}$,之后更新$\boldsymbol{A}_{it}$

以上接口对于$\boldsymbol{Y}$和$\boldsymbol{B}$执行流程和$\boldsymbol{X}$和$\boldsymbol{A}$是一样的。对于一个minibatch,我们调用接口$2I$次:每个权值矩阵两次。

因子调整

在两个NG方法中,我们想阻止Fisher矩阵极大的影响整体的更新,相比于标准的SGD。对此有几点原因:

  • 在训练的早期,$\boldsymbol{x}$和$\boldsymbol{y}$的量可以非常小,甚至是0,这会导致Fisher矩阵的逆很大甚至是无穷。
  • 常规的收敛防御技术要求学习率矩阵的矩阵分量应具有预先知道的常数的上下限的特征值,但是我们不确定我们使用的Fisher矩阵是否是不变的。
  • 从经验的角度来看,如果使用不调整的Fisher矩阵,我们很难阻止参数发散。

我们的方法是调整$\bar{\boldsymbol{X}}_{it}$和$\bar{\boldsymbol{Y}}_{it}$的量,使得它们的Frobenius范数和输入$\boldsymbol{X}_{it}$和$\boldsymbol{Y}_{it}$一样,我们会在附录中介绍它。

因子调整会引入证明上的微小的问题,那就是每一个样本都会影响它们自己的学习率矩阵了(通过缩放矩阵的缩放因子)。我们之前也提过,虽然使用单个样例的学习率也是允许的,但是在实际问题中这不是一个问题,因为我们基本上不适用小于100的minibatch大小。

使用单位矩阵平滑Fisher矩阵

在两种方法中,我们通过在转置之前给Fisher矩阵加上一个单位矩阵的倍数来平滑它。在简单方法中,这是必要的,因为通常从minibatch中估计的Fisher矩阵都不是满秩的。在在线方法中不是必要的,因为Fisher矩阵的分解已经包含了加上单位矩阵这一操作。但是我们发现通过加上单位矩阵的倍数这一操作,比如对简单方法而言,我们可以调高SGD方法的收敛性。在两种情况下,平滑操作实现如下,如果$\boldsymbol{S} \in \boldsymbol{R}^{D \times D}$是一个从$\boldsymbol{x}$或者$\boldsymbol{y}的协方差中估计的$Fisher矩阵因子,那么我们使用$\boldsymbol{S} + \beta\boldsymbol{I}$来代替Fisher矩阵因子$\boldsymbol{A}$或者$\boldsymbol{B}$,其中

$\epsilon = 10^{-20}$用来防止平滑的$\boldsymbol{S}$值为0。也就是说,我们通过将单位矩阵缩放$\alpha$乘以$\boldsymbol{S}$的对角元素的平均值倍来平滑Fisher矩阵。我们在调参实验中发现,在大部分情况下,$\alpha = 4$对于两种方法,甚至是$\boldsymbol{S}$中存在影响不大的噪声时是合适的,比如minibatch比较大时。我们的理解是,$\alpha$相当大时,我们在$\boldsymbol{x}$和$\boldsymbol{y}$的值协相关性比较大的方向上使用了一个比正常情况小的学习率,在其他方向上使用了一个相对稳定的学习率。