学界 | 深度梯度压缩：降低分布式训练的通信带宽

选自arXiv

机器之心编译

参与：Panda

ICLR 2018 将在当地时间明年 4 月 30 日～5 月 3 日在加拿大温哥华举办。目前正处于评审阶段的会议论文也正陆续公开。近日，机器之心发现了这样一篇有望极大改善分布式训练效率的研究论文。以下是我们对该论文的摘要介绍。

大规模分布式训练可以提升越来越深和越来越大的模型的训练效率（Chilimbi et al., 2014; Xing et al., 2015; Moritz et al., 2015; Zinkevich et al., 2010）。分布式训练广泛使用了同步随机梯度下降（SGD）。通过增加训练节点的数量和利用数据并行的优势，可以极大减少在同等规模训练数据上的前向-反向通过的总计算时间。但是，梯度交换（gradient exchange）的计算成本很高，很占用节省下来的计算时间（Li et al., 2014; Wen et al., 2017），尤其是对于计算通信比（computation-to-communication ratio）很低的循环神经网络（RNN）而言。因此，网络带宽就变成了用于扩展分布式训练的显著瓶颈。当在移动设备上执行分布式训练时，比如 federated learning（McMahan et al., 2016; Konecny et al., 2016)），这种带宽问题甚至还会变得更加严重。在移动设备上完成训练是一个引人瞩目的目标，因为这能提供更好的隐私性和更好的个性化（Google, 2017），但移动设备也面临着一些关键的问题：网络带宽甚至还更低、网络连接断断续续以及移动数据套餐价格高昂。

深度梯度压缩（Deep Gradient Compression/DGC）是通过压缩梯度的方式来解决通信带宽问题。为了确保无损于准确度，DGC 在梯度稀疏化（gradient sparsification）之上使用了动量校正（momentum correction）和局部梯度裁剪（local gradient clipping）来维持模型的表现水平。DGC 还使用了动量因子掩蔽（momentum factor masking）和 warmup training 来克服由通信减少所导致的过时问题（staleness problem）。

我们在范围广泛的任务、模型和数据集上对 DGC 进行了实验验证，其中有：用于图像分类的 CNN（使用了 Cifar10 和 ImageNet）、用于语言建模（使用了 Penn Treebank）和语音识别（使用了 Librispeech Corpus）的 RNN。这些实验表明：可以将梯度压缩 600 倍而不造成准确度损失，这比之前的研究成果（Aji & Heafield, 2017）高一个数量级。

图 1：深度梯度压缩可以减少通信时间、提高可扩展性和加速分布式训练

我们通过仅发送重要梯度的方式减少了通信带宽（稀疏更新）。我们使用了梯度幅度作为重要性的简单指示：只有大于一个阈值的梯度才会被传输。为了避免丢失信息，我们在本地累积梯度的其它部分。最后，这些梯度会变得足够大，从而得到传输。因此，我们会立即发送大梯度，但最终会随时间发送完所有的梯度。算法 1 展示了这种方法。

算法 1：在节点 k 上的梯度稀疏化

图 2：带有动量校正的梯度稀疏化

表 1：深度梯度压缩中的技术

我们首先在图像分类任务上检查了深度梯度压缩。图 3(a) 和 3(b) 是 ResNet-110 在 Cifar10 上使用 4 个节点得到的 Top-1 准确度和训练损失。其梯度稀疏度是 99.9%（仅有 0.1% 是非零的）。Gradient Dropping（Aji & Heafield, 2017）（红色）的学习曲线由于梯度过时问题而比基准更差。使用动量校正（黄色），学习曲线的收敛稍微快了一些，准确度也接近基准得多。使用动量因子掩蔽和 warmup training（蓝色），梯度过时问题被消除了，而且学习曲线也紧跟基准。表 2 给出了详细的准确度。在使用深度梯度压缩时，ResNet-110 的准确度得到了完全保持。

图 3：ResNet 在图像分类任务中的学习曲线（梯度稀疏度为 99.9％）

表 2：在 Cifar10 数据集上训练的 ResNet-110

表 3：在 ImageNet 数据集上的梯度压缩比的比较

图 4：在 PTB 数据集上的 LSTM 语言模型的困惑度和训练损失（梯度稀疏度为 99.9％）

表 4：使用 4 个节点的语言建模和语音识别的训练结果

对于语言建模，图 4 给出了在梯度稀疏度为 99.9％ 时，使用 4 个节点训练得到的语言模型的困惑度和训练损失。使用深度梯度压缩所得到的训练损失非常接近基准，验证困惑度也一样。从表 4 可以看到，深度梯度压缩将梯度压缩了 462 倍，而困惑度仅稍微有点损失。

图 5：在 AN4 上的 5 层 LSTM 的 WER（词错率）和训练损失（梯度稀疏度为 99.9％）

对于语音识别，图 5 给出了在梯度稀疏度为 99.9％ 时，使用 4 个节点在 AN4 数据集上训练得到的 5 层 LSTM 的词错率和训练损失。对于深度梯度压缩中的技术所带来的提升，这个学习曲线表现出了和在图像网络中类似的结果。表 4 给出了在 LibriSpeech 测试数据集上的词错率（WER）结果，其中 test-clean 包含干净的语音，而 test-other 是有噪声的语音。使用深度梯度压缩训练的模型在干净语音和有噪声语音上都得到了更好的识别能力，即使梯度的规模被压缩了 608 倍。

图 6：深度梯度压缩可以提升分布式训练的速度和可扩展性。每个训练节点有 4 块英伟达 Titan XP GPU 和一个 PCI 开关。

论文：深度梯度压缩：降低分布式训练的通信带宽（DEEP GRADIENT COMPRESSION: REDUCING THE COMMUNICATION BANDWIDTH FOR DISTRIBUTED TRAINING）

论文链接： https://openreview.net/pdf?id=SkhQHMW0W

摘要：大规模分布式训练需要显著的通信带宽来进行梯度交换，这会限制多节点训练的可扩展性并且需要昂贵的高带宽网络基础设施。当在移动设备上进行分布式训练（federated learning）时，情况甚至更加糟糕，会遭遇更高的延迟、更低的吞吐量和断断续续的糟糕连接的问题。在这篇论文中，我们发现在分布式 SGD 中的梯度交换的 99.9% 都是冗余的，并且提出了可以极大降低通信带宽的深度梯度压缩（DGC）。为了在这种压缩过程中保持准确度，DGC 使用了 4 种方法：动量校正、局部梯度裁剪、动量因子掩蔽和 warmup training。我们已经将深度梯度压缩应用到了图像分类、语音识别和语言建模任务上，使用了 Cifar10、ImageNet、Penn Treebank 和 Librispeech Corpus 等多个数据集。在这些场景中，深度梯度压缩实现了 270 倍到 600 倍的梯度压缩比，并且没有给准确度造成损失；将 ResNet-50 的梯度规模从 97MB 减少到了 0.35MB，对于 DeepSpeech 则将梯度规模从 488MB 压缩到了 0.74MB。深度梯度压缩能够在廉价的商品级 1Gbps 以太网上实现大规模分布式训练，并能促进在移动设备上的分布式训练。

本文为机器之心编译，转载请联系本公众号获得授权。

✄------------------------------------------------

加入机器之心（全职记者/实习生）：[e-mail]

投稿或寻求报道：[e-mail]

广告&商务合作：[e-mail]

点击「阅读原文」，在 PaperWeekly 参与对此论文的讨论。