Pea-KD: Parameter-efficient and Accurate Knowledge Distillation on BERT

TL;DR

Pea-KD introduces a novel knowledge distillation method for BERT that enhances student model capacity and initialization, leading to significant performance improvements across multiple NLP tasks.

Contribution

The paper proposes Pea-KD, combining Shuffled Parameter Sharing and Teacher's Predictions for better model compression and performance in knowledge distillation.

Findings

Improves BERT student performance by 4.4% on average across GLUE tasks.

Outperforms existing KD baselines significantly.

Enhances student model capacity and initialization effectively.

Abstract

How can we efficiently compress a model while maintaining its performance? Knowledge Distillation (KD) is one of the widely known methods for model compression. In essence, KD trains a smaller student model based on a larger teacher model and tries to retain the teacher model's level of performance as much as possible. However, existing KD methods suffer from the following limitations. First, since the student model is smaller in absolute size, it inherently lacks model capacity. Second, the absence of an initial guide for the student model makes it difficult for the student to imitate the teacher model to its fullest. Conventional KD methods yield low performance due to these limitations. In this paper, we propose Pea-KD (Parameter-efficient and accurate Knowledge Distillation), a novel approach to KD. Pea-KD consists of two main parts: Shuffled Parameter Sharing (SPS) and Pretraining with Teacher's Predictions (PTP). Using this combination, we are capable of alleviating the KD's limitations. SPS is a new parameter sharing method that increases the student model capacity. PTP is a KD-specialized initialization method, which can act as a good initial guide for the student. When combined, this method yields a significant increase in student model's performance. Experiments conducted on BERT with different datasets and tasks show that the proposed approach improves the student model's performance by 4.4\% on average in four GLUE tasks, outperforming existing KD baselines by significant margins.