[論文レビュー] Distilling BERT into Simple Neural Networks with Unlabeled Transfer Data

Subhabrata Mukherjee, Ahmed Hassan Awadallah|arXiv (Cornell University)|Oct 4, 2019

Natural Language Processing Techniques参考文献 28被引用数 9

ひとこと要約

この論文では、ドメイン内ラベルなし転送データを用いて、大規模な事前学習済み BERT モデルを単純な RNN ベースの学生ネットワークに蒸留する手法を提案している。この手法により、教師モデルと同等またはそれ以上の性能が達成され、パラメータ数を最大26倍圧縮し、待機時間の遅延を51倍速くする。41言語にわたる F1 スコアの95％を維持する一方、ラベル付きデータは限定的である。

ABSTRACT

Recent advances in pre-training huge models on large amounts of text through self supervision have obtained state-of-the-art results in various natural language processing tasks. However, these huge and expensive models are difficult to use in practise for downstream tasks. Some recent efforts use knowledge distillation to compress these models. However, we see a gap between the performance of the smaller student models as compared to that of the large teacher. In this work, we leverage large amounts of in-domain unlabeled transfer data in addition to a limited amount of labeled training instances to bridge this gap for distilling BERT. We show that simple RNN based student models even with hard distillation can perform at par with the huge teachers given the transfer set. The student performance can be further improved with soft distillation and leveraging teacher intermediate representations. We show that our student models can compress the huge teacher by up to 26x while still matching or even marginally exceeding the teacher performance in low-resource settings with small amount of labeled data. Additionally, for the multilingual extension of this work with XtremeDistil (Mukherjee and Hassan Awadallah, 2020), we demonstrate massive distillation of multilingual BERT-like teacher models by upto 35x in terms of parameter compression and 51x in terms of latency speedup for batch inference while retaining 95% of its F1-score for NER over 41 languages.

研究の動機と目的

• BERT のような大規模で計算コストの高い事前学習済み言語モデルを、リソースが限られた環境や本番環境に導入する課題に対処すること。
• 大量のドメイン内ラベルなし転送データを活用することで、リソースが限られた環境における小規模な学生モデルの性能を向上させること。
• 単純な RNN ベースの学生アーキテクチャが、知識蒸留により大規模な BERT 教師モデルと同等またはそれ以上の性能を達成できるかどうかを検証すること。
• トレーニングスケジュール、蒸留目的、および教師モデルの容量が学生モデルの性能に与える影響を調査すること。

提案手法

• 2段階のトレーニングスケジュールを用いる：まずラベルなしデータ上で表現損失を最適化し、次に段階的に層をアンフリークする形で、ラベル付きデータ上でログイット損失と交差エントロピー損失を同時に最適化する。
• 教師モデルの予測されたクラスラベルを学生モデルのターゲットとして使用するハード蒸留を適用する。
• 教師の出力から得られる温度調整済みのソフトラベルを用いて一般化性能を向上させる、ソフト蒸留による性能向上。
• 最終ログイットを超えた中間隠れ表現を教師ネットワークから活用し、学生の学習を支援する。
• ラベルなしデータ上で事前学習を行わず、代わりに教師からの知識蒸留に依存して学生モデル（BiLSTM ベース）をトレーニングする。
• ラベル付きデータ上の交差エントロピー損失と、教師の出力および隠れ状態からの蒸留損失の組み合わせを用いて学生モデルを最適化する。

実験結果

リサーチクエスチョン

• RQ1 知識蒸留とドメイン内ラベルなしデータを用いてトレーニングされた単純な RNN ベースの学生モデルは、大規模な BERT 教師モデルと同等の性能を達成できるか？
• RQ2 大規模なラベルなし転送データの導入が、リソースが限られた環境における学生モデルの性能にどのように影響するか？
• RQ3 小規模な学生モデルの蒸留性能を最適化するには、どのトレーニングスケジュールと損失の組み合わせが最良か？
• RQ4 ソフト蒸留や中間表現蒸留は、ハード蒸留を上回る学生モデルの性能向上をもたらすか？
• RQ5 学生モデルのサイズが固定された場合、教師モデルのサイズが蒸留性能に与える影響はいかほどか？

主な発見

• 単純な BiLSTM ベースの学生モデルは、パラメータ数を最大26倍圧縮しても、BERT Base や BERT Large と同等の性能を達成する。
• ラベル付きデータが限定的なリソースが限られた環境では、学生モデルの性能が教師モデルをわずかに上回る。
• ソフト蒸留と中間表現蒸留は、ハード蒸留を上回る学生モデルの性能向上をもたらす。
• 多言語 BERT に類似したモデルでは、パラメータ数を最大35倍圧縮し、待機時間の遅延を51倍速くする一方で、41言語にわたる F1 スコアの95％を維持する。
• 段階的な層のアンフリークを伴う段階的トレーニングは、同時に最適化する方法と比較して、最も優れた蒸留性能を達成する。
• 学生モデルのサイズが飽和した段階で、より大きな教師モデルは、蒸留性能にわずかな向上をもたらすにとどまる。