Source-linked AI summary
MiniLMv2: Multi-Head Self-Attention Relation Distillation for Compressing Pretrained Transformers
Wenhui Wang, Hangbo Bao, Shaohan Huang, Li Dong, Furu Wei
TL;DR
The paper addresses restrictions and inefficiencies in task-agnostic Transformer compression, particularly the requirement that teacher and student attention-head counts match. It trains students with multi-head query, key, and value self-attention relations and examines teacher-layer selection, with distilled BERT, RoBERTa, and XLM-R models outperforming state-of-the-art methods.
Problem
Most previous distillation methods rely on self-attention distributions, which restrict student models to the teacher’s number of attention heads during task-agnostic compression.
Method
The method distills scaled dot-product relations among query, key, and value vectors and investigates which teacher layer should provide the transferred self-attention knowledge.
Results
Distilled models from BERT, RoBERTa, and XLM-R outperform state-of-the-art models across different parameter sizes and monolingual or multilingual settings.
Takeaways & Limitations
The approach provides fine-grained self-attention distillation without restricting the student’s attention-head count, while upper-middle layers benefit large-size teacher distillation.
Takeaways & Limitations
The model performs relatively worse on CoLA than MobileBERT, while MLM fine-tuning improves CoLA but slightly reduces performance on other GLUE tasks.
Abstract
from arXiv · showhide
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.
1 Introduction
Pretrained Transformers offer strong NLP performance but create computation and latency challenges as their size grows. The paper proposes multi-head self-attention relation distillation to compress them with flexible student attention-head counts and reports strong monolingual and multilingual results.
- Pretrained Transformers contain hundreds of millions of parameters, creating challenges for fine-tuning and real-life online applications under computation and latency constraints.
- Most prior task-agnostic distillation methods rely on self-attention distributions, restricting students to the same number of attention heads as teachers.
- The method transfers fine-grained relations among query, key, and value vectors, aligning teacher and student relations even when their attention-head counts differ.
- Distilled models from BERT, RoBERTa, and XLM-R outperform state-of-the-art models across monolingual and multilingual settings.
- Experiments find that an upper-middle teacher layer performs better for large-size teachers, while the last layer performs better for base-size teachers.
2 Related Work
The paper situates its approach within Transformer self-attention and task-agnostic knowledge distillation. It replaces conventional attention-distribution transfer with query, key, and value relations while examining teacher-layer selection.
- Transformer backbone: Each Transformer layer combines a multi-head self-attention sub-layer with a position-wise feedforward sub-layer.
- Transformer backbone: Self-attention projects the previous layer’s output into queries, keys, and values, then uses scaled dot-products to weight value vectors.
- Pretrained language models: Pretrained language models learn from large text corpora before fine-tuning for specific tasks, including monolingual and multilingual applications.
- Relation distillation: The proposed method transfers Q-Q, K-K, and V-V relations, and uses an upper-middle teacher layer for large-size models rather than relying only on the last layer.
- Knowledge distillation: Knowledge distillation compresses large models by transferring soft targets, hidden states, attention distributions, or value relations to students.
3 Multi-Head Self-Attention Relation Distillation
The approach deeply mimics teacher self-attention using multi-head relations among queries, keys, and values, avoiding equal attention-head counts between teacher and student. Experiments also examine teacher-layer selection for distillation.
- Multi-Head Self-Attention Relation Distillation: The method distills self-attention relations instead of attention distributions, targeting the dependency-drawing module in Transformers.The relations are formed from pairs of queries, keys, and values.
- Multi-Head Self-Attention Relations: Concatenating and splitting vectors lets teacher and student models with different attention-head counts align to the same relation-head count.The same operation is applied to queries, keys, and values.
- Multi-Head Self-Attention Relations: The training objective uses KL divergence between teacher and student multi-head self-attention relations.The relations are computed from queries, keys, and values across multiple relation heads.
- Multi-Head Self-Attention Relations: Query-query, key-key, and value-value relations are transferred to balance performance and training cost.The method can use multiple relation types rather than only self-attention distributions.
- Teacher Layer Selection: Experiments find that transferring knowledge from an upper-middle teacher layer achieves better performance than relying only on the last layer.The study examines large-size teachers and transfers one teacher layer to the student’s last layer.
4 Experiments
Experiments evaluate MiniLMv2 on English and multilingual downstream benchmarks across teacher models, student sizes, layers, and relation-head configurations. The method consistently outperforms strong baselines while supporting efficient compression and flexible relation alignment.
- Evaluation setup: Experiments cover BERT, RoBERTa, and XLM-R teachers across GLUE, SQuAD 2.0, XNLI, and MLQA.English students are evaluated on GLUE and extractive question answering, while multilingual students are evaluated on cross-lingual inference and question answering.
- Main results: Our 6×768 model outperforms DistilBERT, TinyBERT, MINILM, and two BERT baselines across most GLUE and SQuAD 2.0 tasks.The comparison uses students distilled from BERTBASE, BERTLARGE, and RoBERTaLARGE.
- Main results: 6×768 BERTBASE students retain more than 99% of teacher accuracy while using 50% of Transformer parameters.The corresponding student distilled from BERTLARGE compares favorably with BERTBASE.
- Main results: Base-size students distilled from RoBERTaLARGE and BERTLARGE outperform BERTBASE and RoBERTaBASE, while RoBERTaLARGE students use much fewer training examples.The RoBERTaLARGE student uses an almost 32× smaller training batch size and fewer training steps than RoBERTaBASE.
- Multilingual results: The 6×384 XLM-R student outperforms mBERT with 5.3× speedup and performs better than MINILM on multilingual evaluations.The results further validate multi-head self-attention relation distillation for multilingual models.
- Ablations: Upper-middle teacher layers work best for BERTLARGE and XLM-RLARGE, whereas the last BERTBASE layer performs best.The same upper-middle-layer trend appears for BERTLARGE-WWM and RoBERTaLARGE.
- Ablations: Larger relation-head counts improve performance by capturing finer-grained self-attention knowledge, without requiring a positive-multiple relationship between teacher and student heads.Relation heads may contain fragments of one attention head or fragments from multiple attention heads.
5 Discussion
The discussion compares MiniLMv2 with MobileBERT and examines a task-specific limitation and the computational trade-off of adding more relations. Results support broad teacher and student flexibility, but gains are not uniform across tasks or configurations.
- Comparison with MobileBERT: MiniLMv2 uses standard Transformer students for compressing BERT and RoBERTa, so it is not directly comparable with MobileBERT’s specially designed architecture.A same-parameter-size reference student is provided using BERTLARGE-WWM as teacher.
- Comparison with MobileBERT: Our method can be applied to different teachers with fewer restrictions on student architectures and attention-head configurations.The method’s flexibility follows from transferring self-attention relations rather than requiring matching attention distributions.
- Comparison with MobileBERT: Our model performs relatively worse on CoLA than MobileBERT, whose MLM-objective fine-tuning improves CoLA but slightly reduces other GLUE tasks in preliminary experiments.CoLA requires fine-grained linguistic knowledge that can be learned from language-modeling objectives.
- Computational trade-offs: Adding more self-attention relations improves results across most tasks, especially for BERT students, but increases computational cost.The authors choose a configuration intended to balance performance and computational cost.
6 Conclusion
The paper trains students with multi-head self-attention relations, removes the student-head-count restriction, and finds upper-middle-layer transfer beneficial for large teachers. Models distilled from BERT, RoBERTa, and XLM-R achieve competitive performance and outperform state-of-the-art methods.
- Multi-head self-attention relations train the student with more fine-grained knowledge and no restriction on the number of student attention heads.
- Transferring self-attention knowledge from an upper middle teacher layer achieves better performance for large-size teachers.
- Models distilled from BERT, RoBERTa, and XLM-R achieve competitive performance and outperform state-of-the-art methods.The experiments cover both monolingual and multilingual models.
A GLUE Benchmark
This section presents the GLUE benchmark dataset summary in Table 11. The table concerns dataset statistics and evaluation metrics for GLUE.
- Table 11 presents the dataset statistics used for the GLUE benchmark.
- The reported datasets are used for the General Language Understanding Evaluation benchmark.
- GLUE dataset statistics and metrics are summarized together in Table 11.
B SQuAD 2.0
This section presents dataset statistics and metrics for SQuAD 2.0. The summary is provided in Table 12 for the extractive question-answering evaluation.
- Table 12 presents the dataset statistics and metrics for SQuAD 2.0.
C Hyper-parameters for Fine-tuning
The fine-tuning setup specifies sequence lengths, batch sizes, learning-rate ranges, epochs, and regularization settings for SQuAD 2.0, GLUE, XNLI, and MLQA.
- Extractive Question Answering: SQuAD 2.0 uses maximum sequence length 384, batch size 32, three epochs, and learning rates from {3e-5, 6e-5, 8e-5, 9e-5}.Warmup ratio and weight decay are 0.1 and 0.01.
- GLUE: GLUE uses maximum sequence length 128, batch size 32, task-specific learning rates, and task-specific epoch counts.CoLA is fine-tuned for 25 epochs; warmup ratio and weight decay are 0.1 and 0.01.
- Cross-lingual Natural Language Inference: XNLI uses maximum sequence length 256, batch size 64, 10 epochs, and learning rates from {3e-5, 4e-5, 5e-5, 6e-5}.
- Cross-lingual Question Answering: MLQA uses maximum sequence length 512, batch size 32, four epochs, and learning rates from {3e-5, 4e-5, 5e-5, 6e-5}.