[論文レビュー] LightPAFF: A Two-Stage Distillation Framework for Pre-training and Fine-tuning

Kaitao Song, Hao Sun+5 more|arXiv (Cornell University)|Apr 27, 2020

Topic Modeling参考文献 17被引用数 9

ひとこと要約

LightPAFF は、大規模な事前学習モデル（例：BERT、GPT-2、MASS）から、事前学習段階および微調整段階の両方で、より小さな学生モデルへ知識を転送する二段階の知識蒸留フレームワークであり、最大5倍小さいモデルサイズと5〜7倍速い推論速度を実現し、性能は元のモデルとほぼ同等であるため、リソース制限のあるシステムへの効率的なデプロイが可能である。

ABSTRACT

While pre-training and fine-tuning, e.g., BERT~\citep{devlin2018bert}, GPT-2~\citep{radford2019language}, have achieved great success in language understanding and generation tasks, the pre-trained models are usually too big for online deployment in terms of both memory cost and inference speed, which hinders them from practical online usage. In this paper, we propose LightPAFF, a Lightweight Pre-training And Fine-tuning Framework that leverages two-stage knowledge distillation to transfer knowledge from a big teacher model to a lightweight student model in both pre-training and fine-tuning stages. In this way the lightweight model can achieve similar accuracy as the big teacher model, but with much fewer parameters and thus faster online inference speed. LightPAFF can support different pre-training methods (such as BERT, GPT-2 and MASS~\citep{song2019mass}) and be applied to many downstream tasks. Experiments on three language understanding tasks, three language modeling tasks and three sequence to sequence generation tasks demonstrate that while achieving similar accuracy with the big BERT, GPT-2 and MASS models, LightPAFF reduces the model size by nearly 5x and improves online inference speed by 5x-7x.

研究の動機と目的

• オンラインまたはモバイル環境での大規模事前学習モデル（例：BERT、GPT-2）のデプロイを困難にする高メモリ使用量および高レイテンシの課題に対処する。
• 微調整段階でのみモデルを圧縮する従来の知識蒸留手法では、性能の著しい低下が生じるという限界を克服する。
• マスク言語モデリング（BERT）、因果的言語モデリング（GPT-2）、マスクされたシーケンス対シーケンスモデリング（MASS）といった多様な事前学習目的に適用可能な統一されたフレームワークを開発する。
• 事前学習済みおよび微調整済みの教師モデルからの知識を活用することで、軽量モデルにおける一般言語理解および生成能力を維持する。
• 複数の NLP タスク（分類、言語モデリング、シーケンス対シーケンス生成）において、性能を損なわず、効率的かつ低レイテンシの推論を実現する。

提案手法

• 事前学習段階では事前学習データセットを用い、微調整段階ではタスク固有のデータを用いて、二段階の知識蒸留を適用する。
• 教師モデルよりパラメータ数が少ない学生モデルを、両段階で教師モデルのログイットおよび隠れ表現を模倣するように訓練する。
• 訓練中に知識蒸留損失と正解ラベルからの標準的な交差エントロピー損失のバランスを取るために、ハイパーパrameter λ を導入する。
• パラメータごとの知識蒸留を実装し、学生モデルと教師モデルのパラメータ間の L2 距離を最小化するが、耐性を高めるためにガウスノイズを追加する。
• BERT、GPT-2、MASS の全モデルに対して一貫した形で適用可能な、統一された蒸留目的を定式化する。
• 二段階パイプラインを採用する：(1) 教師モデルの事前学習、(2) 下流タスクにおける教師モデルの微調整、(3) 事前学習済み教師モデルから軽量な事前学習学生モデルへの蒸留、(4) 下流タスクにおける学生モデルの微調整（蒸留付き）

実験結果

リサーチクエスチョン

• RQ1 事前学習段階および微調整段階の両方で知識蒸留を適用することで、モデルサイズと推論レイテンシを著しく削減しつつ、性能を維持できるか？
• RQ2 微調整段階でのみ蒸留を行う単一段階蒸留と比較して、二段階蒸留は軽量モデルの精度および耐性においてどのように異なるか？
• RQ3 事前学習タスクの難易度（例：マスク vs. 因果的言語モデリング）が最適な蒸留ハイパーパrameterやモデル性能に影響を与えるか？
• RQ4 提案された LightPAFF フレームワークは、異なる事前学習目的を持つ BERT、GPT-2、MASS といった多様な事前学習モデルに一貫して適用可能か？
• RQ5 微調整段階でのみ蒸留を行う場合と比較して、両段階での蒸留が一般化性能および耐性をどの程度向上させるか？

主な発見

• LightPAFF は、BERT、GPT-2、MASS のモデルサイズをほぼ 5 倍小さくしつつ、元の大きなモデルと同等の性能を維持している。
• 評価済みのすべてのモデルにおいて、オンライン推論速度が 5 倍から 7 倍速くなり、エッジデバイスへの実用的デプロイが可能である。
• 言語理解（SST-2、QQP、多義語解消）、言語モデリング（WikiText-2、PTB、WikiText-103）、シーケンス対シーケンス生成（WMT Zh-En、En-De、En-Fr）の合計 9 つの下流タスクにおける実験では、圧縮モデルにおいて一貫した性能向上が確認された。
• 二段階蒸留アプローチは、パラメータ摂動に対する耐性向上と、事前学習段階および微調整段階の両方で安定した性能を示すことで、より優れた一般化性能を実現している。
• 最適な蒸留ハイパーパrameter λ はモデルによって異なる：BERT では λ = 1.0、MASS では λ = 0.7、GPT-2 では λ = 0.4 であり、事前学習タスクの難易度および情報保持度の差を反映している。
• 事前学習タスクの難易度と性能には相関がある：BERT（平均で入力トークンの 85% が観測可能）が最も簡単で、次に MASS（75%）、GPT-2（50%）が続く。これは、マスク予測タスクにおける観察された精度低下と整合的である。