MLKD-BERT: Multi-level Knowledge Distillation for Pre-trained Language Models
Ying Zhang, Ziheng Yang, Shufan Ji
TL;DR
This work addresses the challenge of deploying compact yet capable PLMs by introducing MLKD-BERT, a two-stage, multi-level knowledge distillation framework for BERT compression. It transfers embedding-layer and Transformer-layer knowledge in Stage 1 and prediction-layer knowledge in Stage 2, incorporating relation-level cues via embedding-token similarity, self-attention relations with MHA-splits, and sample-level relations including similarity and contrastive signals, formalized through losses $ abla ext{Stage,1}$ and $ abla ext{Stage,2}$ where $ abla ext{Stage,1} = abla ext{EMB} + abla ext{MHA} + abla ext{FFN}$ and $ abla ext{Stage,2} = abla ext{SS} + abla ext{SC} + abla ext{KD}$. Empirically, MLKD-BERT outperforms existing BERT distillation methods on GLUE and SQuAD, while allowing smaller student heads to achieve substantial inference-time speedups with minimal accuracy loss. The method demonstrates that combining feature- and relation-level knowledge with flexible architectural constraints yields efficient, high-performing PLMs suitable for resource-constrained deployment. This approach thus provides a practical pathway to scalable, transformer-based NLP systems beyond vanilla distillation techniques, with potential applicability to other PLMs beyond BERT.
Abstract
Knowledge distillation is an effective technique for pre-trained language model compression. Although existing knowledge distillation methods perform well for the most typical model BERT, they could be further improved in two aspects: the relation-level knowledge could be further explored to improve model performance; and the setting of student attention head number could be more flexible to decrease inference time. Therefore, we are motivated to propose a novel knowledge distillation method MLKD-BERT to distill multi-level knowledge in teacher-student framework. Extensive experiments on GLUE benchmark and extractive question answering tasks demonstrate that our method outperforms state-of-the-art knowledge distillation methods on BERT. In addition, MLKD-BERT can flexibly set student attention head number, allowing for substantial inference time decrease with little performance drop.