[論文レビュー] MixKD: Towards Efficient Distillation of Large-scale Language Models

Kevin J Liang, Weituo Hao+5 more|arXiv (Cornell University)|Nov 1, 2020

Topic Modeling参考文献 52被引用数 30

ひとこと要約

MixKD は、mixup によるデータ拡張と内挿サンプルに対する教師の照会を組み込むことで大規模言語モデルの知識蒸留を改善し、GLUE タスクでのコンパクトな学生モデルの一般化を強化し、より良い性能を引き出します。

ABSTRACT

Large-scale language models have recently demonstrated impressive empirical performance. Nevertheless, the improved results are attained at the price of bigger models, more power consumption, and slower inference, which hinder their applicability to low-resource (both memory and computation) platforms. Knowledge distillation (KD) has been demonstrated as an effective framework for compressing such big models. However, large-scale neural network systems are prone to memorize training instances, and thus tend to make inconsistent predictions when the data distribution is altered slightly. Moreover, the student model has few opportunities to request useful information from the teacher model when there is limited task-specific data available. To address these issues, we propose MixKD, a data-agnostic distillation framework that leverages mixup, a simple yet efficient data augmentation approach, to endow the resulting model with stronger generalization ability. Concretely, in addition to the original training examples, the student model is encouraged to mimic the teacher's behavior on the linear interpolation of example pairs as well. We prove from a theoretical perspective that under reasonable conditions MixKD gives rise to a smaller gap between the generalization error and the empirical error. To verify its effectiveness, we conduct experiments on the GLUE benchmark, where MixKD consistently leads to significant gains over the standard KD training, and outperforms several competitive baselines. Experiments under a limited-data setting and ablation studies further demonstrate the advantages of the proposed approach.

研究の動機と目的

• 精度を犠牲にせず、モデルサイズと推論コストを削減することによって大規模言語モデルの効率的な展開を動機づける。
• タスクデータが限られている場合の過学習と記憶化に対処するため、データ拡張によって学習データを充実させる。
• データに依存しない知識蒸留フレームワーク（MixKD）を提案し、mixup を活用して教師-生徒学習のための拡張サンプルを生成する。
• 拡張された KD による一般化の改善を理論的に正当化する。
• GLUE タスクで実証的な利得を示し、特にデータが少ない領域で基準となる KD および関連手法と比較する。

提案手法

• 大規模な教師（BERT）から小さな生徒（BERT-3 または BERT-6）へ知識蒸留を行う。
• 単語埋め込みに mixup を適用して、生徒と教師の監督の両方のための拡張入力を生成する（x' = λx_i + (1-λ)x_j; y' = λy_i + (1-λ)y_j）。
• 元データに対する標準交差エントロピー、ミックスアップデータに対する生徒の交差エントロピー、ミックスアップデータ上の教師と生徒の蒸留損失を組み合わせた結合目的関数で訓練する（L = L_MLE + α_SM L_SM + α_TMKD L_TMKD）。
• ミックスアップサンプルに対して教師を照会し、それらのサンプルに対する生徒の予測と比較して蒸留損失を最小化する（L_TMKD）。
• 拡張データが経験的誤差と一般化誤差のギャップを縮小する条件を示す理論的な結果を提供し、一般化を改善する。

実験結果

リサーチクエスチョン

• RQ1 MixKD における mixup ベースのデータ拡張は、標準的な KD と比較して小さな Student モデルの一般化を改善し、一般化ギャップを縮小しますか？
• RQ2 データ可用性の異なる状況で、GLUE タスクにおける MixKD の性能はどうか、DistilBERT、PKD、および他のベースラインとどう比較されるか？
• RQ3 KD における mixup 拡張とバック翻訳および他のデータ拡張技術の組み合わせの影響は何か？
• RQ4 特にデータが限られた状況で、ハイパーパラメータと mixup 比率に対する MixKD の感度はどの程度か？

主な発見

• MixKD は GLUE タスクで標準の KD および関連ベースラインを一貫して上回り、データが限定された設定で顕著な利得を示す。
• SM+TMKD+BT を組み合わせた 6 層 MixKD の生徒は、推論を大幅に高速化しパラメータ数を削減したまま、教師の性能の大部分をキャプチャすることが多い。
• SST-2 では、SMTKD とバック翻訳を組み合わせた MixKD が、スクラッチ学習や標準 KD と比べて教師との差をかなり縮小する。
• 理論的結果は、mixup によるデータ拡張が、いくつかのデータ生成シナリオにおいて一般化ギャップを縮小できることを示しており、経験的な利得を裏付ける。
• 埋め込みの可視化は、MixKD が生徒を介在サンプルとデータマニフォールドに沿うように導き、潜在空間の整理を改善することを示している。