1 Introduction↩︎

Pre-trained Language Models (PLMs), such as BERT [5], RoBERTa [6] and XLNet [7] have shown remarkable abilities to match and even surpass human performances on many Natural Languages Understanding (NLU) tasks [8]–[10]. 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 [11], [12], pruning [13]–[15], optimizing the Transformer architecture [16]–[19], and knowledge distillation [3], [20]–[26] 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 [27]–[29], 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 [21], [30], [31], adversarial training [24], [32], [33], and intermediate layer distillation (ILD)  [3], [23], [34], [35].

When it comes to Bert compression, ILD leads to clear gains in performances [20], [21], [23] 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³, and then matching them via regression [1] or cosine similarity [20] 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 [3]. There are some solutions in the literature mostly rely on layer combination [2], attention-based layer projection [3] and contrastive learning [4].

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 [9]. 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.

2 Related Work↩︎

Recent years, have seen a wide range of methods have emerged aiming to expand knowledge transfer of transformer-based [36] NLU models beyond logits matching. DistillBert [20] added a cosine similarity loss between teacher and student embeddings layer. TinyBERT [21], MobileBERT [22], and MiniLM [34] matched the intermediate layers representations and self-attention distributions of the teacher and the student.

In PKD, [1] 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.⁴

Researchers have found that tuning the layer mapping scheme can significantly improve the performance of ILD techniques [1]. Nevertheless, finding the optimal mapping can be challenging, which is referred to as the layer skip and search problems by  [3]. To address the layer skip problem, CKD [2] 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 [3] 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 [6] or GPT-2 [37]. Alternatively, CODIR [4] exploited contrastive learning [38] 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 [1], CKD [2], 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 [3] and CoDIR [4] 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 ⁵. Otherwise said, we don’t compare with works like TinyBERT [21] or MiniLM [34], which use extra losses like self-attention distribution matching. However, we expect that these methods, as well as state-of-the-art [24], [39] one can take full advantage of RAIL-KD, since they use deterministic layer mapping scheme.

3 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 [4], 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}\)  [4]:

\[\begin{align} \bar{h}_{i,X}^T = \frac{1}{L} \sum_{k=0}^{L-1} H_{i,x_k}^T ; \bar{h}_{i,X}^{S_\theta} = \frac{1}{L} \sum_{k=0}^{L-1} H_{i,x_k}^{S_\theta} \end{align}\] 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)).

3.1 Layer-wise RAIL-KD↩︎

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. \[\begin{align} & L_{\text{RAIL-KD}^l}= \sum_{X\in \mathcal{X}} \sum_{i=1}^m \alpha_i \left( \lvert\lvert \frac{\hat{h}_{i,X}^T}{\lvert\lvert\hat{h}_{i,X}^T \rvert\rvert_2}-\frac{\hat{h}_{i,X}^{S_\theta}}{\lvert\lvert\hat{h}_{i,X}^{S_\theta} \rvert\rvert_2}\rvert\rvert_2^2 \right) \end{align} \label{eq1}\tag{1}\] where \(\mathcal{X}\) denotes the set of training samples, and \(\alpha_i\) is a hyper-parameter to weigh the layer-wise distillation loss.

3.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. \[L_{\text{RAIL-KD}^c}=\sum_{X\in \mathcal{X}} \lvert\lvert \frac{\hat{h}_X^T}{\lvert\lvert\hat{h}_X^T \rvert\rvert_2}-\frac{\hat{h}_X^{S_\theta}}{\lvert\lvert\hat{h}_X^{S_\theta} \rvert\rvert_2}\rvert\rvert_2^2 \label{eq2}\tag{2}\] Any type of loss such as contrastive [4], or mean-square-error (MSE) [1], [3] can be applied for our RAIL-KD approach.

3.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: \[\mathcal{L}=\lambda_1 L_{\text{CE}} + \lambda_2 L_{\text{KD}} + \lambda_3 L_{\text{RAIL-KD}^{l/c}}\] 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\).

4 Experimental Protocol↩︎

4.1 Datasets and Evaluation↩︎

We evaluate RAIL-KD on 8 tasks from the GLUE benchmark [9]: 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 [1], [3], [4], [21], 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 [40], PAWS (Paraphrase Adversaries from Word Scrambling) [41], and IMDb (Internet Movie Database) [42] test sets to evaluate the models fine-tuned on MNLI, QQP, and SST-2 tasks respectively.

4.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 [5] as teacher (Bert\(_{12}\)) and 6 layer DistilBert [20] as student (DistillBert\(_6\)) to compare with PKD [1] and ALP-KD [3]. Also, we use 24 layers RoBERTa-large [6] and 6 layers DistilRoberta [43] as the backbone for teacher (RoBERTa\(_{24}\)) and student (DistilRoberta\(_6\)) respectively to compare models when \(n >> m\). Furthermore, we perform evaluation using 12 layers RoBERTa-base (RoBERTa\(_{12}\)) as teacher to be able to directly compare our numbers with the ones of CoDIR.

We re-implement PKD [1] and ALP-KD [3] 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 BERT₁₂ is \(\{2, 4, 6, 8, 10\}\) to distill into DistilBERT₆. For DistilRoBERTa₆, we choose the intermediate layers 4, 8, 12, 16, 20 from the teacher RoBERTa₂₄ 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 BERT₁₂ and 1 to 23 for RoBERTa₂₄ 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 RoBERTa₂₄ and DistilRoBERTa₆ are different, we linearly transform them into same lower-dimensional space. We train the PKD and ALP-KD models following [1], [3].

For RAIL-KD\(^l\), at each epoch we randomly select 5 layers from the intermediate layers of the teacher (i.e., from layers 1 to 11 for BERT₁₂ model and 1 to 23 for RoBERTa₂₄ model). Then, we sort the layer indexes and perform layer-wise distillation (Figure 1 (a)) for RAIL-KD\(^l\). For RAIL-KD\(^c\), 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 BERT₁₂ and RoBERTa₂₄ distillations. We fixed \(\alpha_i=1\), \(\lambda_1, \lambda_2, \lambda_3 = 1/3\) for our proposed approaches ⁶. 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 [44] and PyTorch [45] 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.

5 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. Bert\(^{12}\) and DistilBERT\(^6\) 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-KD\(^{lc}\) 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-KD\(^{l/c}\) obtained the highest per-task performances. Third, we observe that RAIL-KD\(^{l/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 Roberta\(_{24}\) and DistillRoberta\(_6\) 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 Bert\(^{12}\) 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 RoBERTa\(_{24}\), and the computational cost of ALP-KD attention weights over a large number of teacher layers might produce smaller weights on Roberta\(_{24}\) compared to Bert\(_{12}\).

Furthermore, we demonstrate the effectiveness of RAIL-KD by directly comparing it with CoDIR [4], 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 RoBERTa\(_{12}\) to DistillRoberta\(_6\). 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.

5.1 Impact of Random Layer Selection↩︎

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 DistilBERT\(_6\) 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.

5.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 [46]–[48]. Performances of Bert\(_{12}\)/Roberta\(_{24}\) teacher and DistilBERT\(_6\)/DistilRoBERTa\(_6\) 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 [48]–[51]. Still, RAIL-KD\(^{l/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-KD\(^l\) performs consistently better (1.1% on average) than RAIL-KD\(^c\) on Roberta\(_{24}\) compression, while RAIL-KD\(^c\) perform better (1.1% on average) on Bert\(_{12}\). 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.

6 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. 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 DistilBERT\(_6\) ALP-KD studnet on CoLA (left), RTE (middle), and MRPC (right) ⁷. 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 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.

7 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 [37] distillation using intermediate layers, as well as to improve robustness and generalization on a wider range of NLU tasks [52], [53].

Acknowledgments↩︎

We thank Mindspore⁸ for the partial support of this work, which is a new deep learning computing framework.

References↩︎