Distilling Robustness into Natural Language Inference Models with Domain-Targeted Augmentation (2305.13067v3)

Abstract: Knowledge distillation optimises a smaller student model to behave similarly to a larger teacher model, retaining some of the performance benefits. While this method can improve results on in-distribution examples, it does not necessarily generalise to out-of-distribution (OOD) settings. We investigate two complementary methods for improving the robustness of the resulting student models on OOD domains. The first approach augments the distillation with generated unlabelled examples that match the target distribution. The second method upsamples data points among the training set that are similar to the target distribution. When applied on the task of natural language inference (NLI), our experiments on MNLI show that distillation with these modifications outperforms previous robustness solutions. We also find that these methods improve performance on OOD domains even beyond the target domain.

Improving Robustness in Knowledge Distillation Using Domain-Targeted Data Augmentation

The paper "Improving Robustness in Knowledge Distillation Using Domain-Targeted Data Augmentation" by Joe Stacey and Marek Rei introduces innovative strategies to enhance the robustness of knowledge distillation, with a primary focus on handling out-of-distribution (OOD) data in the context of Natural Language Inference (NLI). The authors address a significant challenge in knowledge distillation wherein a student model, distilled from a larger teacher model, struggles to maintain comparable performance in OOD scenarios despite successful in-distribution imitation.

Summary of Methods and Findings

The paper proposes two distinct strategies to improve OOD robustness:

1. Domain-Targeted Data Augmentation: This approach employs a LLM to generate unlabeled, task-specific data from potential OOD domains, which is then used in the distillation process. The aim is to encourage the student model to not only mimic the teacher model on in-distribution data but also on the generated OOD examples. The effectiveness of this method was demonstrated through performance improvements over previous robustness methods on datasets such as MNLI and surprising generalization benefits beyond the targeted domains.
2. Distilled Minority Upsampling (DMU): This technique identifies and up-samples minority examples that challenge prevalent spurious correlations during distillation. The method is complementary to domain-targeted augmentation and specifically enhances performance on harder subsets of data, such as the SNLI-hard dataset.

Technical Insights and Results

• The domain-targeted data augmentation was shown to outperform traditional distillation methods by enhancing OOD performance without requiring labeled OOD data. This is attributed to the fact that generating balanced, task-specific data provides the student models with a broader view of potential data variations they might encounter in real-world applications.
• The incorporation of DMU achieved substantial improvements on datasets with adversarial characteristics, suggesting its strength in addressing biases and improving model fairness. The use of teacher-student ensembles for identifying and learning from minority instances further amplified the benefits of DMU.
• Experiments conducted using various combinations of teacher and student models (TinyBERT, BERT, and DeBERTa) validated the flexibility and effectiveness of the proposed solutions across different architectures.

Implications and Future Directions

The findings of this paper have both practical and theoretical implications. Practically, the proposed methods offer a cost-efficient technique to bolster model robustness which is essential for deploying NLP models in dynamic and diverse real-world environments. Theoretical implications include insights into the role of domain-specific data augmentation in enhancing model generalization and the potential of ensemble methods in refining the distillation process.

Future research could extend these methodologies to other NLP tasks and explore the automated generation of more nuanced OOD data with further advancements in LLMs. Moreover, integrating more sophisticated bias detection and mitigation techniques could further improve the robustness and fairness of distilled models. The paper sets a solid foundation for ongoing advancements in enhancing the robustness of knowledge distillation frameworks.