[論文レビュー] MiniLMv2: Multi-Head Self-Attention Relation Distillation for Compressing Pretrained Transformers

Wenhui Wang, Hangbo Bao+3 more|arXiv (Cornell University)|Dec 31, 2020

Topic Modeling参考文献 45被引用数 15

ひとこと要約

本稿では、教師モデルと生徒モデルのアテンションヘッド数が一致しなくてもよいように設計された知識蒸留手法であるMiniLMv2を提案する。教師モデルと生徒モデル間のクエリ、キー、バリューベクトル間のスケールドドット積関係を細分化して蒸留することで、事前学習済みトランスフォーマーを圧縮する。この手法により、単言語および多言語タスクで最先端の性能を達成し、BERT Largeから蒸留された6×768モデルはBERT Baseを上回りながら2倍の高速化を実現する。

ABSTRACT

We generalize deep self-attention distillation in MiniLM (Wang et al., 2020) by only using self-attention relation distillation for task-agnostic compression of pretrained Transformers. In particular, we define multi-head self-attention relations as scaled dot-product between the pairs of query, key, and value vectors within each self-attention module. Then we employ the above relational knowledge to train the student model. Besides its simplicity and unified principle, more favorably, there is no restriction in terms of the number of student's attention heads, while most previous work has to guarantee the same head number between teacher and student. Moreover, the fine-grained self-attention relations tend to fully exploit the interaction knowledge learned by Transformer. In addition, we thoroughly examine the layer selection strategy for teacher models, rather than just relying on the last layer as in MiniLM. We conduct extensive experiments on compressing both monolingual and multilingual pretrained models. Experimental results demonstrate that our models distilled from base-size and large-size teachers (BERT, RoBERTa and XLM-R) outperform the state-of-the-art.

研究の動機と目的

• 教師モデルと生徒モデルのアテンションヘッド数が同一でなければならないという知識蒸留の制限を解消すること。
• アテンション分布だけでなく、自己注意関係の細分化された特徴を活用することで、モデル圧縮を向上させること。
• 特に多くの層を有する大規模モデルにおいて、蒸留に最適な教師モデルの層を同定すること。
• タスクに依存しない蒸留により、単言語および多言語設定で最先端の性能を達成すること。
• 教師モデルのアテンションヘッド数とは独立して、生徒モデルの設計を柔軟に行えるようにすること。

提案手法

• 本手法は、各自己注意モジュール内におけるクエリ、キー、バリューベクトルのペア間のスケールドドット積を、マルチヘッド自己注意関係として定義する。
• アテンションヘッドの出力を連結し、ユーザーが指定する数の関係ヘッドに分割する関係ヘッド機構を導入することで、ヘッド数が異なるモデル間でのアライメントを可能にする。
• 教師モデルと生徒モデルの関係行列間のL2損失を最小化することで、知識蒸留を実行する。対象はQ-Q、K-K、V-V関係である。
• 最終層だけでなく、教師モデルの任意の層から蒸留が可能であり、実験では大規模モデルにおいて上位～中央層がより優れた性能を示すことが判明している。
• マスク言語モデル学習と関係蒸留を統合した一貫した訓練目的を採用し、生徒モデルのアーキテクチャに制限を課さない。
• 関係ヘッド数は、生徒または教師のヘッド数の倍数である必要がなく、柔軟かつ細分化された知識移転を可能にする。

実験結果

リサーチクエスチョン

• RQ1 教師モデルと生徒モデルのアテンションヘッド数が同一でなくても、自己注意関係蒸留はモデル圧縮を改善できるか？
• RQ2 大規模教師モデルの上位～中央層（例：24層中12層目）から蒸留することで、最終層を使用する場合よりも優れた性能が得られるか？
• RQ3 関係ヘッド数を増やすことで、蒸留された生徒モデルの性能は向上するか？
• RQ4 マルチヘッド自己注意関係蒸留は、単言語および多言語タスクにおいて、既存の知識蒸留手法を上回る性能を発揮できるか？
• RQ5 自己注意関係を増やすことと、蒸留における計算コストのトレードオフは何か？

主な発見

• BERT Largeから蒸留された6×768モデルは、GLUEベンチマークでBERT Baseを上回り、推論速度は2.0倍高速である。
• RoBERTa Largeから蒸留された生徒モデルは、はるかに少ない学習例でRoBERTa Baseを上回る性能を示す。
• BERT Large や XLM-R Large などの大規模教師モデルにおいて、上位～中央層（例：24層中12層目）から蒸留すると、最終層からの蒸留よりも優れた性能が得られる。
• 関係ヘッド数を増やすことで性能が向上し、より細分化された知識移転が教師の自己注意モジュールをよりよく模倣することにつながる。
• Q-Q、K-K、V-V関係の転送は、Q-KおよびV-V関係の転送よりも優れた結果をもたらし、ほとんどのGLUEタスクでMobileBERTを上回りながらも推論が高速である。
• 本手法は柔軟かつ汎用的であり、アーキテクチャ的制約なしに、さまざまな大規模モデル（例：BERT、RoBERTa、XLM-R）から蒸留が可能である。