[論文レビュー] Interlocking Backpropagation: Improving depthwise model-parallelism

Aidan N. Gomez, Oscar Key+5 more|arXiv (Cornell University)|Oct 8, 2020

Advanced Neural Network Applications参考文献 8被引用数 5

ひとこと要約

本論文は、深層ネットワークの定期的な間隔に補助分類ヘッドを挿入することで、勾配伝播の流れを制限し、計算効率を向上させる、ハイブリッドトレーニング戦略「インターロッキングバックプロパゲーション」を提案する。この手法は、エンドツーエンドバックプロパゲーションに比べて顕著に高速なトレーニングを実現しながら、完全に局所学習にした場合に失われる性能の大部分を回復し、固定時間制約下でのResNetsおよびTransformerにおいて、トレーニング効率とテスト精度の両面で優れた性能を発揮する。

ABSTRACT

The number of parameters in state of the art neural networks has drastically increased in recent years. This surge of interest in large scale neural networks has motivated the development of new distributed training strategies enabling such models. One such strategy is model-parallel distributed training. Unfortunately, model-parallelism can suffer from poor resource utilisation, which leads to wasted resources. In this work, we improve upon recent developments in an idealised model-parallel optimisation setting: local learning. Motivated by poor resource utilisation in the global setting and poor task performance in the local setting, we introduce a class of intermediary strategies between local and global learning referred to as interlocking backpropagation. These strategies preserve many of the compute-efficiency advantages of local optimisation, while recovering much of the task performance achieved by global optimisation. We assess our strategies on both image classification ResNets and Transformer language models, finding that our strategy consistently out-performs local learning in terms of task performance, and out-performs global learning in training efficiency.

研究の動機と目的

• グローバルモデル並列トレーニングにおけるリソースの非効率な利用と、局所学習における性能の低下を解消すること。
• 局所学習とグローバル学習の間の妥協戦略を開発し、計算効率を保ちつつタスクパフォーマンスを向上させること。
• 大規模分散トレーニングにおける、トレーニング速度とモデル精度のトレードオフを評価すること。
• インターロッキングバックプロパゲーション戦略を研究するための汎用的でオープンソースのフレームワークを提供すること。

提案手法

• 深層ネットワークの定期的な間隔に補助分類ヘッドを追加することで、局所的なトレーニング信号を生成するnワイズインターロッキングバックプロパゲーションを導入する。
• モジュール間の勾配伝播を制限することで、完全な逆伝播依存関係を回避し、より早い重み更新を可能にする。
• 補助ヘッドからの局所損失を用いて各モジュールを独立して学習させ、無駄な待機時間と通信ボトルネックを低減する。
• 固定時間トレーニングの視点を採用し、ステップ時間の短い戦略を優遇して性能を評価する。
• ResNetsおよびTransformerベースの言語モデルにこの手法を適用し、Adamを用いてトレーニングし、Perplexityを測定する。
• マイクロバッチ処理とTPUベースの分散トレーニングを活用して、複数のアクセラレータに渡るスケーリングを実現する。

実験結果

リサーチクエスチョン

• RQ1 局所学習とグローバル学習の中間戦略としてのハイブリッドトレーニング戦略は、深さ方向モデル並列性におけるトレーニング効率とモデルパフォーマンスの両面を向上させられるか？
• RQ2 インターロッキングバックプロパゲーションは、エンドツーエンド学習および1ワイズ局所学習と比較して、トレーニング速度とテスト精度の面でどのように異なるか？
• RQ3 補助分類ヘッドは、完全に局所学習にした場合に失われる性能をどの程度回復できるか？
• RQ4 固定時間制約下でも、インターロッキングバックプロパゲーションはTransformerのような大規模モデルに効果的にスケーリングできるか？

主な発見

• 2ワイズインターロッキングバックプロパゲーションは、4モジュールのTransformerモデルにおいて、標準的なエンドツーエンドトレーニングの半分未満のトレーニングステップ時間を達成しながら、同等のテストパーセプルリティを達成した。
• 固定時間制約下では、インターロッキングバックプロパゲーションはエンドツーエンドトレーニングを上回るモデルパフォーマンスを発揮しており、同じ時間内により多くの重み更新が可能だからである。
• CIFAR-10およびCIFAR-100において、2ワイズインターロッキングは、エンドツーエンド学習および1ワイズ学習を上回るテスト精度を示したが、それにもかかわらずトレーニングがより高速であった。
• Transformer言語モデルにおいて、インターロッキングバックプロパゲーションは1ワイズとエンドツーエンド学習の間のパフォーマンスギャップを大幅に縮小したが、パーセプルリティにおいてエンドツーエンドを上回ることはできなかった。
• この手法は、nワイズインターロッキングのnパラメータを変化させることで、計算効率とモデル性能の滑らかなトレードオフを実現できることを示している。
• https://github.com/oscarkey/interlocking-backprop に公開されたオープンソースフレームワークにより、この最適化アルゴリズムクラスの再現可能な研究が可能になった。