Learn From One Specialized Sub-Teacher: One-to-One Mapping for Feature-Based Knowledge Distillation (2023)

Abstract: Knowledge distillation is known as an effective technique for compressing over-parameterized language models. In this work, we propose to break down the global feature distillation task into N local sub-tasks. In this new framework, we consider each neuron in the last hidden layer of the teacher network as a specialized sub-teacher. We also consider each neuron in the last hidden layer of the student network as a focused sub-student. We make each focused sub-student learn from one corresponding specialized sub-teacher and ignore the others. This will facilitate the task for the sub-student and keep it focused. Our proposed method is novel and can be combined with other distillation techniques. Empirical results show that our proposed approach outperforms the state-of-the-art methods by maintaining higher performance on most benchmark datasets. Furthermore, we propose a randomized variant of our approach, called Masked One-to-One Mapping. Rather than learning all the N sub-tasks simultaneously, we focus on learning a subset of these sub-tasks at each optimization step. This variant enables the student to digest the received flow of knowledge more effectively and yields superior results.

## introduction large language models, also known as general purpose language models, have revolutionized the nlp domain @xcite @xcite . they are large architectures composed of several transformer blocks @xcite , typically trained on large unlabeled corpora in a selfsupervised way @xcite . they achieved state-of-the-art (sota) performance on downstream tasks through finetuning when the data is scarce @xcite . however, since these models are typically large and computationally expensive, such as bert with millions of parameters @xcite and gpt-3 with billions of parameters @xcite , they are not highly adapted to real-world applications like mobile computing. model compression is an active field of research that focuses on effectively reducing the model size without significant performance degradation @xcite @xcite . knowledge distillation (kd) by @xcite is one of the effective compression techniques in nlp where the knowledge of a highly capable large model, i.e., teacher, is transferred to a smaller efficient model, i.e., student. kd essentially requires designing a loss function to minimize the distance of the output or the intermediate representations between the student and the teacher @xcite @xcite . to distill the intermediate representations, previous research relied on the mean square error (mse) as an objective function between the student and the teacher global representations @xcite . however, this metric is sensitive to scale @xcite and it is not accurate in high dimensional space @xcite . other works used the cosine distance as an alternative, but, it also has several limitations @xcite such as not performing well with sparse data. in this work, we focus on distilling the teacher's last hidden layer representation into the student's last hidden layer representation. to achieve that, we propose a novel kd approach where we reformulate this feature distillation task from a global problem to n local sub-problems. here, the term "local" refers to the consideration of each dimension or neuron individually. in this work, our main assumption is that the student last hidden layer has the same dimension as the teacher last hidden layer. we note that this is a common assumption in feature distillation @xcite . in our proposed framework, we consider each neuron in the teacher's last hidden layer as a specialized sub-teacher and each neuron in the student's last hidden layer as a focused sub-student. each specialized sub-teacher is in charge of distilling its knowledge to the corresponding focused sub-student. we call this a one-to-one matching between the teacher and the student last hidden layers units. to accomplish this one-to-one distillation, we propose a novel objective function where we maximize the per-batch correlation between the outputs of each specialized sub-teacher and its corresponding focused sub-student. empirical results show that studying the global feature distillation task from a local viewpoint helped the student to meet the global teacher's features representation. in addition, we propose an augmented variant of our approach where the model learns a subset of sub-tasks at a time. to sum up, our contributions are the following: ## related work model compression techniques are employed to effectively reduce the size of a neural network while maintaining good performance. various approaches have been proposed such as pruning @xcite . it aims to remove unimportant structures, e.g., weights, neurons, and even entire layers from the model @xcite . another technique is weight sharing @xcite , where different parts of the model use the same set of weights to perform computations @xcite . quantization @xcite ) is yet another strategy used in model compression, where the weights and activations in the network are represented using lower bit integers instead of higher precision floating-point numbers @xcite . knowledge distillation @xcite @xcite , which is the focus of this paper, is also a successful compression technique. it involves training a smaller, more efficient student model to mimic the behavior of a larger, knowledgeable teacher model. by leveraging the knowledge of the teacher model to the student model, the student can achieve comparable performance to the teacher while maintaining a lower size. knowledge distillation has been proven to be an effective technique for large language model compression @xcite . it can be applied during the pre-training stage to generate general-purpose distilled models @xcite and during the fine-tuning stage to generate task-specific distilled models @xcite @xcite . in @xcite , the authors show that large language models, e.g., bert, suffer from over-parametrization in domain-specific tasks. thus, previous work has been improving the taskspecific distillation. several methods focused on enhancing the objectives of the distillation process. these improvements mainly focused on which part of the teacher architecture can be distilled into the student architecture such as the attention matrices @xcite , the different hidden states @xcite , and the prediction layer @xcite . coming up with effective objective functions to distill the knowledge from any part of the teacher into the student is critical. for the logit-based kd, the kl divergence or the mse are used as objective functions to minimize the distance between the logits of the student and the logits of the teacher @xcite . following that, in @xcite , the authors provided a novel viewpoint to study the logit distillation by reformulating the classical kl divergence loss into two parts, which showed a good improvement. in the feature-based kd, the mse and the cosine distance are mainly used as objective functions to align the intermediate layers representations of the teacher and the student @xcite @xcite . in this work, we argue that relying on the mse and the cosine distance as loss functions to align the intermediate representations between the student and the teacher is not an optimal choice. mse is sensitive to scale and does not perform well in highdimensional space. in a high dimensional space, which is the case in a neural network, the data tend to be sparse and all the data points become uniformly distant from each other @xcite . in fact, the embeddings (features representation) provided by bert-like models are high-dimensional and sparse tensors in nature @xcite . thus, using mse or cosine distance to measure the distance between these sparse representations is not optimal as explained and advocated in @xcite . so, splitting the task into sub-tasks and using oneto-one metric alleviates this problem and breaks the curse of dimensionality. in @xcite , it is also shown that the cosine distance is not an accurate measurement of similarity between bert embeddings. moreover, the cosine distance measure can be also affected by sparsity. in fact, in high-dimensional space, it can output large angles between two sparse vectors although they are similar in the non-zero components @xcite . another important limitation of the cosine distance is that it is a global metric and sensitive to unit permutation. for example, let w 1 be the output tensor of a given layer in the teacher and w 2 be the output of the corresponding layer in the student. in the following example, we compute the cosine distance (cosd) from w 1 and w 2 : it is a high value but the layers actually learned the same representation. in our work, instead of treating feature distillation as a global task @xcite between teacher-student layers, we reformulate it as multiple local sub-tasks. furthermore, for each per-dimension kd sub-task, instead of using mse or cosine distance @xcite @xcite , we propose a noval loss function that utilizes a per-batch correlation function. this improves the student's ability to meet the global representation provided by the teacher. we note that the initial results of our work on the one-on-one mapping approach are presented in @xcite . ## our approach: one-to-one mapping in this work, we provide a new viewpoint on how to study the feature distillation task. we break down the global kd task into multiple local kd sub-tasks. specifically, in the last hidden layer of the teacher, we consider each unit as a specialized sub-teacher. in the last hidden layer of the student, we also treat each unit as a focused sub-student. each substudent should concentrate and learn only from the corresponding specialized sub-teacher. we refer to this as one-to-one mapping, which is reformulated as the cross-correlation between the outputs of each specialized sub-teacher and its corresponding focused sub-student. the mapping between each sub-student and each sub-teacher is fixed, i.e. based on an index, and thus remains consistent over batches and epochs. although our approach can be applied to different student-teacher hidden layers, as in @xcite , this work specifically focuses on the last hidden layer. typically, in a kd framework, as illustrated in figure 1 , we have the teacher network modeled by f θ , which is an over-parameterized knowledgeable model. the student network is modeled by we assume that y t and y s are the features representation of the last hidden layer of the teacher h last t and the last hidden layer of the student h last s , respectively. y s and y t are assumed to be meancentered over the batch dimension. we assume that h last t and h last s have n hidden units. the unit i in h last t represents the specialized sub-teacher i and the unit i in h last s represents the focused substudent i . the sub-task i is distilling the knowledge from the specialized sub-teacher i to the focused sub-student i by reducing the distance between the features learned by each of them. to simplify the task for the sub-student i and to keep it focused, we force it to learn only from the sub-teacher i and ignore the other sub-teachers. we reformulate the objective function in this one-toone mapping as maximizing the cross-correlation between the two variables y t i and y s i . y t i and y s i are the output values of the sub-teacher i and the substudent i , respectively. the variables y t i and y s i have n samples coming from the different examples in the input batch. maximizing the correlation across batches between the two aforementioned variables, i.e., minimizing the following loss function: as we want to make the task easier for the substudent i , so it can effectively digest the received information, we force it to mimic only the teacher i and repel all knowledge coming from the other teachers. this results in minimizing the following term: thus, our final per-dimension kd loss function, as shown in figure 1 , is: in l i , the first term is for maximizing the crosscorrelation over batches between the output of the sub-student i and the output of the sub-teacher i . the second term is for minimizing the cross-correlation between the sub-student i and each sub-teacher j given j ̸ = i and @xmath0, 2, ..., n }. our end distillation loss is the sum of the n local kd losses. additionally, we introduce λ 1 and λ 2 as the weights to control the contribution of the first term and the second term of the loss function, respectively: the final training loss of the original student is: where l ce is the classical cross entropy loss between the student predictions and the ground truth labels: where x i is an input sample and y i is its ground truth label. n is the number of samples per batch. empirical results will show that our one-to-one mapping approach effectively facilitates the alignment of the student features representation with the global teacher feature representation. our approach can be implemented at a low computational cost and can be combined with other knowledge distillation methods. ## experimental results in our experiments, the teacher is the bert-base model, with 110 million parameters, after being fine-tuned on each of the datasets for 3 epochs. the student is the distilbert-base with 66 million parameters. n , which is the number of neurons, in the last hidden layer of the teacher and the student, is equal to 768. all experiments are repeated for 5 random seeds, the learning-rate is set to 5e-5, the maximum sequence length is set to 128, and the batch size is set to 16. ## ablation study see results in rather than looking at our approach as a competing method to the kd state-of-the-art approaches, it is proposed to replace any mse or cosine distance-based feature distillation methods between two hidden layers included in any knowledge distillation approach. to show that our approach can outperform mse and cosine distance, applied globally, in the existing pipelines, we experiment with srrl @xcite , one of the latest kd approaches. srrl kd loss has two components: the feature loss, which is an mse between the last layer's features of the teacher and the last layer's features of the student, ## masked one-to-one mapping in our one-to-one mapping feature distillation approach, the goal for the student is to mimic the last hidden layer representation of the teacher. to achieve this, we break down the global feature distillation task into n local sub-tasks. to further enhance the student performance, we propose an augmented variant of our approach called masked one-to-one mapping. instead of learning all n sub-tasks simultaneously, it is more effective for the student to focus on a subset of the total subtasks at a time. this approach enables the student to effectively absorb and manage the incoming flow of knowledge from the teacher. as shown in figure 4 , during each iteration (optimization step), we randomly mask a subset of units in the last hidden layer of the teacher and in the last hidden layer of the student. afterwards, we will have m (where @xmath1, i.e., m feature distillation sub-tasks left for the student to learn. we accomplish this by applying a same binary mask m to the last hidden layer features representation of both the student and the teacher simultaneously. formally, given y @xmath2, the features representation of the student and the teacher, respectively. n is the number of neurons in the last hidden layers of the teacher and the student. n is the batch size. to generate a random binary mask m , we create a vector of size n containing independent binomial random variables. each variable in this vector takes the value 1 with a probability p and 0 otherwise. we multiply element-wise this binary mask m by each column of y t and y s to generate the masked features representation y ′ t of the teacher and y ′ s of the student, respectively: we conduct experiments on 8 glue datasets using the same setup as described in subsection 4.2. however, in these new experiments, we extend the training time to 5 epochs because as shown in tables 3 and 4 , the average performance on the glue tasks of our student model is higher when trained for 5 epochs compared to when trained for 3 epochs. we introduce p as the percentage (%) of sub-tasks to be learned at a given optimization step. we conduct experiments for all values of @xmath3, 20, 30, 40, 50, 70, 90}. however, we report the results of the two best setups, in function of p, in terms of the average performance across all 8 glue datasets. additionally, we provide the best performance achieved for each individual dataset. for each dataset, the experiment is repeated for 5 seeds and the average is reported. as illustrated in figure 4 : masked one-to-one mapping: in each optimization step, the neurons that are marked by the red color will be masked. ## conclusion and future work in this paper, we reformulated the global feature distillation problem into n local sub-problems. we proposed a one-to-one mapping between each neuron in the last hidden layer of the teacher, i.e., specialized sub-teacher, and each neuron in the last hidden layer of the student. i.e., focused substudent. to achieve this goal, we proposed a local correlation-based loss across batches between each specialized sub-teacher output and its corresponding focused sub-student output. our approach only requires the teacher and the student to have the same last hidden layer size. several experiments proved the effectiveness and consistency of our method. our approach can be added to any kd method in nlp or vision. moreover, we proposed an augmented variant of our one-to-one mapping approach called masked one-to-one mapping. instead of learning all the sub-tasks simultaneously, we make the student learn a subset of the tasks at a time. future work includes exploring the same distillation process with several intermediate layers and experimenting with our approach beyond the nlp tasks.