MiniLMv2: Multi-Head Self-Attention Relation Distillation for Compressing Pretrained Transformers

Findings of the Association for Computational Linguistics: ACL-IJCNLP 2021, pages 2140–2151 August 1–6, 2021. ©2021 Association for Computational Linguistics 2140 MINILMv2: Multi-Head Self-Attention Relation Distillation for Compressing Pretrained Transformers Wenhui Wang Hangbo Bao Shaohan Huang Li Dong Furu Wei∗ Microsoft Research {wenwan,t-habao,shaohanh,lidong1,fuwei}@microsoft.com Abstract We generalize deep self-attention distillation in MINILM (Wang et al., 2020) by only us- ing self-attention relation distillation for task- agnostic compression of pretrained Transform- ers. In particular, we define multi-head self- attention relations as scaled dot-product be- tween the pairs of query, key, and value vec- tors within each self-attention module. Then we employ the above relational knowledge to train the student model. Besides its simplic- ity 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 num- ber 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 ex- amine the layer selection strategy for teacher models, rather than just relying on the last layer as in MINILM. We conduct extensive ex- periments on compressing both monolingual and multilingual pre-trained models. Exper- imental results demonstrate that our models1 distilled from base-size and large-size teachers (BERT, RoBERTa and XLM-R) outperform the state-of-the-art. 1 Introduction Pretrained Transformers (Radford et al., 2018; De- vlin et al., 2018; Dong et al., 2019; Yang et al., 2019; Joshi et al., 2019; Liu et al., 2019; Bao et al., 2020; Radford et al., 2019; Raffel et al., 2019; Lewis et al., 2019a) have been highly successful for a wide range of natural language processing tasks. However, these models usually consist of hundreds of millions of parameters and are getting bigger. It brings challenges for fine-tuning and online serv- ∗Contact person. 1Distilled models and code will be publicly available at https://aka.ms/minilm. ing in real-life applications due to the restrictions of computation resources and latency. Knowledge distillation (KD; Hinton et al. 2015, Romero et al. 2015) has been widely em- ployed to compress pretrained Transformers, which transfers knowledge of the large model (teacher) to the small model (student) by minimizing the differences between teacher and student features. Soft target probabilities (soft labels) and interme- diate representations are usually utilized to per- form KD training. In this work, we focus on task-agnostic compression of pretrained Transform- ers (Sanh et al., 2019; Tsai et al., 2019; Jiao et al., 2019; Sun et al., 2019b; Wang et al., 2020). The student models are distilled from large pretrained Transformers using large-scale text corpora. The distilled task-agnostic model can be directly fine- tuned on downstream tasks, and can be utilized to initialize task-specific distillation. DistilBERT (Sanh et al., 2019) uses soft target probabilities for masked language modeling pre- dictions and embedding outputs to train the stu- dent. The student model is initialized from the teacher by taking one layer out of two. Tiny- BERT (Jiao et al., 2019) utilizes hidden states and self-attention distributions (i.e., attention maps …

Content truncated to 500 words