[Paper Review] Robustness Challenges in Model Distillation and Pruning for Natural Language Understanding

Mengnan Du, Subhabrata Mukherjee+4 more|arXiv (Cornell University)|Oct 16, 2021

Topic Modeling50 references4 citations

TL;DR

This paper identifies that knowledge distillation and pruning in NLP model compression lead to reduced out-of-distribution (OOD) robustness despite strong in-distribution performance. It proposes RMC, a regularization framework using sample uncertainty to smooth training, which improves OOD generalization by reducing overfitting to easy/shortcut samples without sacrificing in-distribution accuracy across multiple NLU tasks.

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. Experimental results on several natural language understanding tasks demonstrate that our bias mitigation framework improves the OOD generalization of the compressed models, while not sacrificing the in-distribution task performance.

Motivation & Objective

• To investigate whether compressed models via knowledge distillation and pruning retain robustness on out-of-distribution (OOD) data compared to full-sized pre-trained language models (PLMs).
• To analyze why compressed models generalize poorly on OOD data, particularly their tendency to overfit on easy or shortcut samples.
• To develop a general-purpose regularization framework that enhances OOD generalization during model compression while preserving in-distribution performance.
• To quantify the impact of compression levels on bias and robustness, especially in relation to sample difficulty.

Proposed method

• Estimate sample difficulty using prediction variance across multiple sub-networks derived from pruned versions of the teacher model.
• Use this uncertainty measure as a sample-specific weighting factor to smooth knowledge distillation and fine-tuning processes.
• Apply instance-wise smoothing during distillation and pruning by adjusting loss weights based on estimated sample uncertainty.
• Integrate uncertainty-based regularization into both knowledge distillation and iterative magnitude pruning pipelines.
• Train compressed models with uncertainty-aware loss functions to reduce overfitting on easy/shortcut samples and improve learning on hard samples.
• Evaluate the framework across multiple NLU tasks using standard in-distribution and OOD test sets, including HANS and textual entailment benchmarks.

Experimental results

Research questions

• RQ1 Are distilled and pruned models as robust as their PLM counterparts on out-of-distribution (OOD) test sets?
• RQ2 How does varying the level of compression affect the OOD generalization and bias of compressed models?
• RQ3 Can sample uncertainty be leveraged to regularize model compression and improve generalization across samples of varying difficulty?
• RQ4 Does uncertainty-based smoothing reduce overfitting on easy/shortcut samples while improving performance on hard samples?

Key findings

• RMC improves OOD generalization accuracy on the HANS benchmark by up to 6.9 percentage points over vanilla fine-tuning in the MNLI task, with consistent gains across all sparsity levels.
• The RMC framework reduces the relative bias metric $\mathcal{F}_{\text{bias}}$ by 10% over vanilla tuning in iterative magnitude pruning and by 10.4% in knowledge distillation.
• On the MNLI task, RMC reduces the performance gap between hard and easy samples by 11.3% under knowledge distillation and 10.6% under pruning at 40% sparsity.
• For the QQP task, RMC improves performance on both in-distribution development sets and OOD test sets simultaneously, demonstrating no trade-off with standard performance.
• Compressed models consistently underperform full PLMs on OOD data, with significant overfitting observed on shortcut/easy samples and poor generalization on hard samples.
• The proposed uncertainty-based regularization effectively mitigates overfitting to easy samples and enhances learning on hard samples, improving overall robustness.