Learning What to Fail On: Failure-Mode Contextual Bandits for Adversarial Data Curation

Roie Kazoom; Ofir Cohen; Rami Puzis; Asaf Shabtai; Ofer Hadar

cs.CLAug 19, 2026

Learning What to Fail On: Failure-Mode Contextual Bandits for Adversarial Data Curation

Authors: Roie Kazoom, Ofir Cohen, Rami Puzis, Asaf Shabtai, Ofer Hadar

Organizations: Electrical and Computers Engineering, Ben Gurion University of The Negev · Software and Information Systems Engineering, Ben Gurion University of The Negev

Abstract

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.

Figures & tables

Figure 1: Overview of the proposed failure-mode contextual bandit curation framework. Given an input-label pair, retrieval-augmented prompting generates candidate outputs. The target model filters candidates by retaining only incorrect predictions, which are then automatically validated by an LLM judge ensemble and clustered into recurring failure modes. A contextual-bandit policy selects failure modes under an adversarial budget, and selected examples are mixed with original data to retrain the target model. Validation feedback provides reward $G_{t}$ for updating the policy $\pi_{\theta}$ and critic $R_{\phi}$ , enabling adaptive selection of high-impact failure modes across rounds.

Figure 2: Example of the progressive construction of our failure-aware policy. From left to right: using only an LLM for hypothesis generation without retrieval or feedback; adding few-shot retrieval to condition generation on similar failures; and the full reinforcement-guided policy, which combines retrieval, multi-model validation, and reward-based selection to identify and reinject high-impact adversarial updates.

Figure 5: Task performance

A_{m}(\lambda_{\mathrm{mix}})

versus the mixing ratio of selected adversarial examples to original training data.

Dataset	RoBERTa Base	Additional Data	Paraphrasing	GNLI	Method	Reward Guided	$\bm{r=0}$	$\bm{r=1}$	$\bm{r=2}$	$\bm{r=3}$	$\bm{r=4}$
SNLI	88.48%	89.42%	84.73%	-	BGE	No	90.98%	91.17%	91.51%	91.54%	91.55%
				-	BGE	Yes	90.10%	91.21%	92.13%	92.12%	92.13%
				-	BM25	No	90.03%	91.02%	91.14%	91.18%	91.19%
				-	BM25	Yes	90.09%	91.20%	92.00%	92.11%	92.12%
				-	BGE+BM25	No	90.11%	91.19%	91.35%	91.61%	91.68%
				-	BGE+BM25	Yes	90.54%	90.78%	92.33%	92.41%	92.60%

Table 1: Accuracy (%) on each test set under adversarial mixing. Method names list only BGE , BM25 , and BGE+BM25 , denoting their respective generation methods. “Reward-Guided” indicates unanimous LLM validation.

Method	Backbone	Dev FEVER	Dev Acc.	Test FEVER	Test Acc.
GEAR Zhou et al. (2019)	BERT-base	70.69	74.84	67.19	71.60
KGAT Liu et al. (2020)	RoBERTa-large	76.11	78.29	70.38	74.07
WgtSum Tymoshenko and Moschitti (2021)	RoBERTa-large	79.02	81.30	73.44	77.18
BEVERS DeHaven and Scott (2023)	DeBERTa-v2-XL	–	–	77.70	80.24
Ours	RoBERTa-base	81.24	83.17	76.58	79.42
Ours	DeBERTa-v3	82.91	84.76	78.12	81.03

Table 2: Comparison on the FEVER benchmark. We report FEVER score and label accuracy.

Dataset	Method	Filtered?	0-shot	3-shot	6-shot	9-shot
SNLI	BGE	No	87.51%	90.05%	91.55%	91.56%
	BGE	Yes	88.18%	90.69%	92.13%	92.15%
	BM25	No	87.51%	89.69%	91.19%	91.22%
	BM25	Yes	88.18%	90.67%	92.12%	92.11%
	BGE+BM25	No	87.51%	89.98%	91.68%	91.51%
	BGE+BM25	Yes	88.18%	90.71%	92.60%	92.51%

Table 3: Few-shot accuracy (%) of our generation methods on each set. Columns indicate the number of few-shot examples.

Figure 6: Few-shot accuracy of generation methods by dataset.

Appendix figures & tables13 assets

Supplementary material from the paper’s appendix.

Appendix

Dataset	SmolLM2-360M			Qwen3-0.6B			RoBERTa-base			DeBERTa-v3			RoBERTa-large
Dataset	No FT	Para	Ours	No FT	Para	Ours	No FT	Para	Ours	No FT	Para	Ours	No FT	Para	Ours
SNLI	79.42	82.36	84.17	82.95	85.41	87.32	88.48	90.72	92.60	87.91	89.84	91.94	89.37	91.12	93.21
Adversarial NLI	67.18	70.04	72.48	70.93	73.41	75.91	75.04	78.12	80.95	74.26	77.03	79.87	75.88	78.94	81.62
MultiNLI	61.94	64.37	66.73	64.82	67.05	69.48	54.67	68.42	71.99	66.71	69.12	71.42	68.09	70.33	72.86

Appendix

Table 4: Accuracy (%) across different backbone models and training strategies. “No FT” denotes evaluation without fine-tuning, “Paraphrasing” denotes augmentation via paraphrase-based data, and “Ours” denotes reinforcement-guided training. All experiments use fixed 6-shot prompting and an original-to-generated data ratio of

1{:}4

(

r=4

Method	Failure Modes	Adaptive Policy	Critic $R_{\phi}$	Accuracy (%)
Full (Ours)	✓	✓	✓	92.60
Random Failure-Mode Policy	✓	$\times$	$\times$	89.70
Heuristic Entropy Policy	✓	$\times$	$\times$	90.45
Frozen Policy and Critic	✓	$\times$	✓	91.20
No Failure-Mode Clustering	$\times$	✓	✓	90.85
Oracle Failure-Mode Policy	✓	GT	GT	93.40

Appendix

Table 5: Ablation of the contextual-bandit curator on SNLI.

Selection Policy	SNLI	ANLI	MultiNLI
Confidence-Based Policy	90.84	78.12	69.21
Loss-Based Policy	91.02	78.45	69.68
Margin-Based Policy	90.67	77.94	68.97
Learned Bandit Policy $\pi_{\theta}$	92.60	80.95	71.99

Appendix

Table 6: Comparison of heuristic failure-mode policies and the learned contextual-bandit policy.

Method	SNLI	ANLI	MultiNLI
Full method	92.60	80.95	71.99
No bandit, random clusters	89.70	76.50	66.20
No bandit, top-loss clusters	91.02	78.45	69.68
No clustering, per-example selection	90.85	78.20	69.10
No judge validation	91.68	78.91	70.02
No failure filtering	89.15	76.90	66.50

Appendix

Table 7: Ablation study of the proposed failure-mode contextual bandit curation framework.

# Judges	# Examples	Accuracy (%)	Judges
1	16,147	91.02	G
2	9,312	91.49	G + Q
3	6,438	92.13	G + Q + P

Appendix

Table 8: Filtering and accuracy under different judge ensemble sizes (SNLI test, 1:4 gen:orig,

\alpha=0.83

). Judges: G = Gemma-3-27B-IT ( Google Research, 2025 ) , Q = Qwen3-32B ( Qwen Team, 2025 ) , P = Phi-4 ( Microsoft Research, 2025 ) .

Figure 7: Accuracy vs. number of judges (SNLI test,

\alpha=0.83

, 1:4 generated:original). Points are annotated with the number of filtered examples.

# Judges	# Examples	Accuracy (%)	Judges
1	17,982	89.74	S
2	11,436	90.21	S + P
3	7,905	90.96	S + P + Q

Appendix

Table 9: Filtering and accuracy under different small judge ensemble sizes (SNLI test, 1:4 gen:orig,

\alpha=0.83

). Judges: S = SmolLM2-360M, P = Phi-2, Q = Qwen2.5-1.5B.

Figure 8: Pairwise TF-IDF cosine similarity between datasets.

Figure 9: Pairwise BERTScore F1 between datasets.

Figure 10: Comparison of average hypothesis lengths (in characters and words) across datasets: Generated-BM25 , Generated-BGE , SNLI train ( Bowman et al., 2015 ) , SNLI test ( Bowman et al., 2015 ) , Adversarial NLI ( Nie et al., 2019 ) , and Multi-NLI ( Williams et al., 2018 ) .

Figure 11: Retrieval accuracy (%) by similarity metric for BGE+BM25, BM25, and BGE.

Figure 12: Illustration of one failure-mode contextual bandit curation iteration on a Natural Language Inference (NLI) example. Given a premise and label, retrieval-augmented prompting generates candidate hypotheses, which are first filtered by the target model to retain incorrect predictions and then validated by an automated LLM judge ensemble. Validated failures are clustered into failure modes, and a contextual-bandit policy selects which modes to sample for retraining under an adversarial budget. Validation reward $G_{t}$ updates the policy $\pi_{\theta}$ and critic $R_{\phi}$ , enabling adaptive selection of high-impact failure modes across iterations.

$\tau$ (Percentile)	Noise $\sigma$	Accuracy (%)
60%	0.00	91.94
70%	0.00	92.21
80% ( $\tau^{\star}$ )	0.00	92.60
90%	0.00	92.17
80%	0.05	92.41
80%	0.10	92.05

Appendix

Table 10: Sensitivity to selection threshold

\tau

and reward noise (SNLI, 1:4 gen:orig,

\alpha=0.83

Adversarial Robustness

Adversarial robustness studies how well AI systems withstand deliberately crafted inputs, training-data corruption, and other attacks.

See all 479 papers

Where Does Robustness Live? Neuron-Guided Adaptation for Retrieval-Augmented Language Models
Apr 2, 2026Jae O Lee, Jaemin Kim, Sumyeong Ahn +1Retrieval-Augmented Generation Fine-Tuning
1Hanyang University · 2KENTECH
Closing the Distribution Gap in Adversarial Training for LLMs
Feb 16, 2026Chengzhi Hu, Jonas Dornbusch, David Lüdke +2Adversarial Training Adversarial Robustness of Language Models
Technical University of Munich, Germany.
Where Do Prompt Perturbations Break Generation? A Segment-Level View of Robustness in LoRA-Tuned Language Models
May 2, 2026Zhuoyun Li, Boxuan Wang, Jinwei Hu +6Neural Network Robustness Noise Robustness
School of Computer Science and Informatics, University of Liverpool, UK

Apr 2, 2026J/K move · Enter open · S save

Where Does Robustness Live? Neuron-Guided Adaptation for Retrieval-Augmented Language Models

Jae O Lee, Jaemin Kim, Sumyeong Ahn, Seo Yeon Park

1Hanyang University · 2KENTECH

Retrieval-Augmented Generation Fine-Tuning Knowledge Augmentation for Language Models LLM Fine-Tuning Neural Network Robustness Retrieval Robustness+2

Open paper AskPDF arXiv

Retrieval-Augmented Language Models (RALMs) have shown strong potential in knowledge-intensive tasks, yet they remain vulnerable when retrieved contexts are noisy or irrelevant. Robustness against such contexts requires two distinct capabilities: abstention when contexts are uninformative, and selective extraction when relevant evidence is buried in noise. Yet existing methods face two key limitations: they do not train separately for these two capabilities, and they adapt the model at a coarse layer- or module-level granularity, overlooking that only a small subset of neurons is strongly activated for a given input. We propose NeuRIT, a Neuron-guided Robust Instruction-Tuning framework built on a localization-first perspective. NeuRIT mines context-aware neurons associated with relevant and irrelevant context processing, and uses them as anchors to selectively adapt both the identified neuron groups and the layers in which they concentrate. NeuRIT then performs two-stage instruction tuning that teaches complementary behaviors: suppress generation when there is nothing to extract, and extract relevant evidence when there is. NeuRIT consistently outperforms strong baselines across diverse QA benchmarks and generator backbones. Our code is available at https://github.com/HYU-ARK-Lab/NeuRIT.

Abstract

Adversarial Robustness

Explore similar work

Where Does Robustness Live? Neuron-Guided Adaptation for Retrieval-Augmented Language Models

Closing the Distribution Gap in Adversarial Training for LLMs

Where Do Prompt Perturbations Break Generation? A Segment-Level View of Robustness in LoRA-Tuned Language Models