[Paper Review] 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 Modeling57 references632 citations

TL;DR

This paper introduces MiniLM, a task-agnostic, deep self-attention distillation method to compress large Transformer LMs by mimicking only the teacher’s last layer self-attention, plus value-relations, enabling flexible student architectures and strong performance with significantly fewer parameters.

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.

Motivation & Objective

• Motivate compressing large pre-trained Transformer LMs (e.g., BERT) for faster fine-tuning and serving.
• Propose a task-agnostic distillation framework that deeply mimics the teacher’s self-attention from the last layer.
• Introduce self-attention value-relations as additional deep knowledge to transfer without extra parameters.
• Demonstrate that a smaller student (e.g., 6-layer, 768-d) can achieve near-teacher performance with substantial speedups.
• Show that a teacher assistant can further boost performance, especially for very small students.

Proposed method

• Train a student to deeply mimic the self-attention module of the teacher’s last Transformer layer.
• Transfer both self-attention distributions (queries–keys) and the scaled dot-product between values (value-relations) as knowledge.
• Compute attention map transfer loss via KL-divergence between teacher and student attention distributions.
• Compute value-relations transfer loss via KL-divergence between teacher and student value-relations matrices; no extra parameters are needed for this transfer.
• Optionally use a teacher assistant (intermediate-size student) to bridge large teacher–student gaps and improve performance.
• Compare with prior task-agnostic distillation methods and demonstrate benefits of last-layer, value-relations, and TA.

Experimental results

Research questions

• RQ1 Can task-agnostic distillation be effective when only the teacher’s last-layer self-attention is mimicked?
• RQ2 Does transferring the value-relations in addition to attention distributions yield deeper mimicry and better student performance?
• RQ3 Does introducing a teacher assistant improve distillation, particularly for smaller students?
• RQ4 Can the approach support flexible student architectures (varying number of layers and hidden sizes) without requiring layer-to-layer mappings?

Key findings

• A 6-layer, 768-hidden MiniLM student distilled from BERT-BASE is significantly faster while preserving high performance on SQuAD 2.0 and GLUE tasks.
• Transferring both attention distributions and value-relations from the teacher’s last layer yields measurable gains over using attention distributions alone or other baselines.
• Value-relations transfer provides deeper self-attention mimicry without introducing extra parameters, improving results across multiple tasks and student configurations.
• A teacher assistant further boosts performance for smaller students, helping bridge the gap between teacher and student.
• MiniLM enables multilingual models with competitive performance using substantially fewer Transformer parameters.