Organizations: Electrical and Computers Engineering, Ben Gurion University of The Negev · Software and Information Systems Engineering, Ben Gurion University of The Negev
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 Gt for updating the policy πθ and critic Rϕ , 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 3
Figure 5: Task performance Am(λmix) versus the mixing ratio of selected adversarial examples to original training data.
Dataset
RoBERTa Base
Additional Data
Paraphrasing
GNLI
Method
Reward Guided
r=0
r=1
r=2
r=3
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
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ϕ
Accuracy (%)
Full (Ours)
✓
✓
✓
92.60
Random Failure-Mode Policy
✓
×
×
89.70
Heuristic Entropy Policy
✓
×
×
90.45
Frozen Policy and Critic
✓
×
✓
91.20
No Failure-Mode Clustering
×
✓
✓
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 πθ
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, α=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, α=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, α=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 Gt updates the policy πθ and critic Rϕ , enabling adaptive selection of high-impact failure modes across iterations.
τ (Percentile)
Noise σ
Accuracy (%)
60%
0.00
91.94
70%
0.00
92.21
80% ( τ⋆ )
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 τ and reward noise (SNLI, 1:4 gen:orig, α=0.83 ).
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.
Adversarial training for LLMs is one of the most promising methods to reliably improve robustness against adversaries. However, despite significant progress, models remain vulnerable to simple in-distribution exploits, such as rewriting prompts in the past tense or translating them into other languages. We argue that this persistent fragility stems from a fundamental limitation in current adversarial training algorithms: they minimize adversarial loss on their training set but inadequately cover the data distribution, resulting in vulnerability to seemingly simple attacks. To bridge this gap, we propose Distributional Adversarial Training, DAT. We leverage Diffusion LLMs to approximate the true joint distribution of prompts and responses, enabling generation of diverse, high-likelihood samples that address generalization failures. By combining optimization over the data distribution provided by the diffusion model with continuous adversarial training, DAT achieves substantially higher adversarial robustness than previous methods.
Large language models are sensitive to minor prompt perturbations, yet existing robustness methods usually enforce consistency at the whole-sequence level. This holistic view can hide an important failure mode: a perturbed response may remain globally similar to the clean one while drifting on a critical entity, relation, or conclusion. We introduce S2R2, a segment-level framework for robust LoRA fine-tuning. S2R2 decomposes clean and perturbed generations into semantic segments, aligns them with an optimal-transport objective, and penalises the segments with the largest meaning drift. To connect this output-side objective with model adaptation, we add an adapter-stability regulariser motivated by segment-level attention reallocation, using LoRA norm control as a tractable proxy for limiting perturbation-amplified evidence shifts. A PAC-Bayesian complexity view further explains why controlling adapter growth may support transfer beyond observed perturbations. Experiments on summarisation benchmarks show that S2R2 improves robustness under typographical noise, deletion, synonym replacement, and paraphrasing, while maintaining competitive clean performance and stronger cross-dataset transfer than consistency-based baselines.
Zhuoyun Li, Boxuan Wang, Jinwei Hu +6
School of Computer Science and Informatics, University of Liverpool, UK