[論文レビュー] Progressively Stacking 2.0: A Multi-stage Layerwise Training Method for BERT Training Speedup

Cheng Yang, Shengnan Wang+4 more|arXiv (Cornell University)|Nov 27, 2020

Topic Modeling参考文献 22被引用数 10

ひとこと要約

本稿では、BERTの事前学習を高速化するためのマルチステージレイヤーワイズトレーニング（MSLT）という手法を提案する。この手法は、段階的に浅いモデルを積み上げていき、各段階で上位レイヤーのみを微調整しながら、以前に学習済みのレイヤーは固定する。このアプローチにより、バックワードパスの計算と通信オーバーヘッドを削減でき、標準的なBERT学習と比較して110%以上の学習スピードアップを達成しながら、性能の低下は最小限に抑えられる。

ABSTRACT

Pre-trained language models, such as BERT, have achieved significant accuracy gain in many natural language processing tasks. Despite its effectiveness, the huge number of parameters makes training a BERT model computationally very challenging. In this paper, we propose an efficient multi-stage layerwise training (MSLT) approach to reduce the training time of BERT. We decompose the whole training process into several stages. The training is started from a small model with only a few encoder layers and we gradually increase the depth of the model by adding new encoder layers. At each stage, we only train the top (near the output layer) few encoder layers which are newly added. The parameters of the other layers which have been trained in the previous stages will not be updated in the current stage. In BERT training, the backward computation is much more time-consuming than the forward computation, especially in the distributed training setting in which the backward computation time further includes the communication time for gradient synchronization. In the proposed training strategy, only top few layers participate in backward computation, while most layers only participate in forward computation. Hence both the computation and communication efficiencies are greatly improved. Experimental results show that the proposed method can achieve more than 110% training speedup without significant performance degradation.

研究の動機と目的

• 大規模なBERTモデルにおける高い計算コストと長い学習時間の問題に対処すること。
• 分散環境におけるバックワードパスと通信オーバーヘッドを削減することで、学習の効率を向上させること。
• 既存の段階的積み上げ手法では、後段のステージで全モデルの更新が必要なため、遅延が生じるという非効率性を克服すること。
• 下位レイヤーが初期に安定することに着目し、継続的な再学習が無駄であることを活用すること。
• モデルの性能を維持しつつ、学習時間を著しく短縮する学習戦略を開発すること。

提案手法

• 本手法はBERT学習を複数の段階に分解し、最初は数層のエンコーダーレイヤーを含む浅いモデルから開始する。
• 各段階において、新たに追加された上位レイヤーのみを更新し、以前に学習済みのすべてのレイヤーは固定され、前方伝播のみに使用される。
• 段階を重ねるごとに、同じ事前学習目的を維持したまま、新たなレイヤーを追加することでモデルの深さを段階的に増加させる。
• すべてのレイヤーが追加された後、最終的な再トレーニングフェーズで全パラメータを数ステップにわたり同時に更新し、モデルを精緻化する。
• このアプローチは、交互最適化の原則を活用しており、一度に小さなパラメータサブセットのみを更新することで、収束効率を向上させる。
• 未学習の上位レイヤーを最初に学習し、その後、小さな学習率で全レイヤーを同時に再トレーニングすることで、学習率の競合を回避する。

実験結果

リサーチクエスチョン

• RQ1 深層BERTモデルにおいて、上位レイヤーのみを更新することで、学習効率を著しく向上させることができるか？
• RQ2 後段のステージで以前に学習済みのレイヤーを固定することで、学習時間を短縮できるが、モデル性能に悪影響を及えないか？
• RQ3 段階的積み上げ戦略は、標準的な学習と比較して、スピードと精度の両面でどのように異なるか？
• RQ4 下位レイヤーのアテンション分布が初期に安定することに着目し、これを学習を高速化するのに活用できるか？
• RQ5 マルチステージレイヤーワイズトレーニングにおいて、スピードアップと性能低下の最適なトレードオフは何か？

主な発見

• 提案されたMSLT手法は、同じ数の学習ステップを用いた標準的なBERT学習と比較して、110%以上の学習スピードアップを達成した。
• 同じ学習時間内に、標準的な学習よりも優れた性能を達成しており、優れたサンプル効率を示している。
• 標準BERTと同等のモデル性能を維持しており、GLUEベンチマークでのわずかな低下（例：BERT-Base：平均スコア79.5 vs. 79.8）にとどまる。
• 最終的な再トレーニングフェーズは、全学習ステップの約20%を占めており、この段階以前にモデルはすでに十分に最適化されていることが示された。
• 特に分散学習環境において、各段階でのパラメータ更新が限定的であるため、バックワード計算と勾配同期のオーバーヘッドを顕著に削減した。
• 実験結果から、下位レイヤーのアテンション分布が後段のステージでも大きく変化しないことが確認され、レイヤーを固定するという仮定の妥当性が裏付けられた。