RW - KD : Sample-wise Loss Terms Re-Weighting for Knowledge Distillation (2021)

Abstract: Knowledge Distillation (KD) is extensively used in Natural Language Processing to compress the pre-training and task-specific fine-tuning phases of large neural language models. A student model is trained to minimize a convex combination of the prediction loss over the labels and another over the teacher output. However, most existing works either fix the interpolating weight between the two losses apriori or vary the weight using heuristics. In this work, we propose a novel sample-wise loss weighting method, RW-KD. A meta-learner, simultaneously trained with the student, adaptively re-weights the two losses for each sample. We demonstrate, on 7 datasets of the GLUE benchmark, that RW-KD outperforms other loss re-weighting methods for KD.

## introduction knowledge distillation @xcite has proven highly effective for compressing a large-scale nlp model @xcite , called teacher in kd terms, into a smaller one, the student. a key factor behind kd's success is the use of teacher output as soft labels for supervising the training of the student @xcite . the latter model is trained by jointly minimizing the losses on both hard and soft labels. the contribution of each loss term is conventionally controlled by a balancing hyperparameter. however, recent studies suggested that hard and soft label importance is sample-wise @xcite , and only a subset of training samples are crucial for distillation @xcite . for instance, teacher outputs may be of poor quality for some samples @xcite , but highly informative for others @xcite . also, researchers have found that adjusting loss weights during training greatly benefits performance of kd @xcite @xcite . however, the contribution of loss terms is heuristically decayed by an annealing factor, yet another hyperparameter. we argue that using the same weights for all training samples, referred to in our work as singleweight, prevents exploiting the full advantage of kd, because each data sample might have different optimal weights for the loss terms. we propose a meta-learning approach to learn samplewise weights of loss terms. we revisit learning to weight approaches @xcite , initially proposed for noisy sample downweighting, and adapt it for loss terms weighting in kd. experimental results show that our kd loss weighting scheme consistently outperforms its counterparts on 7 tasks from the glue benchmark @xcite . a fine-grained analysis of the learned weights shows that, compared to the baselines, our meta-learner explores a greater range of kd weights to find the sample-wise optimal values. ## related work in recent years, knowledge distillation for bertlike models @xcite has been extensively studied, leveraging intermediate layer matching @xcite @xcite , data augmentation @xcite @xcite , adversarial training @xcite @xcite , and lately loss terms re-weighting @xcite @xcite . in this work, we explore the latter direction with a meta learning approach @xcite . learning to weight approaches @xcite were mainly proposed to learn per-sample loss weights in order to discount noisy samples thanks to an auxiliary meta-learner which re-weights training samples of the main model. such approaches often train a meta-learner on a clean validation set, or on small-loss training samples if no clean data is available. the meta-learner architecture varies from a simple multi-layer perceptron (mlp) as in meta-weight-net @xcite to lstm-based encoder as in mentornet @xcite . the work of @xcite on multi-modal model compression with kd is the most similar to ours. the authors train a mlp meta learner @xcite , on the validation set, which assigns samplelevel weights for 3 loss terms that are calculated when text, image, and both modalities are given as input. in our work, we use a transformer-based meta learner to estimate the sample-wise optimal weights for kd with gradient similarity (see section 3.2). ## methodology let t p¨q be a fine-tuned fixed teacher, and s θ p¨q the student model parameterized with θ. given a training set of tx i , y i u| n reweighting kd we propose a sample-wise reweighting method for kd to learn a balance between the ce and kd loss for every training sample. the new training loss is computed as follows: finding the optimal weights for each loss is intractable. our solution is inspired by @xcite and @xcite . these works investigate which training samples are most responsible for the generalization performance. we follow this line of works and perturb different losses in kd training to identify which loss is more influential and informative. ## conclusion in this paper, we show the importance of samplewise loss term weighting in knowledge distillation and propose rw-kd a method which does this and leads to better distillation performance on 7 glue tasks. future work involves combining rw-kd with state of the art kd methods that use extra loss terms such as intermediate layer similarity @xcite , attention matching @xcite , and adversarial @xcite losses. we expect that these methods can take full advantage of rw-kd, since they use single-weight loss terms weights. in addition to kd training, we will investigate apply-ing our reweighting method to multi-task learning (mtl) scenarios @xcite @xcite , where learning to balance losses from different tasks is critical to benefit all tasks involved.