[論文レビュー] A Closer Look at Codistillation for Distributed Training

Shagun Sodhani, Olivier Delalleau+4 more|arXiv (Cornell University)|Oct 6, 2020

Topic Modeling参考文献 20被引用数 4

ひとこと要約

本稿では、標準的なデータ並列学習の通信効率の高い代替手法として、codistillationを提案している。複数のモデルを並列に学習させ、パラメータの平均化ではなく知識蒸留損失を用いて同期させる。2ウェイのcodistillationは、2倍のデバイスを用いて標準的なデータ並列学習と同等の精度を達成するが、通信量を最大1000倍まで削減できる。一方、n>2のモデルへのスケーリングは、モデルおよびデータセットの特性に起因する課題を伴う。

ABSTRACT

Codistillation has been proposed as a mechanism to share knowledge among concurrently trained models by encouraging them to represent the same function through an auxiliary loss. This contrasts with the more commonly used fully-synchronous data-parallel stochastic gradient descent methods, where different model replicas average their gradients (or parameters) at every iteration and thus maintain identical parameters. We investigate codistillation in a distributed training setup, complementing previous work which focused on extremely large batch sizes. Surprisingly, we find that even at moderate batch sizes, models trained with codistillation can perform as well as models trained with synchronous data-parallel methods, despite using a much weaker synchronization mechanism. These findings hold across a range of batch sizes and learning rate schedules, as well as different kinds of models and datasets. Obtaining this level of accuracy, however, requires properly accounting for the regularization effect of codistillation, which we highlight through several empirical observations. Overall, this work contributes to a better understanding of codistillation and how to best take advantage of it in a distributed computing environment.

研究の動機と目的

• Codスティレーションがスケーラブルで通信効率の高い標準的なデータ並列学習の代替手段として機能できるかどうかを調査すること。
• 従来のall-reduceベースのデータ並列学習と比較して、2ウェイのcodistillationの性能および通信効率を評価すること。
• 分散学習における2つ以上のモデル（n>2）へのcodistillationのスケーリングにおける制限要因と課題を理解すること。
• codistillationが学習ダイナミクスに与える影響、特に正則化に似た効果を分析すること。
• nウェイのcodistillationが2ウェイのcodistillationや標準的な学習を上回る条件を特定すること。

提案手法

• Codistillationは、複数のモデルを並列に学習させ、各モデルがデータの部分集合を処理する。学習損失には、モデル間の予測の不一致をペナルティとして課す知識蒸留損失が含まれる。
• すべてのモデルパラメータをall-reduceで同期する代わりに、知識蒸留による関数的整合性を介してモデルを同期させる。
• 標準的な交差エントロピー損失に加え、モデル出力が一致するよう促す対称的な蒸留損失を組み合わせた学習目的を採用する。
• 完全な勾配やパラメータではなく、蒸留された予測（ソフトラベル）のみを送信することで通信オーバーヘッドを削減する。
• ResNet50（ImageNet）およびトランスフォーマーに基づく機械翻訳モデルを用いて評価し、nウェイのcodistillationに関するアブレーションスタディも実施した。
• ハイパーパrameter α は蒸留損失の重みを制御し、最適なパフォーマンスを得るために調整された。

実験結果

リサーチクエスチョン

• RQ1 2ウェイのcodistillationは、通信量を桁違いに削減しながら、標準的なデータ並列学習と同等の精度を達成できるか？
• RQ2 1モデルあたりのデバイス数を増やした場合、codistillationはどのようにスケーリングするか？また、線形的なスループット向上を維持するか？
• RQ3 2つ以上のモデル（n>2）へのcodistillationのスケーリングが、さまざまなワークロードで一貫して性能向上をもたらさないのはなぜか？
• RQ4 モデルアーキテクチャおよびデータセットの特性が、nウェイのcodistillationの有効性にどのように影響するか？
• RQ5 codistillationは学習ダイナミクスにどのような影響を与え、どのような正則化効果をもたらすか？

主な発見

• 2ウェイのcodistillationは、2倍のデバイスを用いることで、標準的なデータ並列学習と同等の精度を達成し、1イテレーションあたりの通信量を最大1000倍まで削減できる。
• 1モデルあたりのデバイス数に対して線形的なスケーリングが可能であり、デバイス数を2倍にした場合、収束に必要な学習更新回数を半分にできる。
• n>2のモデルへのcodistillationのスケーリングは、混合した結果を示す。一部のケースでは性能向上が見られるが、一貫性はなく、しばしば2ウェイのcodistillationと同等かそれ未満にとどまる。
• 2つ以上のモデルへのスケーリングにおける課題は、モデルアーキテクチャおよびデータセット固有の特性（一般化挙動や損失関数の形状）に強く関連している。
• Codistillationは学習ダイナミクスに影響を与える正則化効果をもたらし、最適なパフォーマンスを得るにはこれを考慮する必要がある。
• 通信量が削減されているにもかかわらず、ウォールクロック時間の高速化は保証されず、モデルのサイズと通信帯域幅に依存する。大きなモデルでは、より大きな恩恵を受ける。