ICLR 2018 | 深度梯度压缩:减少分布式训练的通信带宽

大型的分布式训练需要较好的通信带宽以便进行梯度的交换，这限制了多节点训练的可扩展性同时也需要昂贵的高带宽设备。这种情况在移动端进行分布式训练时会变得更加糟糕。这篇文章发现分布式SGD中有99.9%的梯度交换都是冗余的，并提出了深度梯度压缩（DGC）用来减少通信带宽。为了保证压缩过程中的精度，DGC使用了4中方法：动量修正、本地梯度剪裁、动量因子遮蔽和warm-up训练。文章中将DGC使用在了图像分类、语音识别和语言模型上面，这些任务中使用的数据集有Cifar10、ImageNet，Penn Treebank 和 Librispeech Corpus。这些情景中在精度未受损失的情况下，DGC的压缩比例达到了270倍至600倍，并将ResNet-50的梯度从97MB剪切至0.35MB，将DeepSpeech的梯度由488MB剪切至0.74MB。深度梯度压缩使得在1Gbps的网络下进行大规模训练成为可能并且促进了分布式训练在移动端的发展。

论文地址:
https://arxiv.org/abs/1712.01887

引言

大规模机器学习提高了模型生产力，同步的SGD被广泛的应用于分布式训练。网络带宽成为了增大训练规模的显著瓶颈，并且带宽问题在移动端进行分布式训练时变得更加糟糕。深度梯度压缩（DGC）通过压缩梯度的方式解决了通信带宽问题，并且为了保证精度不受影响还使用动量修正、本地梯度剪裁、动量因子遮蔽和 warm-up 训练等方法。文章还在不同的任务、模型和数据集上进行了验证包括：图像分类（模型为CNN，数据集为Cifar10 和 ImageNet）、语言模型（模型为RNN，数据集为Penn Treebank）、语音识别（Librispeech Corpus）。这些实验表明梯度可以被压缩600倍而没有造成精度损失，这比前面的一些研究高出了一个数量级。

深度梯度压缩算法

梯度稀疏化

本文通过只发送重要梯度信息来减少通信带宽，梯度的重要程度通过其大小来指定：只有大于某一阈值的梯度才会被传送。为了避免信息损失，剩余的梯度会在本地累积直到累积到足够大然后被传送。于是出现了两种情况：立即发送了大的梯度，延时发送了小梯度。假使 $F(\omega )$为需要优化的损失函数。

上式中$\chi$为训练数据集，$\omega$为网络的权重，$N$是训练的节点数，$B_{k,t}$是minibatch中的一个序列其大小为$b$。在经过$T$轮迭代后有下式。

从上式可以发现本地梯度累积可以被视为batch size由$Nb$增大到$NbT$。

动量修正

当稀疏化非常高时会影响收敛，本文中发现动量修正及本地梯度剪裁会减小这种影响。在N个训练节点上使用vanilla momentum SGD 进行分布式训练时。

上式中$m$为动量，经过T轮迭代后权重系数$\omega$的变化为。

如果SGD的动量直接应用在稀疏梯度上，更新规则变成。

经过动量修正后有。

本地梯度剪裁

梯度剪裁被广泛的采用以防止梯度爆炸。因为需要在每个节点经过迭代来累积梯度，文中在当前的梯度$G_t$与前面的累积梯度（$G_{t-1}$）相加之前进行梯度剪裁。并且将当前节点的梯度乘以系数$N^{-1/2}$进行规范化。假定原始的梯度的L2-norm为$||G|{|}_2$为$th{r}_G$，本地的梯度L2-norm $||G^k|{|}_2$为$th{r}_{G^k}$，有如下公式。

动量因子遮蔽

因为小的梯度被延时更新，所以当这些更新发生时，它们是过时的或者陈旧的。在本文的实验中当梯度稀疏化为99.9%时，大多数的参数是每隔600至1000个迭代来进行更新，这相对于一个epoch中的迭代轮数来说太大了。陈旧的梯度会减慢模型收敛速度并削弱模型的表现。本文引入了动量因子遮蔽这种方法以减轻陈旧梯度的影响。不同于搜索新的动量系数，本文中对积累梯度和动量因子使用了相同的遮蔽。

Warm-up 训练

在早期的训练阶段网络被迅速的更新，并且这些梯度变换多样并且很有效。梯度稀疏化限制了模型的变化范围因此延长了模型显著性改变的周期。Warm-up 训练方法在进行大的minibatch训练很有帮助，在进行warm-up训练时，使用较小的学习率来减慢模型最开始的更新速度并减少有效梯度的稀疏化。在最开始的几个epoch中，本文将梯度稀疏化率由一个较小的值通过指数化增长到最终值。
下图为本文中提出的深度梯度压缩的算法流程。

实验

本文在3种机器学习任务上进行了验证：图像分类任务（Cifar10和ImageNet），语言模型（Penn Treebank），语音识别（AN4和Librispeech）等。唯一的超參就是warm-up训练策略中引入的参数。在所有的与DGC相关的实验中，作者将warm-up周期中的稀疏化率设置为：75%，93.75%，98.4375%，99.6%，99.9%。
1.ResNet-110 在 Cifar10 数据集上训练结果

2.在ImageNet数据集上的梯度压缩率对比

3.在语言模型和语音识别上的实验

4.DGC在各个实验中的加速对比

从上面的图中可以发现，本文中所使用的深度梯度压缩方法在不影响模型精度或者影响很小的情况下对梯度进行了大幅压缩且获得了很大的压缩比率，并有效减小了模型训练所需要的通信带宽。

结论

深度梯度压缩（DGC）在CNNs和RNNs中压缩比率达到了270-600倍。为了达到这种压缩比率并且不降低模型收敛速度，DGC引入了动量修正、本地梯度剪裁、动量因子遮蔽和warm-up训练机制。作者又进一步提出了分层的阈值策略来加速梯度稀疏化处理流程。DGC减少了需求的通信带宽并且提高了在廉价、商用的网络设备上进行分布式训练的可扩展性。

扫码识别关注，获取更多论文解读