Universal- KD : Attention-based Output-Grounded Intermediate Layer Knowledge Distillation (2021)

Abstract: Intermediate layer matching is shown as an effective approach for improving knowledge distillation (KD). However, this technique applies matching in the hidden spaces of two different networks (i.e. student and teacher), which lacks clear interpretability. Moreover, intermediate layer KD cannot easily deal with other problems such as layer mapping search and architecture mismatch (i.e. it requires the teacher and student to be of the same model type). To tackle the aforementioned problems all together, we propose Universal-KD to match intermediate layers of the teacher and the student in the output space (by adding pseudo classifiers on intermediate layers) via the attention-based layer projection. By doing this, our unified approach has three merits: (i) it can be flexibly combined with current intermediate layer distillation techniques to improve their results (ii) the pseudo classifiers of the teacher can be deployed instead of extra expensive teacher assistant networks to address the capacity gap problem in KD which is a common issue when the gap between the size of the teacher and student networks becomes too large; (iii) it can be used in cross-architecture intermediate layer KD. We did comprehensive experiments in distilling BERT-base into BERT-4, RoBERTa-large into DistilRoBERTa and BERT-base into CNN and LSTM-based models. Results on the GLUE tasks show that our approach is able to outperform other KD techniques.

## introduction despite the great success of deep neural networks in different tasks such as computer vision @xcite and natural language processing (nlp) @xcite @xcite due to huge over-parameterization, when it comes to deploying these models for the end user or on edge devices with limited memory and computational power, this over-parameterization can be very prohibitive. therefore, different neural model compression techniques such as quantization @xcite , pruning @xcite , layer decomposition (mondelli and montanari, 2019) and knowledge distillation (kd) @xcite aim at reducing the number of parameters of the models, improving their memory requirements or running efficiency. kd is one of the most well-known neural model compression techniques. kd provides a way to compress a large model (so-called the teacher model) into a small model (i.e. the student model). kd has shown being reliable in reducing the number of parameters and computations while achieving competitive results on downstream tasks. recently, kd has attracted more attention in the nlp field especially due to large pre-trained language models (plms) @xcite @xcite . however, it is evident that the original kd is not performing very well in maintaining the performance of compressed plms and it is required to be equipped with other auxiliary training objectives @xcite . @xcite refers to a few of these auxiliary techniques such as intermediate layer distillation (ild) and data augmentation. mo-bilebert @xcite deploys a progressive layer-wise training. @xcite proposes a two-stage gradual training by using a dynamic temperature factor during training to address the common capacity gap problem (i.e. training with kd becomes more difficult when the capacity gap between the teacher and student networks become too large) in kd. @xcite and @xcite show the importance of ild and try to improve the search and skip problems in ild, respectively. among all auxiliary techniques, the focus of this paper is on ild as a promising approach in improving kd. existing ild techniques are established on raw intermediate representations. we argue that the physical interpretation of these raw representations are not known. currently in most of ild techniques, distilling these representations from one network to another is done arbitrarily and without any particular reasoning. performing ild in this form leads to other weaknesses such as requiring layer mapping search and requiring the architecture of the two networks to be the same. to address these limitations, we provide a universal solution which maps the raw hidden representations of each layer to the output space using so-called intermediate pseudo classifiers. we demonstrate in this work that our universal-kd technique leads to the following contributions: ## related work recent years have seen a wide array of methods that leverage intermediate layer matching @xcite @xcite , data augmentation @xcite @xcite , adversarial training @xcite @xcite , lately loss terms re-weighting @xcite @xcite in order to reduce the teacher-student performance gap. in this section, we introduce the related work to ild, the capacity gap problem, and intermediate layer crossarchitecture distillation. you can find a comprehensive comparison of universal-kd and related works in ## methodology methodology of our work concerns introducing our universal-kd solution and showing how this technique can address the main discussed problems in kd which are interpretability of ild, the capacity gap issue and intermediate layer cross-architecture distillation.