[論文レビュー] APMSqueeze: A Communication Efficient Adam-Preconditioned Momentum SGD Algorithm

Hanlin Tang, Shaoduo Gan+5 more|arXiv (Cornell University)|Aug 26, 2020

Vehicle License Plate Recognition参考文献 35被引用数 4

ひとこと要約

APMSqueeze は、最初の学習段階で Adam を用いてモーメンタム SGD を前処理し、その後エラー補償付き勾配圧縮を適用することで、通信量を最大 97% 削減（32 倍のデータ削減）し、BERT や ResNet-18 で最大 10 倍のエンドツーエンドの高速化を達成する通信効率の高い最適化アルゴリズムであり、収束性や精度を損なわない。

ABSTRACT

Adam is the important optimization algorithm to guarantee efficiency and accuracy for training many important tasks such as BERT and ImageNet. However, Adam is generally not compatible with information (gradient) compression technology. Therefore, the communication usually becomes the bottleneck for parallelizing Adam. In this paper, we propose a communication efficient {\bf A}DAM {\bf p}reconditioned {\bf M}omentum SGD algorithm-- named APMSqueeze-- through an error compensated method compressing gradients. The proposed algorithm achieves a similar convergence efficiency to Adam in term of epochs, but significantly reduces the running time per epoch. In terms of end-to-end performance (including the full-precision pre-condition step), APMSqueeze is able to provide {sometimes by up to $2-10 imes$ speed-up depending on network bandwidth.} We also conduct theoretical analysis on the convergence and efficiency.

研究の動機と目的

• BERT や ResNet-18 などの大規模モデルの分散学習における通信ボトル neck を解消すること。ここで Adam は不可欠であるが、標準的な圧縮手法と互換性がない。
• 勾配圧縮と Adam の非線形更新ルールの不適合性を克服し、既存手法では収束が崩れる問題を解決すること。
• 通信効率の高いアルゴリズムを開発し、Adam の収束速度と最終的なモデル精度を維持しながら、データ送信量を大幅に削減すること。
• 限られたネットワーク帯域幅下でも、最新のスケーラブルかつ高性能な分散学習を実現するため、新規の前処理と圧縮パイプラインを導入すること。

提案手法

• 学習の最初の段階として、数エポックにわたりフル精度の分散 Adam を用いてモーメンタム SGD を前処理し、その後圧縮学習に移行する。
• 前処理段階の後、勾配圧縮をモーメンタム更新方向に適用し、収束性が保たれるようにエラー補償を実装する。
• 量子化とスパース化の両方をサポートするランダム化圧縮演算子 $\bm{C}_{\omega}(\cdot)$ を用い、圧縮損失を補償する。
• 前処理段階では Adam のモーメンタムベクトル $\bm{m}_t$ と適応的学習率 $\bm{v}_t$ を維持するが、本学習段階では $\bm{v}_t$ を固定し、$\bm{m}_t$ を圧縮する。
• ハイブリッド学習戦略を導入：最初の 15% の学習でフル精度の Adam を使用し、残りの 85% でエラー補償付きの圧縮モーメンタム SGD を実行する。
• 理論的分析により、標準的な仮定の下で、圧縮されたアルゴリズムが元の非圧縮バージョンと同等の漸近的収束速度に達することを示した。

実験結果

リサーチクエスチョン

• RQ1 Adam を用いた学習に勾配圧縮を効果的に適用することで、収束性やモデル精度が低下しないか？
• RQ2 数エポックの Adam を用いたモーメンタム SGD の前処理が、その後の学習段階における安定的かつ効率的な通信圧縮を可能にするか？
• RQ3 Adam によって前処理されたモーメンタム SGD の圧縮版に、理論的な収束保証はあるか？
• RQ4 BERT や ResNet-18 といった大規模モデルにおいて、この手法により通信オーバーヘッドの削減とエンドツーエンドの高速化はどの程度達成できるか？
• RQ5 異なるネットワーク帯域幅や圧縮技術に対しても、提案手法は頑健であるか？

主な発見

• APMSqueeze は、BERT-Base および BERT-Large において、フル精度の Adam と同等の収束特性と最終的な精度を維持しながら、通信量を最大 97% 削減（32 倍の圧縮）した。
• 128 個の GPU を用いた環境では、1 エポックあたりの実行時間が最大 8 倍高速化され、1Gbps のネットワーク条件下でエンドツーエンドの高速化が最大 10 倍達成された。
• 10Gbps の帯域幅でもエンドツーエンドの高速化が 2 倍に達し、ネットワーク環境にかかわらず高いスケーラビリティを示した。
• BERT-Base、BERT-Large、CIFAR-10 における ResNet-18 を含む全評価タスクで、Adam と同等の訓練損失とテスト精度を維持した。
• 理論的分析により、APMSqueeze が非圧縮ベースラインと同等の漸近的収束速度に達することが確認され、ワーカー数に比例した線形の高速化が実現した。
• APMSqueeze は、BERT といった非常に複雑なモデルを勾配圧縮と Adam に類似した最適化手法で効果的に学習できる、通信効率の高い分散アルゴリズムとしての初の実現例である。