[論文レビュー] Communication-efficient distributed SGD with Sketching

Nikita Ivkin, Daniel Rothchild+4 more|arXiv (Cornell University)|Mar 12, 2019

Stochastic Gradient Optimization Techniques参考文献 37被引用数 16

ひとこと要約

本稿では、全勾配の代わりにサイズ 𝒪(log d) の勾配スケッチを送信することで通信効率を高めた分散確率的勾配降下法であるSketched-SGDを提案する。これにより通信量の複雑度が非線形的になる。Transformer、LSTM、残差ネットワークにおいて、モデルの精度を落とさずに通信コストを最大40倍まで削減でき、256ワーカーまで効果的にスケーリング可能である。

ABSTRACT

Large-scale distributed training of neural networks is often limited by network bandwidth, wherein the communication time overwhelms the local computation time. Motivated by the success of sketching methods in sub-linear/streaming algorithms, we introduce Sketched SGD, an algorithm for carrying out distributed SGD by communicating sketches instead of full gradients. We show that Sketched SGD has favorable convergence rates on several classes of functions. When considering all communication -- both of gradients and of updated model weights -- Sketched SGD reduces the amount of communication required compared to other gradient compression methods from $\mathcal{O}(d)$ or $\mathcal{O}(W)$ to $\mathcal{O}(\log d)$, where $d$ is the number of model parameters and $W$ is the number of workers participating in training. We run experiments on a transformer model, an LSTM, and a residual network, demonstrating up to a 40x reduction in total communication cost with no loss in final model performance. We also show experimentally that Sketched SGD scales to at least 256 workers without increasing communication cost or degrading model performance.

研究の動機と目的

• 大規模分散ディープラーニングにおける通信ボトル neck を解決すること。ネットワーク帯域幅が学習速度を制限する。
• 1ワーカーあたりの通信コストを 𝒪(d) および 𝒪(W) 未満に削減し、𝒪(log d) にまで低く抑え、多数のワーカーでの効率的スケーリングを可能にする。
• 高い圧縮率を達成しながらも収束保証とモデル性能を維持する手法を開発すること。
• ワーカー数（W）の増加に伴い性能が著しく低下する既存の勾配圧縮技術の限界を克服すること。
• ストリーミングアルゴリズムからのスケッチ化を用いて、通信オーバーヘッドを最小限に抑えた実用的でスケーラブルな分散学習を可能にすること。

提案手法

• Count Sketchとスパース回復技術を用いて、1ワーカーあたり 𝒪(log d) のサイズのスケッチに勾配を圧縮する。
• 全勾配の代わりにローカル勾配のスケッチのみを送信することで、通信量を 𝒪(d) から 𝒪(log d) に削減する。
• スケッチ回復アルゴリズムを用いて、パラメータサーバーで近似された全勾配を、誤差が有界であるように再構築する。
• 最適化に有効な情報がスケッチに保持されるようにすることで、収束特性を維持する。
• 標準的な同期データ並列SGDにスケッチング機構を統合し、標準的な学習ワークフローを保つ。
• 不偏およびバイアスありの勾配推定をサポートし、適切な仮定の下で理論的収束保証を提供する。

実験結果

リサーチクエスチョン

• RQ1 スケッチ化により、収束性を保ちながら通信コストを 𝒪(d) 未満に低下させることができるか？
• RQ2 提案手法は、通信量の増加や性能低下が生じないよう、ワーカー数（W）が非常に多くても効率的にスケーリングできるか？
• RQ3 実世界のモデルにおいて、スケッチ化が高圧縮比（例：40倍）を達成しながらも、最終的なモデル精度を劣化させないか？
• RQ4 従来の勾配圧縮手法（例：top-k、量子化）と比較して、通信効率およびスケーラビリティにおいてどのように優れているか？
• RQ5 スケッチ化に起因する勾配近似の下で、Sketched-SGDの理論的収束挙動はいかなるものか？

主な発見

• Sketched-SGDは、Transformer、LSTM、残差ネットワークモデルにおいて、最終的なモデル性能に影響を及げず、通信コストを最大40倍まで削減した。
• 本手法は256ワーカーまで効果的にスケーリングでき、通信コストの増加やモデル精度の低下がない。
• 1ワーカーあたりの通信コストは 𝒪(log d) にまで低下し、モデルサイズ d に対して非線形的で、ワーカー数 W に対しては定数である。
• MNIST における実験では、Sketched-SGDが通常のSGDと同等の訓練誤差およびテスト誤差率を達成しており、理論的収束レートと整合的である。
• 小規模な実験では、スケッチサイズ280（40列、7行）が、k=10 および P=10 の条件下で約4倍の圧縮比を達成した。
• 理論的分析により、滑らかで凸関数に対する標準的な仮定の下で、Sketched-SGDが良好な収束レートを維持することが確認された。