ASGD with Delay Compensation

简介

Asynchronous Stochastic Gradient Descent (ASGD)异步的随机梯度下降在深度学习模型的训练中经常被用到,但是会存在delayed gradients的问题,就是当一个worker向参数server端提交它算出的梯度时,server端其实已经被其它worker更新好多次了。因此该工作提出了梯度补偿的概念,主要方法是利用梯度函数的泰勒展开去有效逼近loss函数的Hessian矩阵。通过在cifar和imagenet数据集上验证,实验结果显示,新的方法DC-ASGD性能优于同步SGD和异步SGD,几乎接近序列SGD的性能。

ASGD介绍

传统的SGD,更新公式为:

其中,wt为当前模型,(xt, yt)为随机抽取的数据,g(wt; xt, yt)为(xt, yt)所对应的经验损失函数关于当前模型wt的梯度,η为步长/学习率。

同步随机梯度下降法(Synchronous SGD)在优化的每轮迭代中,会等待所有的计算节点完成梯度计算,然后将每个工作节点上计算的随机梯度进行汇总、平均并上面的公式更新模型。之后,工作节点接收更新之后的模型,并进入下一轮迭代。由于Sync SGD要等待所有的计算节点完成梯度计算,因此好比木桶效应,Sync SGD的计算速度会被运算效率最低的工作节点所拖累。

异步随机梯度下降法(Asynchronous SGD)在每轮迭代中,每个工作节点在计算出随机梯度后直接更新到模型上,不再等待所有的计算节点完成梯度计算。因此,异步随机梯度下降法的迭代速度较快,也被广泛应用到深度神经网络的训练中。然而,Async SGD虽然快,但是用以更新模型的梯度是有延迟的,会对算法的精度带来影响。如下图:

在Async SGD运行过程中,某个工作节点Worker(m)在第t次迭代开始时获取到模型的最新参数 w_{t} 和数据(xt, yt),计算出相应的随机梯度 g_{t} ,并将其返回并更新到全局模型w上。由于计算梯度需要一定的时间,当这个工作节点传回随机梯度g_{t}时,模型w_{t}已经被其他工作节点更新了τ轮,变为了 w_{t+\tau} 。也就是说,Async SGD的更新公式为:

可以看到,对参数w_{t+\tau}更新时所使用的随机梯度是g(wt),相比SGD中应该使用的随机梯度g(wt+τ)产生了τ步的延迟。因而,我们称Async SGD中随机梯度为“延迟梯度”。

延迟梯度所带来的最大问题是,由于每次用以更新模型的梯度并非是正确的梯度,因为g(wt) ≠ g(wt+τ),所以导致Async SGD会损伤模型的准确率,并且这种现象随着机器数量的增加会越来越严重。

因此DC-ASGD算法设计了一种可以补偿梯度延迟的方法,他们首先研究了正确梯度g(wt+τ)和延迟梯度g(wt)之间的关系,我们将g(wt+τ)在wt处进行泰勒展开得到:

其中,∇g(wt)为梯度的梯度(loss fuction的Hessian矩阵,因此梯度g(wt)是loss函数关于参数wt的导数)。H(g(wt))为梯度的Hessian矩阵。那么如果将所有的高阶项都计算出来,就可以修正延迟梯度为准确梯度了。然而,由于余项拥有无穷项,并且计算量十分复杂,所以无法被准确计算。因此,可用上述公式中的一阶项进行延迟补偿:

但是上面的公式还是要计算∇g(wt)(参数的Hessian矩阵),但是在DNN中有上百万甚至更多的参数,计算和存储Hessian矩阵∇g(wt)很困难。因此,寻找Hessian矩阵的一个良好近似是能否补偿梯度延迟的关键。根据费舍尔信息矩阵的定义,梯度的外积矩阵是Hessian矩阵的一个渐近无偏估计:

其实,进一步可以写成:G\left(W_{t} \right)=g\left( W_{t}\right)g\left( W_{t} \right)^{T}

又可知,在DNN中用Hessian矩阵的对角元素来近似表示Hessian矩阵,可在显著降低运算和存储复杂度的同时还可以保持算法精度,于是我们采用外积矩阵的diag(G(wt))作为Hessian矩阵的近似。为了进一步降低近似的方差,我们使用一个(0,1]之间参数λ来对偏差和方差进行调节。另外由于:

综上,带有延迟补偿的异步随机梯度下降法(DC-ASGD):

具体算法

算法1中,worker m从参数服务器中pull最新的模型参数w,然后计算得到梯度 g_{m} 后push到参数服务器中。

算法2中,当参数服务器接收到worker m算出来的梯度 g_{m}后,利用梯度补偿公式算出下一个时间刻参数服务器正确的参数。如果参数服务器接受到worker m的pull 参数请求时,将当前参数服务器的参数wt备份成w_bak,并将wt发送给worker m。

实验

在CIFAR10数据集和ImageNet数据集上对DC-ASGD算法进行了评估,实验结果显示:DC-ASGD算法与Async SGD算法相比,在相同的时间内获得的模型准确率有显著的提升,并且也高于Sync SGD,基本可以达到SGD相同的模型准确率。

参考文献:

编辑于 2019-09-05

文章被以下专栏收录