EMNLP 2017 | Sparse Communication for Distributed Gradient Descent

通过将分布式随机梯度下降（SGD）中的稠密更新替换成稀疏更新可以显著提高训练速度。当大多数更新接近于0时，梯度更新会出现正偏差，因此我们将99％最小更新（绝对值）映射为零，然后使用该稀疏矩阵替换原来的稠密矩阵。该方法可以于梯度量化相结合来进一步压缩梯度。我们探索了不同的的参数配置并将它们应用到神经机器翻译和MNIST图像分类任务中。大多数参数配置适用于MNIST，而不同的配置会降低更复杂的翻译任务的收敛速度。我们的实验说明在不损失精度和BLEU的情况下，可以在MNIST任务上获得49%的加速比，在神经网络翻译（Neural Machine Translation，NMT）任务上获得22%的加速比。

1. Introduction

分布式计算可以在大数据集上高效地训练大规模神经网络模型。本文主要关注数据并行，即每个节点在数据集的不同子集上优化相同的模型，通过网络进行梯度和参数的传递。这种网络通信开销是巨大的，之前有2中方法压缩数据以降低通信开销：1比特量化和通过丢弃小的更新来传输稀疏矩阵。上述两种方法都在语音识别任务和MNIST数据集上测试过。通过把这些方法应用到NMT模型中，我们发现1比特量化在NMT模型中表现不佳。

本文主要比较了NMT和MNIST图像分类在训练时的行为。NMT参数由3个大型嵌入矩阵控制：源语言输入，目标语言输入和目标语言输出。相比于MNIST的3层神经网络，NMT系统包括多种不同尺度的参数。

2. Related Work

Zinkevich等人提出了一种异步架构，每个节点独立地推送或拉取模型而不用等待较慢的节点。Chilimbi和Rechte建议在无锁情况下更新模型，允许竞争条件。此外，Dean在交换梯度之前运行多个小批量迭代，从而降低了通信成本。1比特SGD和之后的量化SGD通过将梯度更新转化为1比特矩阵来降低通信量。Storm提出了阈值量化，即只发送大于某一阈值的梯度。然而最优阈值不好选择，不过阈值也可以随着优化进度变化。

3. Distributed SGD

如图1所示，我们使用带参数共享的分布式SGD。其中N个工作节点即是客户机又是服务器，每个服务器负责\(\frac{1}{N}\)的参数。客户机拥有所有参数的副本，这些参数副本用来计算梯度。计算出的梯度被分成N份，每一份都被推送到对应的服务器中。同样的，每台客户机会从所有的服务器中拉取参数。每个节点与所有\(N\)个节点进行关于\(\frac{1}{N}\)份参数的通信，因此每个节点需要的带宽是固定的。

4. Sparse Gradient Exchange

我们通过丢弃绝对值最小的梯度的\(R\%\)来实现梯度稀疏化。随着时间的推移，小梯度会累积，我们发现将它们归零会损害收敛速度。我们在局部记录残差并将它们与下一次迭代的梯度相加，然后再丢弃。

整个梯度丢弃的过程如算法1所示。算法输入是梯度\(\triangledown\)和丢弃率\(R\)。首先将梯度与上一轮迭代的残差相加，得到新的梯度。然后再选择一个阈值，使梯度绝对值小于该阈值的数目占\(R\%\)。然后初始化\(dropped\)数组为0，将满足\(|\triangledown[i]| \ge threshold\)条件的\(\triangledown[i]\)赋值给\(dropped[i]\)

，再根据\(\triangledown - dropped\)计算残差。

为了计算残差，我们在服务器端保存最后一次拉取的副本。同步SGD保留一份副本，而异步SGD在每台客户机上都有一份副本，但是服务器只对每台客户机负责\(\frac{1}{N}\)的参数，因此内存开销也是固定的。

选择阈值的开销比较昂贵，但我们可以获得一个阈值的近似值。我们对0.1％的梯度进行采样，并通过样本集来选择阈值。我们可以在局部为每个参数矩阵选择一个阈值，或者为所有参数全局选择一个阈值。在实验中，我们发现层标准化使全局阈值起作用，因此默认情况下我们使用具有一个全局阈值的层标准化。如果没有层标准化，全局阈值会降低NMT的收敛性。

5. Experiment

我们在MNIST数据集和罗马语到英语的神经机器翻译系统上做了实验。对于图像分类的实验，我们训练了一个3层的全连接神经网络，每层拥有4096个隐藏单元。我们使用初始学习率为0.005的AdaGrad算法训练模型，batchsize设为40。本次实验的NMT模型基于带注意力机制的编码——解码LSTM，该模型拥有119M参数。默认batchsize为80，每迭代10000步保存和验证模型。我们选择了4个具有比较高的BLEU的模型，然后对它们做平均，得到最终的模型。AmuNMT模型涌来解码，其中beamsize为12。所有的实验都使用异步SGD算法，尽管我们的方法对同步SGD也很管用。

5.1. Drop Raito

为了找到最有的丢弃率R%，我们尝试了90%，99%和99.9%。图3显示，在丢弃掉99.9%的梯度后，尽管BLEU分数很低，但模型仍在学习。然而，尽管使用偏移值编码将数据减少了50倍，但丢弃99％的梯度对收敛或BLEU几乎没有影响。图3的两个子图中的x轴都是batchsize，这表明我们不依赖于速度改进来补偿收敛速度。对于MNIST和NMT来说，梯度丢弃而非正则化提高了早期批次的精度。

5.2. Local vs Global Threshold

因为所有的参数之间可能不能直接比较，因此我们对每个矩阵使用局部阈值对所有梯度使用全局阈值。我们没有实现局部阈值拉取，因此仅示出了局部阈值推送。

结果显示，层标准化对MNIST图像分类基本上没啥影响。另一方面，不使用层归一化时，NMT模型收敛的比较慢。

5.3. Convergence Rate

对于NMT模型，通信时间只占总训练时间的13%。丢弃99%的梯度可以提升11%的训练速度。在batchsize为32时，通信时间占总时间的17%，而丢弃99%的梯度可以获得22%的速度提升。对于MNIST，通信时间占总时间的41%，而我们会有49%的速度提升。

我们还对收敛率进行了实验，它是训练损失和速度的结合体。如图6所示，MNIST的baseline最终达到了99.28%的准确率，而丢弃99%的梯度后我们获得了99.42%的准确率。

表2是NMT的实验结果，最终的BLEU分数基本上没什么变化。在batchsize为32时，我们算法的收敛速度比baseline快23%，当batchsize为80时，两个算法的训练时间基本相同。

5.4. 1-Bit Quantization

使用梯度丢弃后再用1比特量化技术可以进一步地压缩梯度。我们使用3种不同的配置测试了1比特量化：阈值，列式均值和全局均值。在梯度丢弃（99%的丢弃率，使用层标准化以及全局阈值）之后应用1比特量化技术。

在图8中，1比特量化技术降低了NMT的收敛速率，但是对MNIST图像识别任务没有太大影响。因为NMT任务中前1%的梯度分布更为偏斜，所以1比特量化造成了更多的损失，2比特量化就足够了。

6. Conclusion and Future Work

梯度更新具有正偏斜性：大多数梯度都接近于0。我们可以利用其中99%的梯度进行局部更新，来减少50倍的通信量。先前的工作表明可以应用1比特量化来进一步压缩通信。然而，我们发现这在NMT模型上是不成立的，使用2比特量化才足够。此外，我们的NMT模型包含许多不同尺度的参数，因此层标准化和局部阈值都是必要的。未来的工作是在具有昂贵通信成本的系统（如多节点环境）上测试我们提出方法。