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

TL;DR

The paper addresses the robustness gap in knowledge distillation for natural language inference by introducing two orthogonal strategies: domain-targeted augmentation (DTA) that uses language-model generated, unlabelled data to shape teacher distributions toward target domains, and distilled minority upsampling (DMU) that upsamples minority/minority-like examples identified via the student (or ensemble) during distillation. By distilling from teachers (including ensembles) and aligning the student’s predictions to teacher distributions on generated data, the authors achieve improved out-of-distribution generalization on MNLI and related NLI tasks, while also enhancing performance on hard minority cases such as SNLI-hard. DTA and DMU prove complementary, and their combination, often with teacher ensembles, yields robust zero-shot performance across MNLI-matched/mismatched and SNLI-hard, and even improves robustness to known biases like the NLI word-overlap heuristic via WOA. A practical contribution is showing that generated, domain-targeted data can substantially boost OOD robustness, albeit with cost considerations for large-scale data generation; collectively, the work advances domain-aware robustness in KD for NLI and suggests avenues for broader applicability across tasks. $Loss$ formulations for distillation, including ensemble versions, are used to drive probability-matching between student and teacher predictions, enabling robust knowledge transfer under distribution shift.

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.

Figures (5)

• Figure 1: The full process for applying our DTA and DMU methods together. This diagram includes an explanation of how ensembles can be used in both DTA (with an ensemble of student models) and DMU (with an ensemble of student and/or teacher models).
• Figure 2: Our generator model is asked to create a sentence (premise) about a specified genre, before being asked to create a hypothesis that is either implied by the premise, contradicts the premise, or is neutral with respect to the premise. As the hypotheses generated are not faithful to the desired labels (as with this example), we use these examples as unlabelled data during knowledge distillation.
• Figure 3: The process for generating augmented data for our word-overlap augmentation (WOA). In step 4, the model is asked if both the premise and the hypothesis are mostly coherent sentences. In this step, the premise-hypothesis pair is only added to our augmented dataset if the model answers 'yes' for both the premise and the hypothesis. Finally, the sentence pair is only included if the model answers 'no' to the final question in step 5.
• Figure 4: Prompts used to generate hypotheses for MNLI, where <Premise> contains the premise generated by the generator model.
• Figure 5: Prompts used to generate data hypotheses for SNLI, where <Premise> contains the premise generated by the generator model.