Improving Black-box Robustness with In-Context Rewriting

Kyle O'Brien, Nathan Hoyen Ng, Isha Puri, Jorge Mendez-Mendez, Hamid Palangi, Yoon Kim, Marzyeh Ghassemi, Thomas Hartvigsen

Abstract: Machine learning models for text classification often excel on in-distribution (ID) data but struggle with unseen out-of-distribution (OOD) inputs. Most techniques for improving OOD robustness are not applicable to settings where the model is effectively a black box, such as when the weights are frozen, retraining is costly, or the model is leveraged via an API. Test-time augmentation (TTA) is a simple post-hoc technique for improving robustness that sidesteps black-box constraints by aggregating predictions across multiple augmentations of the test input. TTA has seen limited use in NLP due to the challenge of generating effective natural language augmentations. In this work, we propose LLM-TTA, which uses LLM-generated augmentations as TTA's augmentation function. LLM-TTA outperforms conventional augmentation functions across sentiment, toxicity, and news classification tasks for BERT and T5 models, with BERT's OOD robustness improving by an average of 4.48 percentage points without regressing average ID performance. We explore selectively augmenting inputs based on prediction entropy to reduce the rate of expensive LLM augmentations, allowing us to maintain performance gains while reducing the average number of generated augmentations by 57.74\%. LLM-TTA is agnostic to the task model architecture, does not require OOD labels, and is effective across low and high-resource settings.

Submission Length: Regular submission (no more than 12 pages of main content)

Changes Since Last Submission:

The primary changes since the last submission are as follows:
- Prominent modifications to the Introduction, Experiment Setup, and Limitations sections regarding the imprecise notion of OOD. The Limitation section covers this topic in the greatest detail and cites alternative approaches for future work that can more precisely measure the effect that shifting distribution properties have on LLM-TTA performance.
- Added links to the project's codebase, trained models, and data for reproducibility.

Code: https://github.com/Kyle1668/LLM-TTA

Assigned Action Editor: ~Kevin_Swersky1

Submission Number: 2192