[論文レビュー] MiniLM: Deep Self-Attention Distillation for Task-Agnostic Compression of Pre-Trained Transformers

Wenhui Wang, Furu Wei+4 more|arXiv (Cornell University)|Feb 25, 2020

Topic Modeling参考文献 57被引用数 632

ひとこと要約

tldr: この論文は、教師の最後のレイヤーの self-attention のみを模倣することで、大規模 Transformer LMs を圧縮するタスク非依存の深い自己注意蒸留手法である MiniLM を紹介します。さらに値の関係性を利用し、柔軟な student アーキテクチャと、はるかにパラメータ数を減らしつつ高い性能を実現します。

ABSTRACT

Pre-trained language models (e.g., BERT (Devlin et al., 2018) and its variants) have achieved remarkable success in varieties of NLP tasks. However, these models usually consist of hundreds of millions of parameters which brings challenges for fine-tuning and online serving in real-life applications due to latency and capacity constraints. In this work, we present a simple and effective approach to compress large Transformer (Vaswani et al., 2017) based pre-trained models, termed as deep self-attention distillation. The small model (student) is trained by deeply mimicking the self-attention module, which plays a vital role in Transformer networks, of the large model (teacher). Specifically, we propose distilling the self-attention module of the last Transformer layer of the teacher, which is effective and flexible for the student. Furthermore, we introduce the scaled dot-product between values in the self-attention module as the new deep self-attention knowledge, in addition to the attention distributions (i.e., the scaled dot-product of queries and keys) that have been used in existing works. Moreover, we show that introducing a teacher assistant (Mirzadeh et al., 2019) also helps the distillation of large pre-trained Transformer models. Experimental results demonstrate that our monolingual model outperforms state-of-the-art baselines in different parameter size of student models. In particular, it retains more than 99% accuracy on SQuAD 2.0 and several GLUE benchmark tasks using 50% of the Transformer parameters and computations of the teacher model. We also obtain competitive results in applying deep self-attention distillation to multilingual pre-trained models.

研究の動機と目的

• 大規模な事前学習済み Transformer LMs（例：BERT）を圧縮してファインチューニングとデプロイを高速化する動機づけ。
• 最後のレイヤーから教師の自己注意を深く模倣するタスク非依存の蒸留フレームワークを提案。
• 追加パラメータなしで転送可能な自己注意の値-relations（自動的に学習される知識）を導入。
• 小さな student（例：6層、768-d）でも教師に近い性能と大幅なスピードアップを達成できることを示す。
• 教師アシスタントが、特に非常に小さな student に対して性能をさらに向上させることができることを示す。

提案手法

• student が教師の最後の Transformer レイヤーの自己注意モジュールを深く模倣するよう訓練。
• 自己注意分布（クエリ–キー）と値のスカラー付きドット積（値-relations）を知識として転送。
• KL発散を用いて教師と student の自己注意分布間の注意マップ転送損失を計算。
• KL発散を用いて教師と student の値-relations 行列間の転送損失を計算；この転送には追加パラメータは不要。
• 任意で教師アシスタント（中間サイズの student）を使い、教師と student のギャップを埋めて性能を改善。
• 従来のタスク非依存蒸留法と比較し、最終層・値-relations・TA の利点を実証。

実験結果

リサーチクエスチョン

• RQ1 タスク非依存蒸留は、教師の最後のレイヤー自己注意のみを模倣する場合でも効果的か。
• RQ2 自己注意分布とともに値-relations を転送することで、より深い模倣と student の性能向上が得られるか。
• RQ3 教師アシスタントを導入すると、特に小さな student に蒸留の効果が改善されるか。
• RQ4 レイヤー間のマッピングを必要とせず、柔軟な student アーキテクチャ（レイヤー数・隠れ層サイズの変更）をサポートできるか。

主な発見

• 6層・768-hidden の MiniLM student は BERT-BASE から蒸留され、SQuAD 2.0 および GLUE タスクで高いパフォーマンスを保ちつつ大幅な高速化を実現。
• 教師の最後のレイヤーから自己注意分布と値-relations の両方を転送することで、自己注意分布のみや他のベースラインを用いた場合よりも測定可能な改善を得られる。
• 値-relations の転送は追加パラメータを導入せず、より深い自己注意の模倣を提供し、複数のタスクと様々な student 構成での結果を改善。
• 教師アシスタントは、小さな student に対してさらに性能を向上させ、教師と student のギャップを埋めるのに役立つ。
• MiniLM は、はるかに少ない Transformer パラメータで競争力のある性能を持つ多言語モデルを可能にする。