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

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 unlabeled 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.

## introduction large pre-trained language models can achieve impressive performance across a range of natural language understanding tasks @xcite @xcite . however, as performance has increased, so has the number of model parameters @xcite . while large models can be impractical for many applications, knowledge distillation can be used to reduce their size @xcite @xcite . during distillation, a smaller student model is trained to mimic the behaviour of a more complex teacher model on its training data, often improving the performance of the student model on in-distribution examples. however, this does not necessarily lead to robust improvements that generalise to out-of-distribution settings @xcite @xcite . we investigate two methods for improving out-of-distribution performance of the resulting student models: 1) augmenting the distillation by generating new unlabelled task-specific examples that match the target distribution, and 2) upsampling examples among the training data that are similar to the target distribution. we show that these two approaches are orthogonal and can be effectively combined together. we use language models (lm) and a multistep prompting process to generate additional unlabelled examples that target a particular task and domain. while labels for these examples can also be generated @xcite , our experiments show that the lm-generated labels are unreliable and lead to poor performance when used directly for supervised training of the student model. instead, we use these examples to gather predicted probability distributions from the teacher models, then optimise the student models to predict similar distributions during distillation. this approach overcomes the issue of noisy labels and manages to considerably improve student model performance on out-of-distribution examples. in contrast to prior work on generating in-domain examples using the training set @xcite @xcite , this study is the first to investigate the use of generated data to target specific out-of-distribution domains. while domain-targeted augmentation of the data improves performance in many settings, we found that it has little effect on minority examples 2 . therefore, we investigate an additional method of upsampling minority examples during distillation, which substantially improves the performance of the student model on adversarial nli test sets. both of these methods can be combined together to improve robustness across a range of different nli settings. while most distillation is performed with a single teacher model, we also experiment with these methods by distilling from an ensemble of models. ensembles can be used to better identify minority examples, while also increasing the general robustness of the teacher predictions. to the best of our knowledge, this is the first work to investigate model ensembles for better identification of minority examples. we evaluate the distillation methods on the task of natural language inference (nli). in particular, we aim to improve the robustness of models trained on snli @xcite and evaluated on mnli @xcite (and vice versa) -a setting where prior work has consistently found negative results or limited improvements @xcite @xcite @xcite @xcite . we find that our simple but novel approach proves to be highly effective, combining the strengths of both llms and classification models. ## methods given either a large teacher model or an ensemble of teacher models, we aim to distil these models into a single student model that will perform well on different, out-of-distribution datasets, while also performing well on the in-distribution data used to train the teacher model. in the case of mnli, our out-of-distribution data consists of multiple, different domains. by generating additional data for some of these domains (our target domains), we can measure how much performance improves on both the target domains and other out-of-distribution data comprised of different domains. ## conclusion we introduce domain-targeted augmentation and dmu as two methods to improve out-ofdistribution robustness in nli. in the case of the domain-targeted augmentation, using the additional, generated examples during knowledge distillation proves to be a highly effective technique, outperforming all previous work that measures outof-distribution robustness on mnli. not only do we find that performance is better on the targeted domains, but performance is also better for domains that were not included in the augmented data. we also find that our dmu method produces substantial improvements on snli-hard, helping the student model to make better predictions for minority examples. using ensembles can help both methods, improving how minority examples are identified for dmu, and improving the teacher distributions for our domain-targeted augmentation. we also find that our woa method can improve robustness on hans, showing that using unlabelled data during distillation can also target specific, known dataset biases.