Deep Gradient Compression: Reducing the Communication Bandwidth for Distributed Training #1

mitsuhiko-nozawa · 2021-05-07T08:26:00Z

一言でいうと

分散環境下における深層学習モデルの学習において、勾配の同期のための通信量を削減(gradient sparsification)＋4つのテクニックを導入することで、モデルの精度を悪化することなく高いモデルの圧縮率、学習速度の向上を達成

論文リンク

著者/所属機関

Yujun Lin, Song Han, Huizi Mao, Yu Wang, William J. Dally
Tsinghua University Beijing National Research Center for Information Science and Technology
Stanford University
Stanford University NVIDIA
Stanford University Google Brain

投稿日付(yyyy/MM/dd)

ICLR2018

概要

同期方式の分散深層学習において、パラメータ更新時の勾配の同期(All Reduce処理)にかかる時間が全体の計算時間に対して支配的になり、ネットワーク帯域幅がボトルネックとなる。高い帯域幅のネットワークは高価であり、モバイルデバイスなどではさらに状況が悪化する。この問題を解決するために、勾配の圧縮を行い、勾配の同期、ネットワーク帯域幅というボトルネックを解決することを目指す。

新規性・差分

従来の勾配の圧縮方法と比べて、高い勾配の圧縮率(実験では600倍)を達成
提案手法は、モデルに対して何か特別なレイヤーを取り付ける必要などがない
比較手法よりも精度の悪化を防ぐことができている

手法

Gradient Sparsification

1: 各GPUノードkの勾配Gを0で初期化 2: ミニバッチごとのイテレーション 3: 前イテレーションの勾配を引き継ぐ 4~7: ミニバッチで勾配を計算 8~13: Mこのレイヤーのうちjごとに繰り返し、レイヤーjのうち勾配の絶対値が上位s%の値を閾値として選択、閾値以上の部分を1,それ以外を0とするマスクを生成、勾配のうちマスクされた部分を抽出,抽出された部分は0にする 14: 各ノードでマスクされた勾配を加算集約、各ノードに配る 15: 重み更新

すなわち、勾配をどんどん蓄積し、大きくなったら更新対象にしていくことを繰り返す。

4つのテクニック

momentum correction

モメンタム付きSGDにおいて、合計Tイテレーションののちに勾配が閾値を超えて更新される状況を考える。Tが増える=ミニバッチサイズが大きくなることを意味しており、この状況で勾配を更新すると、momentum係数がうまく働かなくなる。(本当はバッチサイズBごとに係数をかけたいが、BTごとに係数がかかるようになる)そのため、蓄積する勾配に対して、以下のように係数をかけながら保存していく。

local gradient clipping

勾配クリッピングは、勾配爆発を抑えるために広く用いられている手法だが、分散環境下でこの手法をそのまま取り入れるとなると、勾配の全集約後(グローバル)に行うことになる。そのままの応用はできないので、代わりに各ノードで勾配クリッピングを適応、その際クリップする閾値は ${N}^{-1/2}$ がよかったそう。

momentum masking

勾配の更新が行われたらモメンタム用の変数も0にリセットすることで、古い勾配を引きずりすぎることを防ぐ。

warm-up training

学習初期は勾配が不安定で大きく変化しがちだが、スパース化された勾配では変化がほとんど行われない勾配も多く、その悪影響をさらに受けやすい。また、そのような勾配が蓄積されることで、最適化される勾配の方向を謝る可能性がある。最初の数エポックに warm-up training を設け、この期間は学習率を低めにし、勾配のスパース率も控えめにする。

導入したテクニックがどの部分に寄与しているかの対応表