Robustness Challenges in Model Distillation and Pruning for Natural Language Understanding (2023)

Abstract: Recent work has focused on compressing pre-trained language models (PLMs) like BERT where the major focus has been to improve the in-distribution performance for downstream tasks. However, very few of these studies have analyzed the impact of compression on the generalizability and robustness of compressed models for out-of-distribution (OOD) data. Towards this end, we study two popular model compression techniques including knowledge distillation and pruning and show that the compressed models are significantly less robust than their PLM counterparts on OOD test sets although they obtain similar performance on in-distribution development sets for a task. Further analysis indicates that the compressed models overfit on the shortcut samples and generalize poorly on the hard ones. We further leverage this observation to develop a regularization strategy for robust model compression based on sample uncertainty.

## introduction large pretrained language models (plms) (e.g., bert @xcite , roberta @xcite , @xcite ) have obtained state-of-the-art performance in several natural language understanding (nlu) tasks. however, recent studies @xcite @xcite indicate that plms heavily rely on shortcut learning/spurious correlations, rather than acquiring higher level language understanding and semantic reasoning in several nlu tasks. specifically, these models often exploit dataset biases and artifacts, e.g., lexical bias and * most of the work was completed while the first author was an intern at microsoft research during summer 2021. on the other hand, it is difficult to use these large plms models in real-world applications with latency and capacity constraints, e.g., on edge devices and mobile phones. thus, model compression emerges as one of the techniques to reduce model size, speed up inference, and save energy without significant performance drop for downstream tasks. state-of-the-art model compression techniques such as knowledge distillation @xcite and pruning @xcite primarily focus on evaluating compressed model performance in in-distribution test data. however, in-distribution testing is insufficient to capture the generalizability of plms @xcite . in contrast to existing work that is geared towards general-purpose plms @xcite @xcite , this work aims to study the impact of compression on the shortcut learning and ood generalization ability of compressed models. towards this end, we conduct comprehensive experiments to evaluate the ood robustness of compressed models, with bert as the base encoder. we focus primarily on two popular model compression techniques in the form of pruning and knowledge distillation @xcite . for pruning, we consider two popular techniques including iterative magnitude pruning @xcite and structured pruning @xcite . specifically, we explore the following research questions: are distilled and pruned models as ro-bust as their plm counterparts for downstream nlu tasks? what is the impact of varying the level of compression on ood generalization and bias of compressed models? we evaluate the performance of several compressed models obtained using the above techniques on both standard indistribution development sets and ood test sets for downstream nlu tasks. experimental analysis indicates that distilled and pruned models are consistently less robust than their plm counterparts. further analysis of the poor generalization performance of compressed models reveals some interesting observations. for instance, we observe that the compressed models overfit on the easy / shortcut samples and generalize poorly on the hard ones. this motivates our second research question: how to regularize model compression techniques to generalize across samples with varying difficulty? this brings some interesting challenges since we do not know which samples are easy or hard apriori. based on the above observations, we propose a bias mitigation framework to improve the ood robustness of compressed models, termed as rmc (robust model compression). first, we leverage the uncertainty of the deep neural network to quantify the difficulty of a training sample. this is given by the variance in the prediction of a sample from multiple sub-networks of the original large network obtained by model pruning. second, we leverage this sample-specific measure for smoothing and regularizing different families of compression techniques. the major contributions of this work can be summarized as follows: ## attribution of low robustness in this section, we explore the factors that lead to low robustness of compressed models. previous work has demonstrated that the performance of different models on the glue benchmark @xcite tends to correlate with the performance on mnli, making it a good representative of natural language understanding tasks in general @xcite . for this reason, we choose the mnli task for a study. for the mnli task, we consider the dataset splits from @xcite . the authors partition the development set into easy/shortcut foot_3 and hard subsets. in this experiment, we use pruned models with varying sparsity to investigate the reason for the low robustness of the compressed models. we have the following key observations. observation 1: the compressed models tend to overfit the easy/shortcut samples and generalize poorly on the hard ones. the performance of pruned models at five sparsity levels (ranging between [0.2 -0.85]) on the easy and hard samples for the mnli task is illustrated in figure 1 . it demonstrates that the accuracy on the hard samples is much lower compared to the accuracy on the easy ones. as the sparsity increases, we observe a larger accuracy drop on the hard samples compared to the easy ones. in particular, the accuracy gap between the two subsets is 22.7% at the sparsity of 0.85, much higher than the 16.1% accuracy gap at the sparsity of 0.4. these findings demonstrate that the compressed models overfit on the easy samples, while generalizing poorly on the hard ones. furthermore, this phenomenon is amplified at higher levels of sparsity for the pruned models. observation 2: compressed models tend to assign overconfident predictions to easy samples. one of the potential reasons is that compressed models are more prone to capture spurious correlations between shortcut features in training samples with certain class labels for their predictions @xcite . ## mitigation framework in this section, we propose a general bias mitigation framework (see figure 2 ), termed as rmc (robust model compression), to improve the robustness of compressed models on downstream tasks. our rmc framework follows the philosophy of taskspecific knowledge distillation @xcite , but with explicit regularization of the teacher network leveraging sample uncertainty. this prevents the compressed model from overfitting in the easy samples that contain shortcut features and helps improve its robustness. this regularized training is implemented in two stages. ## mitigation performance evaluation in this section, we conduct experiments to evaluate the robustness of our rmc mitigation framework. ## conclusions in this work, we conduct a comprehensive study of the robustness challenges in compressing large plms when fine-tuning in downstream nlu datasets. furthermore, we propose a general mitigation framework with instance-level smoothing for robust model compression. experimental analysis demonstrates our framework to improve the generalization and ood robustness of compressed models for different compression techniques, while not sacrificing the in-distribution performance.