RAIL-KD: RAndom Intermediate Layer Mapping for Knowledge Distillation

Introduction

Pre-trained Language Models (PLMs), such as BERT (Devlin et al. 2019), RoBERTa (Liu et al. 2020) and XLNet (Yang et al. 2019) have shown remarkable abilities to match and even surpass human performances on many Natural Languages Understanding (NLU) tasks Rajpurkar et al. 2018; Wang et al. 2018; Wang et al. 2019. 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 (Shen et al. 2019; Zafrir et al. 2019), pruning (Guo et al. 2019; Gordon et al. 2020; Michel et al. 2019), optimizing the Transformer architecture (Fan et al. 2019; Ghaddar and Langlais 2019; Wu et al. 2020b; Lu et al. 2020), and knowledge distillation (Sanh et al. 2019a; Jiao et al. 2019; Sun et al. 2020b; Wang et al. 2020a; Rashid et al. 2021; Passban et al. 2021; Jafari et al. 2021; Kamalloo et al. 2021) 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 Buciluǎ et al. 2006; Hinton et al. 2014; Turc et al. 2019, 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 Fu et al. 2020; Li et al. 2021; Jiao et al. 2019, adversarial training Zaharia et al. 2021; Rashid et al. 2020; Rashid et al. 2021, and intermediate layer distillation (ILD) Wang et al. 2020b; Wang et al. 2020a; Ji et al. 2021; Passban et al. 2021.

When it comes to Bert compression, ILD leads to clear gains in performances Sanh et al. 2019a; Jiao et al. 2019; Wang et al. 2020a 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 In some cases, the representations are directly matched if the teacher and student have the same hidden size., and then matching them via regression Sun et al. 2019 or cosine similarity Sanh et al. 2019a losses. On major problem with ILD is the absence of an appropriate strategy to select layers to be matched on both sides, reacting to the skip and search problem (Passban et al. 2021). There are some solutions in the literature mostly rely on layer combination (Wu et al. 2020a), attention-based layer projection (Passban et al. 2021) and contrastive learning (Sun et al. 2020a).

While these solutions are all effective to some extent, to the best of our knowledge, there is no work in the literature doing a comprehensive evaluation of these techniques in terms of both efficiency and performance. A case in point is that the aforementioned solutions to the layer skip and search problem are not scalable to very deep networks. We propose RAIL-KD (RAndom Intermediate Layer KD), a simple yet effective method for intermediate layer mapping which randomly selects $k$ out of $n$ intermediate layers of the teacher at each epoch to be distilled to the corresponding student layers. Since the layer selection is done randomly, all the intermediate layers of the teacher will have a chance to be selected for distillation. Our method adds no computational cost to the training, still outperforming all aforementioned methods on the GLUE benchmark Wang et al. 2018. Moreover, we observe larger gains distilling larger teacher models, as well as when compressed models are evaluated on out-of-domain datasets. Last, we report the results on 5 random seeds in order to verify the contribution of the random selection process, thus making the comparison fair with previous methods. The main contributions of our paper are as follows:

We introduce RAIL-KD, a more efficient and scalable intermediate layer distillation approach.

To the best of our knowledge, we are the first to perform a comprehensive study of the intermediate layer distillation techniques in terms of both efficiency and performance.

We consider the distillation of models such as Bert and RoBERTa, and compare different up-to-date distillation techniques on out-of-domain test sets. Thus providing new points of comparison.

Related Work

Recent years, have seen a wide range of methods have emerged aiming to expand knowledge transfer of transformer-based Vaswani et al. 2017 NLU models beyond logits matching. DistillBert Sanh et al. 2019a added a cosine similarity loss between teacher and student embeddings layer. TinyBERT Jiao et al. 2019, MobileBERT Sun et al. 2020b, and MiniLM Wang et al. 2020b matched the intermediate layers representations and self-attention distributions of the teacher and the student.

In PKD, Sun et al. 2019 used deterministic mapping strategies to distill a 12-layers Bert teacher to a 6-layers 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. e.g. matching the first (or last) $k$ layers of the student with their corresponding teacher layers.

Researchers have found that tuning the layer mapping scheme can significantly improve the performance of ILD techniques Sun et al. 2019. Nevertheless, finding the optimal mapping can be challenging, which is referred to as the layer skip and search problems by Passban et al. 2021. To address the layer skip problem, CKD (Wu et al. 2020a) 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 $n>>m$), 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 Passban et al. 2021 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 (Liu et al. 2020) or GPT-2 (Radford et al. 2019). Alternatively, CODIR Sun et al. 2020a exploited contrastive learning (Tian et al. 2019) 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.

Table 1 summarizes the main characteristics of the existing state-of-the-art intermediate layer distillation techniques used for pre-trained language models compared with our proposed RAIL-KD. As it is shown in this table, PKD Sun et al. 2019, CKD Wu et al. 2020a, and CoDIR are the most related works to us. However, PKD and CKD treat the mapping as an extra hyperparameter that requires extensive experiments to find the optimal mapping. On the other hand, ALP-KD Passban et al. 2021 and CoDIR Sun et al. 2020a use attention mechanism and contrastive learning respectively to address the issue, but at the expense of extra computational cost.

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 twice faster than CoDIR in a 24 to 6 layers 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 Only work that performs intermediate layer distillation. Otherwise said, we don’t compare with works like TinyBERT Jiao et al. 2019 or MiniLM Wang et al. 2020b, which use extra losses like self-attention distribution matching. However, we expect that these methods, as well as state-of-the-art Rashid et al. 2021; He et al. 2021 one can take full advantage of RAIL-KD, since they use deterministic layer mapping scheme.

RAIL-KD

The RAIL-KD method is sketched in Figure 1. 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. Our method is architecture agnostic can be applied to different domains such as computer vision and NLP.

RAIL-KD transfers intermediate knowledge of a pre-trained teacher $T$ with $n$ intermediate layers to a student model $S_{\theta}$ with $m$ intermediate layers. Let $(X,y)$ denotes a training sample $X=(x_{0},\cdots,x_{L-1})$ which is a sequence of $L$ (sub-)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 $H_{X}^{T}=\{H_{1,X}^{T},\cdots,H_{m,X}^{T}\}$ and $H_{X}^{S_{\theta}}=\{H_{1,X}^{S_{\theta}},\cdots,H_{m,X}^{S_{\theta}}\}$ respectively, where $H_{i,X}^{T}=\cup_{k=0}^{L-1}\{H_{i,x_{k}}^{T}\}\in R^{L\times d_{1}}$ and $H_{i,X}^{S_{\theta}}=\cup_{k=0}^{L-1}\{H_{i,x_{k}}^{S_{\theta}}\}\in R^{L\times d_{2}}$.

Here, $d_{1}$ and $d_{2}$ indicate the hidden dimension of the layers of the teacher and the student models respectively. To obtain $H_{i,X}^{T}$ and $H_{j,X}^{S}$, we need to find an aggregated representation for the sequence of $L$ tokens at each layer of the two networks. In this regard, one can either use the $<$CLS$>$ token representation or use the mean-pooling of the sequence representations of the layer. Since in (Sun et al. 2020a), the mean-pooling representation shows better results, we adopt it to compute the sentence representation of each layer. Mean-pooling is a row-wise average over $H_{i,X}^{T}$, $H_{i,X}^{S_{\theta}}$ to get $\bar{h}_{i,X}^{T}\in R^{d_{1}}$ $\bar{h}_{i,X}^{S_{\theta}}\in R^{d_{2}}$ (Sun et al. 2020a):

RAIL-KD proposes the intermediate layer distillation in two different forms: using layer-wise distillation (see Fig. 1(a)) or by concatenating layer representations (see Fig. 1(b)).

In this setting, the representations $\bar{h}_{i,X}^{T}\in R^{d_{1}}$ and $\bar{h}_{i,X}^{S_{\theta}}\in R^{d_{2}}$ are projected into the same lower-dimensional space $\hat{h}_{i,X}^{T},\hat{h}_{i,X}^{S_{\theta}}\in R^{u}$ using $(d_{1}\times u)$ and $(d_{2}\times u)$ linear mappings respectively to calculate the layer-wise losses.

where $\mathcal{X}$ denotes the set of training samples, and $\alpha_{i}$ is a hyper-parameter to weigh the layer-wise distillation loss.

2 Concatenated RAIL-KD

In this setting, intermediate layer representations are concatenated and then distilled: $\bar{h}_{X}^{T}=[\bar{h}_{1,X}^{T},\cdots,\bar{h}_{m,X}^{T}]$, $\bar{h}_{X}^{S_{\theta}}=[\bar{h}_{1,X}^{S_{\theta}},\cdots,\bar{h}_{m,X}^{S_{\theta}}]$ which are further mapped into the same lower-dimensional space $\hat{h}_{X}^{T},\hat{h}_{X}^{S_{\theta}}\in R^{u}$ using $(md_{1}\times u)$ and $(md_{2}\times u)$ linear mappings to calculate the concatenated distillation loss.

Any type of loss such as contrastive (Sun et al. 2020a), or mean-square-error (MSE) (Passban et al. 2021; Sun et al. 2019) can be applied for our RAIL-KD approach.

3 Training Loss

The intermediate representation distillation loss $L_{\text{RAIL-KD}}$ is combined with the original KD loss $L_{\text{KD}}$, which is used to distill the knowledge from the output logits of the teacher model $T$ to the output logits of the student model $S_{\theta}$, and the original cross-entropy loss $L_{\text{CE}}$. The total loss function for training the student model is:

where $\lambda_{1}$, $\lambda_{2}$, and $\lambda_{3}$ are hyper-parameters of our model to minimize the total loss, and $\lambda_{1}+\lambda_{2}+\lambda_{3}=1$.

Experimental Protocol

We evaluate RAIL-KD on 8 tasks from the GLUE benchmark Wang et al. 2018: 2 single-sentence (CoLA and SST-2) and 5 sentence-pair (MRPC, RTE, QQP, QNLI, and MNLI) classification tasks, and 1 regression task (STS-B). Following prior works (Sun et al. 2019; Passban et al. 2021; Jiao et al. 2019; Sun et al. 2020a), we use the same metrics as the GLUE benchmark for evaluation. Moreover, to further show the generalization capability of our RAIL-KD method on out-of-domain (OOD) across tasks, we use Scitail (Khot et al. 2018), PAWS (Paraphrase Adversaries from Word Scrambling) (Zhang et al. 2019), and IMDb (Internet Movie Database) (Maas et al. 2011) test sets to evaluate the models fine-tuned on MNLI, QQP, and SST-2 tasks respectively.

2 Implementation Details

We run extensive experiments on 3 different teachers in order to ensure a fair comparison with of a wide range of prior works, and also to show the effectiveness of RAIL-KD. We experiment with the 12 layers Bert-base-uncased Devlin et al. 2019 as teacher (Bert12) and 6 layer DistilBert Sanh et al. 2019a as student (DistillBert6) to compare with PKD Sun et al. 2019 and ALP-KD Passban et al. 2021. Also, we use 24 layers RoBERTa-large (Liu et al. 2020) and 6 layers DistilRoberta (Sanh et al. 2019b) as the backbone for teacher (RoBERTa24) and student (DistilRoberta6) respectively to compare models when $n>>m$. Furthermore, we perform evaluation using 12 layers RoBERTa-base (RoBERTa12) as teacher to be able to directly compare our numbers with the ones of CoDIR.

We re-implement PKD (Sun et al. 2019) and ALP-KD (Passban et al. 2021) approaches using the default settings proposed in the respective papers. We used early stopping based on performance on the development set, while making sure that the figures are in line with the ones reported in the papers. More precisely, the best layer setting for PKD teacher BERT12 is $\{2,4,6,8,10\}$ to distill into DistilBERT6. For DistilRoBERTa6, we choose the intermediate layers 4, 8, 12, 16, 20 from the teacher RoBERTa24 model for distillation that we found to work the best on the development set.

Using ALP-KD, we compute attention weights for the intermediate layers of the teacher (i.e., 1 to 11 for BERT12 and 1 to 23 for RoBERTa24 models) to calculate the weighted intermediate representations of the teacher for each intermediate layer of the student model (i.e., 1 to 5 layers of the student models). Since, the hidden dimensions of the RoBERTa24 and DistilRoBERTa6 are different, we linearly transform them into same lower-dimensional space. We train the PKD and ALP-KD models following (Sun et al. 2019; Passban et al. 2021).

For RAIL-KDl, at each epoch we randomly select 5 layers from the intermediate layers of the teacher (i.e., from layers 1 to 11 for BERT12 model and 1 to 23 for RoBERTa24 model). Then, we sort the layer indexes and perform layer-wise distillation (Figure 1(a)) for RAIL-KDl. For RAIL-KDc, we concatenated the representations of the sorted randomly selected intermediate layers and then perform concatenated representation distillation (Figure 1(b)). We use a linear transformation to map the intermediate representations (layer-wise or concatenated representations) into 128-dimensional space and normalize them before computing the loss $L_{\text{RAIL-KD}^{l/c}}$ for both BERT12 and RoBERTa24 distillations. We fixed $\alpha_{i}=1$, $\lambda_{1},\lambda_{2},\lambda_{3}=1/3$ for our proposed approaches We didn’t find a significant improvement when changing these values.. We search learning rate from {1e-5, 2e-5, 5e-5, 4e-6}, batch size from {8, 16, 32}, and fixed the epoch number to 40 for all the experiments. we run all experiments 5 times and report average score, in order to validate the credibility of our results. We ran all the experiments on a single NVIDIA V100 GPU using mixed-precision training (Micikevicius et al. 2018) and PyTorch Paszke et al. 2019 framework.

Our results indicate that random layer mapping not only delivers consistently better results than the deterministic mapping techniques such as PKD, but it has less computational overhead during training, while avoid extensive search experiments to find optimal mapping. On the other hand, using attention for layer selection (ALP-KD) or contrastive learning (CoDIR) leads to slightly worse result than than random selection.

Results

Table 2 shows the performances of models trained on GLUE tasks, and evaluated on their respective dev and test sets for 12 layer to 6 layer distillation. Bert12 and DistilBERT6 are used as backbone for the teacher and student models respectively. The baselines are fine-tuning without KD (w/o KD) and Vanilla KD. Moreover, we directly compare RAIL-KDlc results with PKD and ALP-KD as more competitor techniques.

First, we observe that in the 12 to 6 layer distillation, 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 margin of 0.5% and 0.3% on average on dev and test sets respectively. We notice that, except on RTE test, our RAIL-KDl/c obtained the highest per-task performances. Third, we observe that RAIL-KDl/c 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 Table 3. In this experiment, we used Roberta24 and DistillRoberta6 as teacher and students models respectively. Overall, RAIL-KD outperforms the best baseline by 1.2% and 0.3% on dev and test sets respectively. Interestingly, the gap on dev compared with PKD and ALP-KD is larger than the one reported on Bert12 experiments, and PKD Test socres are much lower from that of ALP and RAIL-KD. This might be because PKD skips a large number of intermediate layers on RoBERTa24, and the computational cost of ALP-KD attention weights over a large number of teacher layers might produce smaller weights on Roberta24 compared to Bert12.

Furthermore, we demonstrate the effectiveness of RAIL-KD by directly comparing it with CoDIR Sun et al. 2020a, the current state-of-the-art ILD method. It uses the contrastive objective and a memory bank to extract a large number of negative samples for contrastive loss calculations. Table 4 shows GLUE test results of both approaches when distilling RoBERTa12 to DistillRoberta6. CoDIR results are adopted from their paper, and we followed their experimental protocol by not reporting scores on STS-B. In addition, we report the overall training time speedup against the teacher for different techniques. On average, RAIL-KD is almost twice faster while performs on par with CoDIR (+0.2%). Moreover, RAIL-KD outperforms CoDIR on 5 out of 8 datasets.

To evaluate the impact of random layer selection on the performance of RAIL-KD compared to the other baselines, we report the standard deviation of the DistilBERT6 student models on the three smallest GLUE tasks, which are known to have the highest variance in Table 5.

Figures shows that RAIL-KD variance is at the same scale compared with PKD and ALP-KD on CoLA and MRPC, and even lower on RTE. This is a strong indicator that the gains from RAIL-KD are not due to random layer selection.

2 Out-of-Distribution Test

We further validate the generalization ability of student models by measuring their robustness to in-domain and out-of-domain evaluation. We do so by evaluating models fine-tuned on MLI, QQP and SST-2 and then evaluated on SciTail, PAWS, and IMDB respectively. These datasets contains counterexamples to biases found in the training data McCoy et al. 2019; Schuster et al. 2019; Clark et al. 2019. Performances of Bert12/Roberta24 teacher and DistilBERT6/DistilRoBERTa6 student variants are reported in Table 6. Also, we compute the unweighted average score of the three tasks.

First, we notice high variability in models rank and some inconsistencies in performances across tasks when compared with in-domain results. This was also reported in prior works on out-of-domain training and evaluation Clark et al. 2019; Mahabadi et al. 2020; Utama et al. 2020; Sanh et al. 2020. Still, RAIL-KDl/c clearly outperforms all baselines across tasks. Surprisingly, we observe that PKD and ALP-KD perform poorly (on all three tasks) compared to the Vanilla KD baseline.

Interestingly, we observe that RAIL-KDl performs consistently better (1.1% on average) than RAIL-KDc on Roberta24 compression, while RAIL-KDc perform better (1.1% on average) on Bert12. These results suggest that layer-wise distillation approach is more effective than concatenated distillation when we have a large capacity gap (layer number) between the teacher and the student, and vice versa.

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 Bert12 teacher (i.e. layers 2, 4 and 6) and the first three layer representations of the student for PKD, ALP-KD, RAIL-KDl/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. The similar trend 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 method gives 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 and 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 DistilBERT6 ALP-KD studnet on CoLA (left), RTE (middle), and MRPC (right) Similar trends found on other datasets.. 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 Bert12. This may be 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 out-of-domain evaluation compared with RAIL-KD.

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 equal 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 in performance improvement. Furthermore, our approach shows better performance for larger model distillation with faster training time, which opens up an avenue to investigate our approach for a super-large model such as GPT-2 (Radford et al. 2019) distillation using intermediate layers, as well as to improve robustness and generalization on a wider range of NLU tasks Ghaddar et al. 2021a; Ghaddar et al. 2021b.

Acknowledgments

We thank Mindspore https://www.mindspore.cn/ for the partial support of this work, which is a new deep learning computing framework.

References