Improving Black-box Robustness with In-Context Rewriting

Fuente: arXiv
Saved in:
Bibliographic Details
Main Authors: O'Brien, Kyle, Ng, Nathan, Puri, Isha, Mendez, Jorge, Palangi, Hamid, Kim, Yoon, Ghassemi, Marzyeh, Hartvigsen, Thomas
Format: Preprint
Published: 2024
Subjects:
Online Access:
Tags: Add Tag
No Tags, Be the first to tag this record!
_version_ 1866914898433277952
author O'Brien, Kyle
Ng, Nathan
Puri, Isha
Mendez, Jorge
Palangi, Hamid
Kim, Yoon
Ghassemi, Marzyeh
Hartvigsen, Thomas
author_facet O'Brien, Kyle
Ng, Nathan
Puri, Isha
Mendez, Jorge
Palangi, Hamid
Kim, Yoon
Ghassemi, Marzyeh
Hartvigsen, Thomas
contents 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. We share our data, models, and code for reproducibility.
format Preprint
id arxiv_https___arxiv_org_abs_2402_08225
institution arXiv
publishDate 2024
record_format arxiv
spellingShingle Improving Black-box Robustness with In-Context Rewriting
O'Brien, Kyle
Ng, Nathan
Puri, Isha
Mendez, Jorge
Palangi, Hamid
Kim, Yoon
Ghassemi, Marzyeh
Hartvigsen, Thomas
Machine Learning
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. We share our data, models, and code for reproducibility.
title Improving Black-box Robustness with In-Context Rewriting
topic Machine Learning
url https://arxiv.org/abs/2402.08225