[论文解读] A Communication-Efficient Distributed Gradient Clipping Algorithm for Training Deep Neural Networks

Mingrui Liu, Zhenxun Zhuang+2 more|arXiv (Cornell University)|May 10, 2022

Stochastic Gradient Optimization Techniques被引用 6

一句话总结

该论文提出了一种用于训练深度神经网络的通信高效分布式梯度裁剪算法，尤其适用于RNN和LSTM等易出现梯度爆炸问题的模型。通过在每台机器上执行本地梯度裁剪，并仅周期性地聚合模型权重，该方法实现了$O(1/N\epsilon^4)$的迭代复杂度和$O(1/\epsilon^3)$的通信复杂度，在弱平滑性假设下实现了线性加速并减少了通信轮次。

ABSTRACT

In distributed training of deep neural networks, people usually run Stochastic Gradient Descent (SGD) or its variants on each machine and communicate with other machines periodically. However, SGD might converge slowly in training some deep neural networks (e.g., RNN, LSTM) because of the exploding gradient issue. Gradient clipping is usually employed to address this issue in the single machine setting, but exploring this technique in the distributed setting is still in its infancy: it remains mysterious whether the gradient clipping scheme can take advantage of multiple machines to enjoy parallel speedup. The main technical difficulty lies in dealing with nonconvex loss function, non-Lipschitz continuous gradient, and skipping communication rounds simultaneously. In this paper, we explore a relaxed-smoothness assumption of the loss landscape which LSTM was shown to satisfy in previous works, and design a communication-efficient gradient clipping algorithm. This algorithm can be run on multiple machines, where each machine employs a gradient clipping scheme and communicate with other machines after multiple steps of gradient-based updates. Our algorithm is proved to have $O\left(\frac{1}{Nε^4} ight)$ iteration complexity and $O(\frac{1}{ε^3})$ communication complexity for finding an $ε$-stationary point in the homogeneous data setting, where $N$ is the number of machines. This indicates that our algorithm enjoys linear speedup and reduced communication rounds. Our proof relies on novel analysis techniques of estimating truncated random variables, which we believe are of independent interest. Our experiments on several benchmark datasets and various scenarios demonstrate that our algorithm indeed exhibits fast convergence speed in practice and thus validates our theory.

研究动机与目标

• 为解决RNN和LSTM等深度网络在训练过程中因梯度爆炸问题导致的分布式SGD收敛缓慢的问题。
• 在通信频率较低的分布式设置中实现梯度裁剪，避免每一步都需要全局梯度平均。
• 在非Lipschitz、非凸且弱平滑的损失曲面下，建立理论保证——特别是线性加速和通信复杂度降低。
• 通过在Wikitext-2、Penn Treebank和ImageNet等基准数据集上的实证评估，验证理论结论。

提出的方法

• 提出一种分布式算法，在每台机器上执行多次本地SGD步骤并进行梯度裁剪，之后才与其他机器通信。
• 基于每台机器的本地梯度范数进行梯度裁剪，而非全局平均梯度，以保持通信效率。
• 采用一种弱平滑性条件（即$(L_0, L_1)$-平滑性），该条件推广了标准平滑性，能更好地刻画LSTM的行为。
• 采用新颖的分析技术来估计截断的随机变量，从而在非Lipschitz和非凸设置下实现收敛性保证。
• 引入一种通信策略，仅在$I$轮本地更新后聚合模型权重，从而降低通信频率。
• 理论分析证明，找到$\epsilon$-驻点的迭代复杂度为$O(1/N\epsilon^4)$，通信复杂度为$O(1/\epsilon^3)$。

实验结果

研究问题

• RQ1 在通信频率较低的情况下，梯度裁剪能否在保持收敛性保证的前提下有效扩展到分布式训练？
• RQ2 在弱平滑性条件下，采用周期性全局聚合的本地梯度裁剪是否能在分布式训练中实现线性加速？
• RQ3 在非凸、非Lipschitz设置下，与朴素并行梯度裁剪相比，该算法能否实现更低的通信复杂度？
• RQ4 在异构、部分参与和同质数据设置下，该算法表现如何？
• RQ5 在真实世界基准如AWD-LSTM和ResNet-50上，该算法是否能保持快速收敛和良好的泛化性能？

主要发现

• 所提算法实现了$O(1/N\epsilon^4)$的迭代复杂度和$O(1/\epsilon^3)$的通信复杂度，表明实现了线性加速并减少了通信轮次。
• 在Wikitext-2数据集上，当$I=32$时，该算法实现了76.72的验证困惑度，接近单GPU基线的75.63，同时使用了分布式训练。
• 在Penn Treebank上，当$I=1024$时，该算法实现了66.5976的验证困惑度，在同质设置下优于朴素并行基线的66.0435。
• 在部分参与设置下，当$I=16$时，该算法保持了优异性能，验证困惑度为67.3863，而朴素基线为65.5083。
• 在异构设置下，当$I=32$时，该算法实现了70.4912的验证困惑度，优于朴素并行基线的68.5727。
• 结合动量后，该算法在$I=4$时实现了64.1062的验证困惑度，显著优于朴素并行基线的68.2008。