Improving Task-Agnostic BERT Distillation with Layer Mapping Search

Abstract

Knowledge distillation (KD) which transfers the knowledge from a large teacher model to a small student model, has been widely used to compress the BERT model recently. Besides the supervision in the output in the original KD, recent works show that layer-level supervision is crucial to the performance of the student BERT model. However, previous works designed the layer mapping strategy heuristically (e.g., uniform or last-layer), which can lead to inferior performance. In this paper, we propose to use the genetic algorithm (GA) to search for the optimal layer mapping automatically. To accelerate the search process, we further propose a proxy setting where a small portion of the training corpus are sampled for distillation, and three representative tasks are chosen for evaluation. After obtaining the optimal layer mapping, we perform the task-agnostic BERT distillation with it on the whole corpus to build a compact student model, which can be directly fine-tuned on downstream tasks. Comprehensive experiments on the evaluation benchmarks demonstrate that 1) layer mapping strategy has a significant effect on task-agnostic BERT distillation and different layer mappings can result in quite different performances; 2) the optimal layer mapping strategy from the proposed search process consistently outperforms the other heuristic ones; 3) with the optimal layer mapping, our student model achieves state-of-the-art performance on the GLUE tasks.

References in corpus (16)

• Distilling the Knowledge in a Neural Network
• Exploring the Limits of Transfer Learning with a Unified Text-to-Text Transformer
• FitNets: Hints for Thin Deep Nets
• MiniLM: Deep Self-Attention Distillation for Task-Agnostic Compression of Pre-Trained Transformers
• Q8BERT: Quantized 8Bit BERT
• Well-Read Students Learn Better: On the Importance of Pre-training Compact Models
• Universal Transformers
• Distilling Task-Specific Knowledge from BERT into Simple Neural Networks
• The Evolved Transformer
• A Tensorized Transformer for Language Modeling
• MobileBERT: a Compact Task-Agnostic BERT for Resource-Limited Devices
• Structured Pruning of a BERT-based Question Answering Model
• FastBERT: a Self-distilling BERT with Adaptive Inference Time
• Evolving Normalization-Activation Layers
• Q-BERT: Hessian Based Ultra Low Precision Quantization of BERT
• Know What You Don't Need: Single-Shot Meta-Pruning for Attention Heads

Cited by in corpus (2)

• Extract then Distill: Efficient and Effective Task-Agnostic BERT Distillation
• AutoTinyBERT: Automatic Hyper-parameter Optimization for Efficient Pre-trained Language Models