Improved knowledge distillation by utilizing backward pass knowledge in neural networks

Aref Jafari, Mehdi Rezagholizadeh, Ali Ghodsi

Keywords: Deep Neural Networks, Knowledge Distillation, Natural Language Processing

Abstract:

Knowledge distillation (KD) is one of the prominent techniques for model compression. In this method, the knowledge
of a large network (teacher) is distilled into a model (student) with usually significantly fewer parameters. KD tries to
better-match the output of the student model to that of the teacher model based on the knowledge extracts from
the forward pass of the teacher network. Although conventional KD is effective for matching the two networks over the given
data points, there is no guarantee that these models would match in other areas for which we do not have enough
training samples. In this work, we address that problem by generating new auxiliary training samples based on
extracting knowledge from the backward pass of the teacher in the areas where the student diverges greatly from the
teacher. We compute the difference between the teacher and the student and generate new data samples that
maximize the divergence. This is done by perturbing data samples in the direction of the gradient of the difference
between the student and the teacher. Augmenting the training set by adding this auxiliary improves the performance of
KD significantly and leads to a closer match between the student and the teacher. Using this approach, when data samples
come from a discrete domain, such as applications of natural language processing (NLP) and language understanding,
is not trivial. However, we show how this technique can be used successfully in such applications. We studied the effect
of the proposed method on various tasks in different domains, including images and NLP tasks with considerably
smaller student networks. The results of our experiments, when compared with the original KD, show 4% improvement on MNIST with a student network that is 160 times smaller, 1% improvement on a CIFAR-10 dataset with
a student that is 9 times smaller, and an average 1.5% improvement on the GLUE benchmark with a distilroBERTa-base
student.

One-sentence Summary: This paper, presents a new idea to improve the performance of existing knowledge distillation methods by enriching the train data samples

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

Reviewed Version (pdf): https://openreview.net/references/pdf?id=I-d0x7Dc-g