Improving Black-box Robustness with In-Context Rewriting

TL;DR

This work tackles the challenge of out-of-distribution robustness in NLP under black-box constraints by introducing LLM-TTA, a test-time augmentation framework that uses large language models to generate faithful input augmentations. It explores two augmentation modes—zero-shot paraphrasing and In-Context Rewriting (ICR)—and adds entropy-based selective augmentation to reduce expensive LLM calls. Across sentiment, toxicity, and news topic tasks with BERT and T5, LLM-TTA generally improves OOD accuracy, with ICR delivering the strongest gains and only modest or no degradation on ID performance; selective augmentation further boosts efficiency by cutting augmentation rates by about 57.74%. The method is architecture-agnostic, does not require OOD labels, and remains effective in both data-scarce and data-rich settings, with data, models, and code shared for reproducibility.

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. We share our data, models, and code for reproducibility.

Figures (7)

• Figure 1: LLM-TTA. In settings where the task model is effectively a black box, we can intervene on the input data to improve robustness. We propose rewriting OOD inputs at test time using an LLM to improve robustness (LLM-TTA). Our experiments find that LLM-TTA improves performance without requiring task model access or OOD labels.
• Figure 2: TTA Inference Steps. This figure shows the three stages of TTA. The process begins with an augmentation function generating multiple altered versions of the current test input. For ICR, the input to the LLM also contains ID examples. The task model then makes predictions over the test input and its augmentation. Lastly, we aggregate the predictions to arrive at a "smoothed” judgment. Standard aggregation methods include mean probability aggregation (demonstrated in this figure) and vote-based aggregation.
• Figure 3: LLM-TTA Prompts. We evaluate LLM-TTA with two prompting methods. "<style_input>" is replaced with the test input and "<style_transfer_exemplars>" with the ID examples during inference. During the course of prompt engineering, we find that instructing the LLM to generate text in specific formats surrounded by brackets and to change the details of the text while preserving semantics leads to the best performance.
• Figure 4: TTA Effectiveness Across Data Scales. This figure shows the absolute improvements in OOD accuracy averaged across shifts and experiment runs with standard deviations. We train five BERT models on 5%, 10%, 20%, 40%, and 80% of the ID training set. We find that LLM-TTA improves robustness across data scales. These results suggest that LLM-TTA can still be helpful for practitioners operating in data-scarce regimes.
• Figure 5: Changed Predictions Across Classes. Results are from BERT with ICR as the TTA augmentation function across all OOD inputs. Variance across classes indicates that TTA affects some classes more than others. TTA can hurt the performance of some classes while improving overall performance.