[論文レビュー] 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の収束が遅くなる原因となる勾配の爆発問題に対処すること。
• 全ステップでグローバルな勾配平均化を必要としない、通信頻度が低い分散環境でも勾配クリッピングを可能にすること。
• 非リプシッツ的・非凸的・緩い滑らかさの損失関数の下で、具体的には線形スループット向上と通信複雑度の低減という理論的保証を確立すること。
• Wikitext-2 や Penn Treebank、ImageNet といったベンチマークデータセットを用いた実験的評価を通じて理論的主張の妥当性を検証すること。

提案手法

• 各マシンで複数回の局所的SGDステップを実行し、勾配クリッピングを適用した後、他のマシンと通信する分散アルゴリズムを提案する。
• グローバルに平均化された勾配ではなく、各マシンの局所的勾配ノルムに基づいて勾配クリッピングを適用することで、通信効率を維持する。
• 標準的な滑らかさを一般化した、$(L_0, L_1)$-滑らかさと呼ばれる緩い滑らかさ条件を採用し、LSTMの挙動をより適切に捉える。
• 非リプシッツ的・非凸的状況下でも収束保証を得るために、切り捨てられた確率的変数の推定に新規の解析手法を用いる。
• モデル重みの集約を $I$ 回の局所更新ステップ後にのみ行う通信戦略を導入し、通信頻度を低減する。
• 理論的解析により、$epsilon$-停留点を求めるための $O(1/N\epsilon^4)$ の反復複雑度と $O(1/\epsilon^3)$ の通信複雑度を証明する。

実験結果

リサーチクエスチョン

• RQ1 通信頻度が低い分散学習に、勾配クリッピングを効果的にスケーリング可能であり、収束保証を維持できるか？
• RQ2 緩い滑らかさの下で、局所的勾配クリッピングに周期的なグローバル集約を組み合わせることで、分散学習における線形スループット向上が達成できるか？
• RQ3 非凸的・非リプシッツ的状況下で、単純な並列勾配クリッピングと比較して通信複雑度を低減できるか？
• 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 よりも顕著に優れた性能を示した。