[Paper Review] Extract then Distill: Efficient and Effective Task-Agnostic BERT Distillation

Cheng Chen, Yichun Yin+5 more|arXiv (Cornell University)|Apr 24, 2021

Topic Modeling30 references4 citations

TL;DR

This paper proposes Extract Then Distill (ETD), a flexible and efficient method for task-agnostic BERT distillation that reuses pre-trained teacher model parameters to initialize smaller student models. By extracting and reusing the most important neurons (FFN, attention heads, hidden states) via importance scores or random selection, ETD reduces training cost by 70% while improving performance over baseline distillation methods.

ABSTRACT

Task-agnostic knowledge distillation, a teacher-student framework, has been proved effective for BERT compression. Although achieving promising results on NLP tasks, it requires enormous computational resources. In this paper, we propose Extract Then Distill (ETD), a generic and flexible strategy to reuse the teacher's parameters for efficient and effective task-agnostic distillation, which can be applied to students of any size. Specifically, we introduce two variants of ETD, ETD-Rand and ETD-Impt, which extract the teacher's parameters in a random manner and by following an importance metric respectively. In this way, the student has already acquired some knowledge at the beginning of the distillation process, which makes the distillation process converge faster. We demonstrate the effectiveness of ETD on the GLUE benchmark and SQuAD. The experimental results show that: (1) compared with the baseline without an ETD strategy, ETD can save 70\% of computation cost. Moreover, it achieves better results than the baseline when using the same computing resource. (2) ETD is generic and has been proven effective for different distillation methods (e.g., TinyBERT and MiniLM) and students of different sizes. The source code will be publicly available upon publication.

Motivation & Objective

• To address the high computational cost of task-agnostic BERT distillation, which requires extensive training on large-scale unsupervised corpora.
• To enable efficient and effective distillation for student models of any size, including those with different widths and depths than the teacher.
• To develop a generic initialization strategy that reuses teacher parameters to accelerate convergence and improve performance.
• To demonstrate that parameter reuse via pruning-based extraction is more effective than random initialization or naive layer copying.

Proposed method

• ETD introduces two variants: ETD-Rand (random neuron extraction) and ETD-Impt (importance-based neuron extraction) for width-wise parameter reuse.
• The method performs uniform layer selection across the teacher model’s Transformer layers to ensure broad knowledge transfer.
• Importance scores are computed for FFN neurons, attention heads, and hidden states to guide selective parameter extraction.
• The extracted parameters are used to initialize the student model before fine-tuning, reducing the need for extensive distillation training.
• The approach is compatible with existing distillation frameworks such as TinyBERT and MiniLM.
• The extraction process is efficient and incurs negligible computational overhead.

Experimental results

Research questions

• RQ1 Can reusing pre-trained teacher parameters through structured extraction improve the efficiency of task-agnostic BERT distillation?
• RQ2 Does importance-based parameter selection outperform random or uniform selection in initializing student models?
• RQ3 Can ETD be applied to student models with arbitrary widths and depths, not just those matching the teacher?
• RQ4 To what extent does ETD reduce computational cost while maintaining or improving model performance?
• RQ5 How does the choice of layer selection strategy (uniform, top, bottom) affect distillation convergence and final performance?

Key findings

• ETD reduces the computational cost of distillation by 70% compared to the baseline, without sacrificing performance.
• When using the same computational resources, ETD achieves better results than baseline distillation on both GLUE and SQuAD benchmarks.
• ETD-Impt, which uses importance scores for parameter selection, consistently outperforms ETD-Rand and random initialization.
• The ablation study confirms that even the least important neurons contribute positively, and reversing the importance ranking degrades performance.
• Uniform layer selection is superior to top- or bottom-layer selection, as it ensures balanced knowledge transfer from shallow to deep layers.
• The method is general and effective across different distillation frameworks, including TinyBERT and MiniLM, and works with student models of varying sizes.