Pea-KD: Parameter-efficient and accurate Knowledge Distillation

IKHYUN CHO, U Kang

Keywords: BERT, Deep Learning, Natural Language Processing, Transformer, Knowledge Distillation, Parameter Sharing

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, the existing KD methods suffer from the following limitations. First, since the student model is small in absolute size, it inherently lacks model complexity. 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 allows greater model complexity for the student model. 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 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.

One-sentence Summary: It introduces a new Knowledge Distillation method. It improves the performance of the student model with 2 main modules: novel parameter sharing and new pretraining which uses the teacher model's predictions.

Code Of Ethics: I acknowledge that I and all co-authors of this work have read and commit to adhering to the ICLR Code of Ethics

Community Implementations: [ 2 code implementations](https://www.catalyzex.com/paper/arxiv:2009.14822/code)

Reviewed Version (pdf): https://openreview.net/references/pdf?id=pgcw7UOeGy