PaperlayerView exact revision on arXiv

Source-linked AI summary

MiniLM: Deep Self-Attention Distillation for Task-Agnostic Compression of Pre-Trained Transformers

Wenhui Wang, Furu Wei, Li Dong, Hangbo Bao, Nan Yang, Ming Zhou

arXiv:2002.10957v2cs.CL

TL;DR

Large pre-trained language models are powerful but expensive to fine-tune and serve because of their size. The paper proposes deep self-attention distillation, training smaller students to mimic a teacher’s last-layer self-attention and value relations. MiniLM outperforms baselines across student sizes and retains more than 99% performance on several benchmarks while using substantially fewer resources.

  • Problem

    Large pre-trained language models create latency and capacity challenges for fine-tuning and online serving because they contain hundreds of millions of parameters.

  • Method

    Deep self-attention distillation trains a student to mimic the teacher’s last-layer self-attention distributions and value relations, with teacher assistants also supported.

  • Results

    MiniLM outperforms state-of-the-art baselines across student sizes; its 6-layer BERTBASE-derived model is 2.0× faster while retaining more than 99% performance on several tasks.

  • Takeaways & Limitations

    Last-layer attention distillation provides an effective and flexible way to compress monolingual and multilingual pre-trained Transformer language models.

  • Takeaways & Limitations

    The reported setup assumes a teacher with L layers and hidden size d_h, while the student has M layers and a different hidden size d′.

Abstract

from arXiv · show

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.

1. Introduction

Large pre-trained language models are effective but costly to fine-tune and serve. MiniLM addresses this by distilling deep self-attention knowledge into flexible, smaller students, achieving strong downstream performance with reduced computation.

  • Pre-trained language models contain hundreds of millions of parameters, creating latency and capacity challenges for fine-tuning and online serving.
  • Task-specific distillation first fine-tunes large language models on downstream tasks before transferring knowledge to smaller models.
  • Earlier task-agnostic methods transfer soft predictions, hidden states, or self-attention knowledge, often through layer-to-layer mappings with architectural constraints.
  • MiniLM distills the teacher’s last-layer self-attention, adding value relations to attention distributions while avoiding difficult layer mapping and supporting flexible student depth.
  • 2.0× faster inference and more than 99% performance are achieved by the 6-layer, 768-dimensional student distilled from BERTBASE.

2. Preliminary

Transformers encode contextualized token representations through stacked self-attention and feed-forward layers. Knowledge distillation trains a smaller student by minimizing differences from teacher features, including self-attention distributions and hidden representations.

  • Backbone Network: Transformer: Transformers encode token context with stacked layers, each combining self-attention, a feed-forward network, residual connections, and layer normalization.
  • Backbone Network: Transformer: BERT token inputs combine token, absolute position, and segment embeddings, with [SEP] separating segments and [CLS] representing the whole input.
  • Backbone Network: Transformer: Each Transformer layer transforms the previous hidden states into contextualized representations for the input tokens.
  • Self-Attention: Self-attention projects previous-layer outputs into queries, keys, and values, then computes attention distributions from scaled query-key dot products.
  • Transformer Distillation: Knowledge distillation trains a student to minimize a loss between student and teacher features over training data, commonly using MSE or KL-divergence.
  • Transformer Distillation: Transformer distillation can transfer masked-language-model probabilities, embeddings, self-attention distributions, and hidden states, with some methods requiring layer-count or dimensional alignment.

3. Deep Self-Attention Distillation

Deep self-attention distillation trains a task-agnostic student by mimicking the teacher’s last-layer self-attention, including attention distributions and value relations. A teacher assistant further supports distillation when the teacher–student size gap is large.

  • Deep Self-Attention Distillation: The method introduces value-relation transfer alongside existing attention-distribution transfer and uses a teacher assistant for large size gaps.The comparison section positions these choices against prior layer-to-layer approaches.
  • Self-Attention Distribution Transfer: The framework transfers the teacher’s last-layer attention distributions instead of mapping knowledge layer-to-layer.This avoids finding a layer mapping and allows more flexible student layer counts.
  • Self-Attention Value-Relation Transfer: Value relations add multi-head scaled dot-products between values to the attention-distribution transfer objective.The value-relation loss uses KL divergence, and the total loss sums attention-distribution and value-relation transfer losses.
  • Self-Attention Value-Relation Transfer: Scaled dot-products convert teacher and student value vectors into same-sized relation matrices, supporting flexible hidden dimensions without extra transformation parameters.This is intended to deepen the student’s mimicry of the teacher’s self-attention behavior.
  • Teacher Assistant: A teacher assistant is an intermediate-size student distilled from the teacher before guiding the final smaller student.The assistant bridges the teacher–student size gap and brings further improvements for smaller students.

4. Experiments

The experiments evaluate MINILM across downstream tasks and student architectures, showing strong performance from deep self-attention distillation. Value-relation transfer and teacher assistants further improve smaller students.

  • Main Results: MINILM exceeds DistillBERT and TinyBERT across most tasks, with gains of 3.0+% F1 on SQuAD 2.0 and 5.0+% accuracy on CoLA.
  • Main Results: 2.0× faster inference lets the 6-layer, 768-dimensional student retain more than 99% performance on SQuAD 2.0 and MNLI.
  • Smaller Students: MINILM outperforms soft-label distillation and TinyBERT on SQuAD 2.0, MNLI, and SST-2, including experiments with smaller students.
  • Smaller Students: Adding a teacher assistant further improves deep self-attention distillation, especially for smaller student models.
  • Ablation Studies: Value-relation transfer positively contributes to student results across different parameter sizes by helping students mimic teacher self-attention behavior.
  • Ablation Studies: The value-relation loss performs better than value-level MSE, improves SQuAD by about 1.0% F1, and requires no additional parameters.
  • Ablation Studies: Distilling only the teacher’s last Transformer layer achieves better results than layer-to-layer distillation while reducing computation and training time.

5. Discussion

MINILM extends its evaluation to larger in-house teachers, natural language generation, and multilingual models. The distilled models remain competitive or outperform baselines while using fewer parameters, and multilingual distillation is faster than soft-label distillation.

  • Better Teacher Better Student: The 12x384 MINILM achieves 2.7× speedup while performing competitively better than BERTBASE on SQuAD 2.0 and GLUE.
  • MINILM for NLG Tasks: MINILM’s 12x384 and 6x384 models achieve competitive performance on answer-aware question generation using SQuAD 1.1.
  • MINILM for NLG Tasks: The 12x384 model outperforms BERTSUMABS and MASSBASE on abstractive summarization with much fewer parameters, while the 6x384 model remains competitive.
  • Multilingual Models: MINILM’s multilingual distillation is much faster than soft-label distillation for multilingual pre-trained models.
  • Multilingual Models: Multilingual MINILM achieves competitive XNLI performance with much fewer Transformer parameters and compares favorably with mBERT and XLM.
  • Multilingual Models: The 12x384 multilingual model performs competitively better than mBERT and XLM on MLQA, while the 6-layer model also remains competitive.

6. Related Work

Prior work compresses pre-trained language models through task-specific or task-agnostic distillation using soft targets, hidden states, and self-attention features. MINILM instead transfers last-layer self-attention distributions and value relations to avoid layer-mapping constraints.

  • Existing Distillation Methods: Earlier methods transfer soft targets, embeddings, hidden states, or self-attention features, often through layer-to-layer matching.
  • Task-Specific Distillation: Task-specific distillation first fine-tunes large language models on downstream tasks, making the procedure costly and time-consuming for large datasets.
  • Task-Agnostic Distillation: Task-agnostic distillation trains a student to mimic the original pre-trained model so it can be directly fine-tuned on downstream tasks.
  • MINILM: MINILM transfers self-attention distributions and value relations from the teacher’s last Transformer layer instead of matching layers.
  • MINILM: Last-layer distillation avoids restrictions on student layer counts and reduces the effort of finding teacher–student layer mappings.

7. Conclusion

The paper proposes task-agnostic deep self-attention distillation for compressing pre-trained Transformers. It reports strong downstream accuracy and competitive multilingual results, while identifying compression of larger models as future work.

  • MINILM deeply mimics the teacher’s self-attention modules using last-layer attention distributions and value relations.
  • Teacher assistants further boost distillation performance, and the method retains high accuracy on SQuAD 2.0 and GLUE while outperforming state-of-the-art baselines.
  • Deep self-attention distillation also applies to multilingual pre-trained models and larger models, although compressing larger models remains future work.

A. GLUE Benchmark

This section points to tables summarizing the GLUE benchmark and SQuAD 2.0 datasets, with the latter presented through dataset statistics and metrics.

  • Table 14 presents a summary of the GLUE benchmark.
  • Table 15 presents dataset statistics and metrics for SQuAD 2.0.

B. Fine-tuning Hyper-parameters

The fine-tuning setup specifies task-dependent sequence lengths, batch sizes, learning-rate choices, epoch counts, and related training settings.

  • SQuAD 2.0: For SQuAD 2.0, sequences use a maximum length of 384 with a sliding window of 128 for longer inputs.
  • SQuAD 2.0: The 12-layer SQuAD 2.0 model uses batch size 48 and peak learning rate 4e-5 for 3 epochs, while other distilled models use batch size 32 and 6e-5.
  • GLUE: GLUE uses maximum sequence length 128, with BERTBASE students tuned over batch size 32, learning rates 2e-5 to 5e-5, and 3 to 5 epochs.
  • GLUE: Some GLUE tasks use up to 10 epochs, producing slight improvements.
  • Dataset references: Tables 14 and 15 provide dataset summaries for GLUE and SQuAD 2.0, respectively.

C. SQuAD 2.0

The supplied task configurations cover question generation, abstractive summarization, cross-lingual inference, and cross-lingual question answering.

  • Question Generation: Question generation uses batch size 32, total length 512, maximum output length 48, and beam size 5.Both models are fine-tuned for 25 epochs with label smoothing rate 0.1 and length penalty 1.3.
  • Abstractive Summarization: Abstractive summarization uses batch size 64 and label smoothing rate 0.1.For CNN/DailyMail, total length is 768 and maximum output length is 160; decoding uses beam size 5 and length penalty 0.7.
  • Cross-lingual Natural Language Inference: Cross-lingual natural language inference on XNLI uses maximum sequence length 128, batch size 128, 5 epochs, and learning rates from 3e-5 to 6e-5.
  • Cross-lingual Question Answering: Cross-lingual question answering on MLQA uses maximum sequence length 512 and a sliding window.
Loading 2002.10957v2…