# Improving Black-box Robustness with In-Context Rewriting

arXiv: https://arxiv.org/abs/2402.08225 · HTML: https://mlsanity.com/paper/2402.08225/

## Key facts

- Cited by 8 papers in the map (data through 14 Oct 2026).
- Builds on 73 of its 96 references that are in the map.
- First posted on arXiv 13 Feb 2024.

## Builds on (most cited first, 30 of 73)

- [BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding](https://mlsanity.com/paper/1810.04805/) (2018)
- [Language Models are Few-Shot Learners](https://mlsanity.com/paper/2005.14165/) (2020)
- [Exploring the Limits of Transfer Learning with a Unified Text-to-Text Transformer](https://mlsanity.com/paper/1910.10683/) (2019)
- [Llama 2: Open Foundation and Fine-Tuned Chat Models](https://mlsanity.com/paper/2307.09288/) (2023)
- [DistilBERT, a distilled version of BERT: smaller, faster, cheaper and lighter](https://mlsanity.com/paper/1910.01108/) (2019)
- [SQuAD: 100,000+ Questions for Machine Comprehension of Text](https://mlsanity.com/paper/1606.05250/) (2016)
- [Benchmarking Neural Network Robustness to Common Corruptions and Perturbations](https://mlsanity.com/paper/1903.12261/) (2019)
- [The Pile: An 800GB Dataset of Diverse Text for Language Modeling](https://mlsanity.com/paper/2101.00027/) (2020)
- [Character-level Convolutional Networks for Text Classification](https://mlsanity.com/paper/1509.01626/) (2015)
- [HuggingFace's Transformers: State-of-the-art Natural Language Processing](https://mlsanity.com/paper/1910.03771/) (2019)
- [The Many Faces of Robustness: A Critical Analysis of Out-of-Distribution Generalization](https://mlsanity.com/paper/2006.16241/) (2020)
- [Pythia: A Suite for Analyzing Large Language Models Across Training and Scaling](https://mlsanity.com/paper/2304.01373/) (2023)
- [WILDS: A Benchmark of in-the-Wild Distribution Shifts](https://mlsanity.com/paper/2012.07421/) (2020)
- [Certified Adversarial Robustness via Randomized Smoothing](https://mlsanity.com/paper/1902.02918/) (2019)
- [Unsupervised Data Augmentation for Consistency Training](https://mlsanity.com/paper/1904.12848/) (2019)
- [What Does BERT Look At? An Analysis of BERT's Attention](https://mlsanity.com/paper/1906.04341/) (2019)
- [EDA: Easy Data Augmentation Techniques for Boosting Performance on Text Classification Tasks](https://mlsanity.com/paper/1901.11196/) (2019)
- [Improving Neural Machine Translation Models with Monolingual Data](https://mlsanity.com/paper/1511.06709/) (2015)
- [A Survey of Uncertainty in Deep Neural Networks](https://mlsanity.com/paper/2107.03342/) (2021)
- [Measuring Robustness to Natural Distribution Shifts in Image Classification](https://mlsanity.com/paper/2007.00644/) (2020)
- [Memory-Based Model Editing at Scale](https://mlsanity.com/paper/2206.06520/) (2022)
- [ToxiGen: A Large-Scale Machine-Generated Dataset for Adversarial and Implicit Hate Speech Detection](https://mlsanity.com/paper/2203.09509/) (2022)
- [Large Language Models Are State-of-the-Art Evaluators of Translation Quality](https://mlsanity.com/paper/2302.14520/) (2023)
- [A Survey of Data Augmentation Approaches for NLP](https://mlsanity.com/paper/2105.03075/) (2021)
- [Nuanced Metrics for Measuring Unintended Bias with Real Data for Text Classification](https://mlsanity.com/paper/1903.04561/) (2019)
- [MEMO: Test Time Robustness via Adaptation and Augmentation](https://mlsanity.com/paper/2110.09506/) (2021)
- [Med-BERT: pre-trained contextualized embeddings on large-scale structured electronic health records for disease prediction](https://mlsanity.com/paper/2005.12833/) (2020)
- [A Comprehensive Survey on Test-Time Adaptation under Distribution Shifts](https://mlsanity.com/paper/2303.15361/) (2023)
- [Pretrained Transformers Improve Out-of-Distribution Robustness](https://mlsanity.com/paper/2004.06100/) (2020)
- [Unsolved Problems in ML Safety](https://mlsanity.com/paper/2109.13916/) (2021)

## Built on by (most cited first, 8 of 8)

- [A Survey on Medical Large Language Models: Technology, Application, Trustworthiness, and Future Directions](https://mlsanity.com/paper/2406.03712/) (2024)
- [A Comprehensive Survey on the Trustworthiness of Large Language Models in Healthcare](https://mlsanity.com/paper/2502.15871/) (2025)
- [SLOT: Sample-specific Language Model Optimization at Test-time](https://mlsanity.com/paper/2505.12392/) (2025)
- [Rephrasing natural text data with different languages and quality levels for Large Language Model pre-training](https://mlsanity.com/paper/2410.20796/) (2024)
- [Global Contrastive Training for Multimodal Electronic Health Records with Language Supervision](https://mlsanity.com/paper/2404.06723/) (2024)
- [On Adversarial Robustness and Out-of-Distribution Robustness of Large Language Models](https://mlsanity.com/paper/2412.10535/) (2024)
- [What are the Right Symmetries for Formal Theorem Proving?](https://mlsanity.com/paper/2605.22257/) (2026)
- [Toward Robust In-Context Learning: Leveraging Out-of-distribution Proxies for Target Inaccessible Demonstration Retrieval](https://mlsanity.com/paper/2606.00014/) (2026)

## Abstract (arXiv)

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.