Search papers, labs, and topics across Lattice.
This paper presents a novel framework for adversarial data curation that leverages failure-mode contextual bandits to enhance the robustness of natural language understanding models. By generating candidate examples through retrieval-augmented prompting and employing an LLM judge ensemble for validation, the method adaptively selects failure modes for retraining based on a stochastic policy that balances robustness gains and data costs. The approach significantly improves the accuracy of RoBERTa-base on multiple benchmarks, demonstrating its effectiveness over previous adversarial augmentation techniques while also providing a theoretical foundation for its benefits in reducing shortcut-aligned gradient contributions.
Failure-mode contextual bandits can boost model accuracy by over 4% on standard benchmarks while eliminating the need for additional human annotation.
We introduce a failure-aware adversarial retrieval-augmented framework for improving robustness in natural language understanding. Rather than selecting synthetic examples with a fixed reward threshold, our method formulates adversarial data curation as a failure-mode contextual bandit problem. Candidate examples are generated with retrieval-augmented prompting, filtered by the current target model, automatically validated by an LLM judge ensemble, and clustered into recurring failure modes. A stochastic policy then selects which failure modes to sample for retraining, and is updated using validation-based reward that balances robustness gains, forgetting, and data cost. This makes the data curator itself the learning agent, enabling adaptive selection of the most useful model failures across training rounds. On standard benchmarks, our approach improves RoBERTa-base accuracy from 88.48% to 92.60% on SNLI, from 75.04% to 80.95% on ANLI, and from 54.67% to 71.99% on MultiNLI, while consistently outperforming prior adversarial augmentation methods. We further demonstrate transfer to FEVER fact verification, achieving up to 79.86\% FEVER score and 82.45\% accuracy with RoBERTa-large. Finally, we provide a theoretical interpretation showing that, under stated assumptions, failure-mode sampling can reduce shortcut-aligned gradient contributions while inducing bounded distributional drift. By combining retrieval, automated validation, contextual-bandit failure selection, and controlled adversarial retraining, our framework enables scalable robustness improvement without additional human annotation.
