RAIL - KD : RA ndom Intermediate Layer Mapping for Knowledge Distillation (2022)

Abstract: Intermediate layer knowledge distillation (KD) can improve the standard KD technique (which only targets the output of teacher and student models) especially over large pre-trained language models. However, intermediate layer distillation suffers from excessive computational burdens and engineering efforts required for setting up a proper layer mapping. To address these problems, we propose a RAndom Intermediate Layer Knowledge Distillation (RAIL-KD) approach in which, intermediate layers from the teacher model are selected randomly to be distilled into the intermediate layers of the student model. This randomized selection enforces that all teacher layers are taken into account in the training process, while reducing the computational cost of intermediate layer distillation. Also, we show that it acts as a regularizer for improving the generalizability of the student model. We perform extensive experiments on GLUE tasks as well as on out-of-domain test sets. We show that our proposed RAIL-KD approach outperforms other state-of-the-art intermediate layer KD methods considerably in both performance and training-time.

## introduction pre-trained language models (plms), such as bert @xcite , roberta @xcite and xlnet @xcite have shown remarkable abilities to match and even surpass human performances on many natural languages understanding (nlu) tasks @xcite @xcite . however, the deployment of these models in real world applications (e.g. edge devices) come with challenges, mainly due to large model size and inference time. in this regard, several model compression techniques such as quantization @xcite @xcite , pruning @xcite @xcite ,matrix factorization @xcite , optimizing the transformer architecture @xcite @xcite , and knowledge distillation @xcite @xcite @xcite @xcite have been developed to reduce the model size and latency, while maintaining comparable performance to the original model. kd, which is the main focus of this work, is a neural model compression approach that involves training a small student model with the guidance of a large pre-trained teacher model. in the original kd technique @xcite @xcite , the teacher output predictions are used as soft labels for supervising the training of the student. there has been several attempts in the literature to reduce the teacher-student performance gap by leveraging data augmentation @xcite @xcite @xcite , adversarial training @xcite @xcite , and intermediate layer distillation (ild) @xcite @xcite . when it comes to bert compression, ild leads to clear gains in performances @xcite @xcite due to its ability to enhance the knowledge transfer beyond logits matching. this is done by mapping intermediate layer representations of both models to a common space foot_0 , and then matching them via regression @xcite or cosine similarity @xcite losses. one major problem with ild is the absence of an appropriate strategy to select layers to be matched on both sides, reacting to the skip ## related work in recent years, a wide range of methods have tried to expand knowledge transfer of transformerbased @xcite nlu models beyond logits matching. distillbert @xcite added a cosine similarity loss between teacher and student embeddings layers. tinybert @xcite , mobilebert @xcite , and minilm @xcite ) matched the intermediate layers representations and self-attention distributions of the teacher and the student. in pkd, @xcite used deterministic mapping strategies to distill a 12-layer bert teacher to a 6-layer bert student. pkd-last and pkd-skip refer to matching layers {1 -5} of the student with layers {7 -11} and {2, 4, 6, 8, 10} of the teacher respectively. however, these works ignored the impact of layer selection, as they used a fixed layer-wise mapping. 2 researchers have found that tuning the layer mapping scheme can significantly improve the performance of ild techniques @xcite . nevertheless, finding the optimal mapping can be challenging, which is referred to as the layer skip and search problems by @xcite . to address the layer skip problem, ckd @xcite is built on top of pkd by partitioning all the intermediate layers of the teacher to the number of student layers. then, the combined representation of the layers of each partition is distilled into a number of subset corresponding to the number of student layers. however, finding the optimal partitioning scheme requires running exhaustive experiments. given teacher and student bert models with n and m layers respectively (where @xmath0, it is not trivial to choose the teacher layers that can be incorporated in the distillation process and how we should map them to the student layers (search). alp-kd @xcite overcomes this issue by computing attention weights between each student layer and all the intermediate layers of the teacher. the learned attention weights for each student layer are used to obtain a weighted representation of all teacher layers. although alp-kd has shown promising results on 12-layer bert-based compression, attending to all layers of the teacher adds considerable computational overhead to the training phase. this can become computationally prohibitive when scaling to very large models such as roberta-large @xcite or @xcite . alternatively, codir @xcite exploited contrastive learning @xcite to perform intermediate layers matching between the teacher and the student models with no deterministic mapping. similar to alp-kd, this approach also requires excessive training time due to the contrastive loss calculation and the use of negative samples from a memory bank. our proposed rail-kd method does not add any computational cost to the distillation process, while empirically outperforming previous methods. for instance, rail-kd is roughly two-times faster than codir in a 24 to 6 layer compression. in addition, it does not require extensive experiments to find the optimal mapping scheme. in this work, we position ourselves to works that tackle the skip and search problem. 3 in other words, we don't compare with works like tinybert @xcite or minilm @xcite , which use extra losses like self-attention distribution matching. however, we expect that these methods, as well as state-of-the-art @xcite ones can take full advantage of rail-kd, since they use a deterministic layer mapping scheme. ## rail-kd the rail-kd method is sketched in figure 1 . rail-kd transfers intermediate knowledge of a pre-trained teacher t with n intermediate layers to a student model s θ with m intermediate layers. in contrast to traditional intermediate layer distillation techniques which keep the selected layers of the teacher for distillation fixed during training, in rail-kd, at each epoch, a few intermediate layers from the teacher model are selected randomly for distillation. here for simplicity, we set the number of selected intermediate layers of the teacher model equal to that of the student model. let (x, y) denote a training sample @xmath1, • • • , x l-1 ) which is a sequence of l tokens and y its corresponding label. in figure 1 , our random selection operator is applied to the intermediate layers of the teacher to randomly select m out of n layers. the intermediate layer representations of the m selected layers of the teacher and the student model corresponding to the x input can be described as here, d 1 and d 2 indicate the hidden dimension of the layers of the teacher and the student models respectively. to obtain h t i,x and h s j,x , we need to find the individual representation of each token x k at each layer i, which is indicated as h t i,x k and h s i,x k for the teacher and student networks respectively. at this stage, we need to obtain an aggregated representation of the sequence x at each layer. in this regard, one can either use @xmath2. since in @xcite , the mean-pooling representation shows better results, we adopt it to compute the sentence representation of each layer. meanpooling is a row-wise average over @xcite : after obtaining aggregated layer representations for both the student and teacher networks, our rail-kd proposal is to randomly select m layer representations from the teacher through training to perform the intermediate layer distillation (ild). ## results results of codir are directly copied from their paper @xcite . first, we observe that the performance gap between ild methods and vanilla-kd is tight (0.8% and 0.3% on dev and test sets respectively). moreover, as we expect, alp-kd performs better on dev and similar on test compared to pkd with 0.2% improvement on the dev results. second, results show that rail-kd outperforms the best ild methods by a margin of 0.5% and 0.3% on average on dev and test sets respectively. we notice that, except on rte test, our rail-kd variants obtained the highest per-task performances. third, we observe that rail-kd variants perform very similarly, which indicates that our method is effective on concatenated as well as layer-wise distillation. similar trends are seen on the 24-to 6-layer model compression experiments, which are reported in ## analysis we run extensive analysis to better understand why rail-kd performs better than the other baselines. we visualize the layer-wise cosine similarity between the intermediate representations of the teacher and the student networks. figure 2 shows the cosine similarity score between three intermediate layer representations of bert 12 teacher (i.e. layers 2, 4 and 6) and the first three layer representations of the student for pkd, alp-kd, rail-kd l/c students on 100 samples randomly selected from the sst-2 dataset. due to space constraints, we only plot the scores for the first three layers of the student model. similar trends are seen from the other layers. we found that rail-kd allows the student to mimic teacher layers similar to pkd and much better than alp-kd, despite that the mapping scheme varies at each epoch. moreover, we observe that alp-kd results have less similarity scores in the upper intermediate layers. pkd gives lower similarity scores in the lower layers while improving in the upper layers. in contrast, our approach gives more stable similarity scores for all layers while getting closer to the teacher representation in the upper layers. we further investigate the attention weights learned by alp-kd, and find out that they mostly focus on few layers (sparse attention). figure 3 illustrates the distribution of weights, averaged on all training samples of distilbert 6 alp-kd studnet on cola (left), rte (middle), and mrpc (right) 5 . the figure clearly shows (light colors) that most of alp weights are concentrated on top layers of the teacher. for instance, layers 1,2,5 of the three students mostly attend to the last layer of bert 12 . this is an indicator that alp-kd overfits to the information driven from last layers. in contrast, the randomness in layer selection of rail-kd ensures a uniform focus on teacher layers. this may explain the poor performance of alp-kd on outof-domain evaluation compared with rail-kd. from figure 3 , we see clearly that alp-kd mostly prefers the upper layers of the teacher. on the other hand, the deterministic nature of pdk allows it to match better particular layers of the teacher (e.g. bottom ones as shown in figure 2 ), but pkd never sees the layers that are skipped by the mapping. consequently, it is expected that even though pdk can mimic bottom layers well, it is worse overall because it completely ignores some layers of the teacher. random layer selection allow rail-kd to mimic all teacher layers while delivering high performances. 5 similar trends found on other datasets. ## conclusion and future work we introduced a novel, simple, and efficient intermediate layer kd approach that outperforms the conventional approaches with performance improvement and efficient training time. rail-kd selects random intermediate layers from the teacher which equals to the number of intermediate layers of the student model. the selected intermediate layers are then sorted to distill their representations into the student model. rail-kd yields better regularization, which helps performance. furthermore, our approach shows better performance for larger model distillation with faster training time, which opens up an avenue to investigate our approach for super-large models.