Why Skip If You Can Combine: A Simple Knowledge Distillation Technique for Intermediate Layers (2020)

Abstract: With the growth of computing power neural machine translation (NMT) models also grow accordingly and become better. However, they also become harder to deploy on edge devices due to memory constraints. To cope with this problem, a common practice is to distill knowledge from a large and accurately-trained teacher network (T) into a compact student network (S). Although knowledge distillation (KD) is useful in most cases, our study shows that existing KD techniques might not be suitable enough for deep NMT engines, so we propose a novel alternative. In our model, besides matching T and S predictions we have a combinatorial mechanism to inject layer-level supervision from T to S. In this paper, we target low-resource settings and evaluate our translation engines for Portuguese→English, Turkish→English, and English→German directions. Students trained using our technique have 50% fewer parameters and can still deliver comparable results to those of 12-layer teachers.

## introduction almost in all deep learning tasks, including neural machine translation (nmt), an ensemble of models outperforms a single model. in fact, ensemble modelling (training multiple models and ensemble decoding) is supported by most publicly available nmt frameworks @xcite @xcite . however, we know that dealing with multiple models could be challenging, especially in deep learning scenarios. to tackle the issue, one effective solution is to compress the knowledge in an ensemble into a single model through distillation @xcite . the core part of any knowledge distillation (kd) pipeline is a component that matches different mod-els' predictions, which is usually implemented via multiple cost functions (see section 2). furthermore, we also need to take care of the architecture mismatch that may exist between student (s) and teacher (t ) models. in kd, these two models can have different architectures @xcite and the motivation is to be able to compress a large teacher into a smaller student. this research focuses on the aforementioned issue. if we distill from intermediate layers of a teacher that has more layers than its student, we have to select a subset of t layers and skip others as there are no peers for all of them on the s side. clearly, we do not benefit from the skipped layers in this scenario. this type of kd introduces a problem of finding an optimal subset of t layers (to distill from). although this might, to some extent, be mitigated via a search mechanism, our experimental results show that the problem is severe in nmt and each layer plays a unique role. therefore, we prefer to keep all layers rather than skip them. kd has recently become popular in nmt but, to the best of our knowledge, all nmt models @xcite are still trained using the original idea of kd @xcite , which is referred to as regular kd (rkd) throughout this paper. rkd only matches s and t outputs, regardless of their internal architecture. however, there exist techniques such as patient kd (pkd) @xcite proposed for other tasks that not only match final predictions but also focus on internal components and distill their information too @xcite . in this research, we borrowed those ideas and adapted them to nmt. this is the first contribution of the paper. pkd and other similar models suffer from the skip problem, which happens when t has more layers than s and some t layers have to be skipped in order to carry out layer-to-layer distillation. in this paper, we propose a model to distill from all teacher layers so we do not have to skip any of them. this is our second contribution by which we are able to outperform pkd. moreover, for the first time we report experimental results for transformer-based @xcite models trained with a layerlevel kd technique in the context of nmt. this set of results is our third and last contribution in this paper. the remainder of the paper is organized as follows: in section 2 we explain the fundamentals of kd. section 3 discusses the methodology. we describe the advantages of our model and accompany our claims with experimental results in section 4. finally, in section 5, we conclude the paper with our future plan. ## background usually, in multi-class classification scenarios the training criterion is to minimize the negative loglikelihood of samples, as shown in equation 1 : with the matching strategy proposed in kd, s learns to mimic its t . a teacher could be a deep model trained on a large dataset but we do not necessarily need to have the same complex architecture for s. we can distill teacher's knowledge into a smaller model and replicate its results with fewer resources. @xcite studied this problem and proposed a sequence-level extension of equation 2 for nmt models. they evaluated their idea on recurrent, lstm-based models @xcite and could run the final model on a cellphone. @xcite extended the original two-class idea (one s with one t ) to distill from multiple teachers. they trained an attentionbased recurrent model @xcite for their experiments. @xcite proposed a setting to train a multilingual transformer for different language directions. in order to have a high-quality multilingual model they distill knowledge from separately trained bilingual models. their work is one of the few papers that reports kd results for nmt on transformers. however, their results are not directly comparable to ours as they benefit from rich, multilingual corpora. @xcite introduced a pipeline where a student model learns from different checkpoints. at each validation step, if the current checkpoint is a better model than the best existing checkpoint, s learns from it, otherwise the best stored checkpoint is considered as the teacher. in all models discussed so far, i) s usually has the same architecture as its teacher(s) but we know that recent nmt models, particularly transformers, are deep models which makes them challenging to run on edge devices. moreover, ii) the training criterion in the aforementioned models is to combine final predictions. transformers have new components (e.g. self-attention) and multiple (sub-)layers that consist of valuable information @xcite and we need more than an output-level combination to efficiently distill for/from these models. therefore, a new technique that is capable of addressing i and ii is required. authors of pkd spotted the problem and focused on internal layers @xcite . they studied the limitations of rkd for bert @xcite models and introduced a layer-to-layer cost function. they select a subset of layers from t whose values are compared to s layers. they also showed that different internal components are important and play critical roles in kd. the layer-level supervision idea was successful for monolingual models but so far, no one has tried it in the context of nmt. in this paper, we investigate if the same idea holds for bilingual models or if nmt requires a different type of kd. moreover, we address the skip problem in pkd (shown in figure 1 ). it seems in deep teacher models we do not need to skip layers and we can distill from all layers. ## methodology in rkd, distillation only happens at the output level whereas pkd introduces layer-wise supervision. this idea is illustrated in figure 1 . figure 1 : the network on the left-hand side is s and the other one is t . in this example, t has 3 hidden layers and kd for intermediate layers can be applied using all layers or a subset of them, e.g. the second layer can be skipped. in pkd, finding a skippable layer is the main challenge. accordingly, we propose a combinatorial idea, ckd, by which we are able to fuse layers and benefit from all information stored in all layers. our idea can be formulated as follows: for f (), a simple concatenation followed by a linear projection provided the best results in our experiments, so in the previous example: ## experimental study although our proposed model is a general kd technique and can be applied in different settings, we narrow down the scope of this paper to lowresource, nmt settings. the incentive idea behind our project was to train nmt models for small datasets, so we report experimental results accordingly. to evaluate ckd, we trained multiple models to translate from english (en) into german (de), and from portuguese (pt) and turkish (tr) into english (en). for the pt|@xmath0, and the @xmath1. in @xmath2, we use the original split of datasets from iwslt 1 with 167k, 7590, and 5388 sentences for training, development, and test sets, respectively. for @xmath3, the split is 142k, 1958, and 1982 for training, development, and test sets. with this dataset selection our t models' results are comparable to publicly reported results. foot_1 on these datasets, our teachers outperform all other existing models so we can ensure that we distill from reliable sources. for @xmath4, the dataset is the same as the original transformer's @xcite , namely the training set includes 4.5m sentences, newstest2013 is used as the validation set and newstest2014 is our test set with 3000 and 3003 sentences, respectively. we selected this dataset to be comparable to a well-known baseline and make sure our training pipeline yields highquality engines. we preprocess datasets with sentence-piece @xcite . for @xmath5, we extracted a shared vocabulary set for both source and target sides with 32k subwords. both s and t are trained using the same training set. @xmath6. for @xmath7, we conduct two experiments. since our focus in this paper is to work with low-resource settings, in @xmath8, s and t are trained on a dataset of 200k sentences randomly sampled from the main dataset (4.5m). 3 for this experiment the vocabulary set size is 15k. in @xmath9, we slightly changed the setting where we use the entire set of 4.5m sentences to train t but s still uses the same 200k dataset. in this scenario, we assumed that there already exists a high-quality teacher trained on a large dataset but we only have a small in-house dataset to train the student. for this experiment the vocabulary size is 37k. ckd makes it possible to reduce the number of parameters in our students by 50% and yet deliver the same high-quality translations. accordingly, this enables us to run these translation engines on edge devices. for results reported in we tried a simple (and somewhat arbitrary) configuration for layer connections and there is no systematic strategy behind it. however, better results can be achieved with better heuristics or through a search process. moreover, as the mappings show there is no connection between student and teacher models' decoder layers. in our experiments, we noticed that any kd technique applied to the decoder considerably decreases performance, so we only use kd on the encoder side. more specifically, each student model has two decoder layers which only receive inputs from the same model's encoder layers and they are not connected to the teacher side. to train students with different kd techniques we use different loss functions. in t and no-kd we only have a single loss function (l) as described in the original transformer model @xcite . for models trained with rkd, an additional loss is involved to match teacher and student predictions (l kd ). the final loss in this case is an interpolation of the aforementioned losses: (β × l) + (η × l kd ) . in our experiments, β = (1 -η) where η = 0.1 is obtained through a search process over the set {0.1, 0.3, 0.5, 0.7, 0.9}. for students trained using pkd and ckd, a third loss is also used in addition to l and l kd . similar to other losses, the third one is also multiplied by a weight value (λ) to incorporate its impact into the training process. in this new setting, β = (1η -λ), η = 0.1, and λ = 0.7. the high value of λ compared to other weights shows the importance of intermediate kd for deep models. all these values are learned through an empirical study in order to minimize the final loss of translation engines. ## conclusion and future work in this paper, we proposed a novel model to distill from intermediate layers as well as final predictions. moreover, we addressed the skip problem of pkd. we applied our technique in nmt and showed its potential in training high-quality and compact models. in our future work, i) we are interested in distilling from deep nmt models into extremely small students with ckd, in the hope of achieving the same results of large models with much smaller counterparts. ii) we also try to improve the combination module and find a better alternative than concatenation. iii) finally, we plan to evaluate ckd in other tasks such as language modeling.