NIPS 2017 | TernGrad: Ternary Gradients to Reduce Communication in Distributed Deep Learning

在深度神经网络的分布式训练中,梯度和参数同步时的网络开销是一个瓶颈。本文提出了一个名为TernGrad的方法,通过将梯度三元化到{-1, 0, 1}上来减少通信量。此外,本文还使用逐层三元化和梯度裁剪技术加速算法的收敛。当把TernGrad应用到AlexNet和GoogLeNet中,测试精度的损失很小。实验表明,TernGrad在不同的深度神经网络模型上拥有显著的加速比。

本文的主要贡献如下:

  1. 将梯度三元化以降低通信量
  2. 通过数学方法证明使用TernGrad后,算法会收敛
  3. 提出逐层三元化和梯度裁剪技术,将梯度的上界逼近标准SGD的上界
  4. 建立性能模型评估TernGrad等梯度压缩方法的加速比

1. Problem Formulation and TernGrad

在介绍TernGrad之前,先来回顾一下传统的数据并行SGD。在数据并行SGD的每次迭代$t$中,训练数据会被分成N份以供N个工作节点进行训练。工作节点i根据输入样本$z_t^{(i)}$计算参数的梯度$\boldsymbol{g}_t^{(i)}$,之后,工作节点将梯度发送给参数服务器。参数服务器接收到所有工作节点的梯度后,对梯度进行聚合,然后把参数返回给工作节点。与传统的数据并行SGD不同的是,TernGrad使用了参数本地化(parameter localization)技术。也就是说,每个工作节点维护一个本地的参数副本,所有工作节点上的参数副本都由同一个随机种子初始化。在整个训练过程中,工作节点和参数服务器之间只传递量化后的梯度。

Figure-1

完整的TernGrad算法由Algorithm 1描述。相比于传统的数据并行SGD,TernGrad在每个工作节点上增加了梯度三元化和参数更新两个操作,参数服务器只需要聚合三元化后的梯度。

Algorithm-1

具体来说,TernGrad按照下式将梯度向量中的每个值三元化到$\{-1,0,+1\}$上,这里$\boldsymbol{b}_t$是一个二元随机向量。

其中$s_t$是一个对$\pm 1$进行缩放的标量。$\circ$是Hadamard乘积,$\text{sign}(\cdot)$和$\text{abs}(\cdot)$分别返回每个元素的符号和绝对值。对于给定的$\boldsymbol{g}_t$,$\boldsymbol{b}_t$的每个元素都独立地服从伯努利分布:

其中$b_{tk}$和$g_{tk}$分别是$\boldsymbol{b}_t$和$\boldsymbol{b}_t$的第$k$个元素。之所以选择这种随机性舍入(stochastic rounding)而不是确定性舍入,是因为随机性舍入具有无偏的期望值且在低精度处理中应用广泛。

理论上,TernGrad至少可以将工作节点传递给参数服务器的通信量减少20.18倍,即使用2bit来编码三元化梯度向量,也能减少16倍的通信量。前面提到,TernGrad使用参数本地化技术来减少参数服务器传递给工作节点的通信量,也就是说参数服务器只对工作节点发送来的三元化梯度进行聚合。这可能导致聚合后的梯度$\overline{\boldsymbol{g}_t}$不再是三元化的,如图2(d)所示,尤其是当工作节点使用不同的缩放因子$s_t^{(i)}$时。为了解决这一问题,我们使用标量共享(scaler sharing)技术,即所有工作节点之间共享一个缩放因子:

通过将参数本地化与标量共享相结合,参数服务器传递给工作节点的通信量至少会降低$32/\log_2(1+2N)$倍。

为了提高TernGrad的精度(即使TernGrad的方差上界接近标准SGD的方差上界),本文提出了逐层三元化(layer-wise ternarizing)与梯度裁剪(gradient clipping)技术。在神经网络的每一层,梯度会随着后向传播而发生改变。因此,TernGrad在每一层使用不同的缩放因子,并且分别对权重和偏置进行三元化。为了进一步提高精度,还可以把梯度分割到不同的桶(bucket)中,再分别进行三元化。但是,这种方法将引入更多的浮点缩放因子并增加通信量。

逐层三元化可以缩小跨层梯度的动态值域所带来的方差间隔。然而,某一层梯度的动态值域仍然是一个问题。因此,TernGrad还使用了梯度裁剪技术,也就是限制梯度$\boldsymbol{g}$中元素$g_i$的范围:

其中$\sigma$是$\boldsymbol{g}$中元素的标准差,$c$是一个需要调整的超参数,本文将该参数设置为2.5。在分布式训练时,每个工作节点先进行梯度裁剪操作,再进行梯度三元化操作。根据图2可以看出,无论是卷积层还是全连接层,所有的梯度都服从高斯分布并且集中在一个很小的范围内。梯度裁剪就是只保留小范围内的梯度,丢弃范围外的梯度。裁剪后,梯度近似服从正态分布,并且方向改变了一个小角度。

Figure-2

上述两种方法之所以有效,是因为当缩放因子$s_t$过大时,大多数梯度被三元化为$0$,只有少数梯度被三元化为$\pm 1$。这就可能使得大多数参数不变而其他参数一直在调整,从而引入较高的训练方差。通过逐层三元化和梯度裁剪,就可以降低$s_t$的值,并且梯度分布近似于正态分布,从而降低了训练时的方差。

2. Experiments

2.1. Intergrating with Various Training Schemes

本文第一组实验主要比较了TernGrad在不同训练模式下的精度。这组实验用到了2个模型,分别是在MNIST数据集上训练的LeNet以及在CIFAR-10数据集上训练的ConvNet,其中ConvNET在训练时使用了数据增强技术。主要用到的优化算法包括传统的SGD,带动量的SGD以及Adam。

Figure-3

图3是LeNet的实验结果。可以看到,使用TernGrad后的传统SGD和带动量的SGD可以在相同的迭代次数内收敛,并且精度不会损失很多。

Table-1

表1是ConvNet的实验结果。实验中,将每个工作节点的batch size设置为固定值。因此,总的batch size会随着工作节点的增加而线性增加。当batch size增大时,TernGrad和baseline的都会略有下降,这是因为较大的batch size会使参数的更新频率降低+,导致模型收敛到尖锐的极小值附近。

2.2. Scaling to Large-scale Deep Learning

本文的第二部分实验主要介绍了将TernGrad在大规模训练时的效果。为了使TernGrad能够成功训练大型神经网络模型,实验中做了以下改动:

  1. 减少了dropout的比率。因为dropout为神经网络添加了随机性(相当于正则化),而TernGrad本身就引入了随机性。过高的dropout比率与TernGrad相结合会使模型效果变差。
  2. 使用较小的权值衰减策略。原因同上。
  3. 不对最后一层进行三元化。因为最后一层的one-hot编码会生成一个倾斜的分布,而三元化后的分布是对称的。

Table-2

实验中所有的网络模型都使用带动量的SGD以及批量归一化(batch normalization)进行训练。AlexNet的结果如表2所示,每个工作节点的batch size固定为128。为了方便实验的进行,所有的深度网络模型都训练相同的epoch。因此,当工作节点增加时,总体的迭代次数就会变少,参数更新频率也会降低。在batch size增大的同时适当增加学习率可以克服这一问题。

Figure-4

TernGrad在不同batch size下都可以收敛到与传统SGD相近的水平。注意到在batch size为1024时,TernGrad甚至将Top-1准确度提升了0.92%,这是因为其固有的随机性使得模型跳出尖锐的极小值区域。图4是batch size为512时的训练情况。图4(a)中,TernGrad的Top-1准确度曲线基本上与baseline重合,这证明了TernGrad的有效性。图4(b)的训练损失曲线说明TernGrad会收敛到比baseline还要低的水平。图4(c)则表明在AlexNet的最后一个全连接层(fc6)中,平均约71.32%的梯度被三元化成0。

Table-3

表3是GoogLeNet的实验结果。平均来说,精确度的损失不超过2%。

3. Performance Model and Analysis

接下来,我们对模型的性能,即吞吐量进行分析。我们主要分析了三个不同的网络模型——AlexNet、GoogLeNet和VggNet-A在使用不同数量的GPU时的吞吐量。

Figure-5

图5是三种网络模型在两个不同的GPU集群上的训练吞吐量。可以看到,TernGrad可以有效地提高训练吞吐量。总的来说,并行加速比取决于网络模型的通信-计算比、GPU的数量和网络带宽。拥有较大通信-计算比的网络模型(如AlexNet和VggNet-A)会更加受益于TernGrad。如图5(a)所示,TernGrad在带宽较小时表现非常好。图5(b)则说明使用高速互连网络(InfiniBand)时,TernGrad仍然可以对训练进行加速。