PaperlayerView exact revision on arXiv

Source-linked AI summary

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

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

arXiv:2608.18681v1cs.CL

TL;DR

NLI and other supervised NLP models remain brittle on adversarial or out-of-domain examples, often relying on spurious lexical cues. This paper adaptively curates validated failure modes with a contextual bandit, improving robustness across NLI benchmarks and FEVER while using less data than large untargeted synthetic corpora.

  • Problem

    Supervised NLP models remain brittle on adversarial or out-of-domain examples, often relying on spurious lexical cues despite progress.

  • Method

    A failure-mode contextual bandit clusters validated model errors and learns which modes to sample for retraining without additional human annotation.

  • Results

    Across NLI benchmarks and FEVER, failure-mode curation improves robustness; RoBERTa-large reaches 79.86% FEVER score.

  • Takeaways & Limitations

    Prioritizing validated, model-specific failure modes is more effective than simply adding generated examples and uses substantially less data than untargeted synthetic corpora.

  • Takeaways & Limitations

    The theoretical interpretation depends on stated assumptions, while retrieval-weight calibration uses 1,000 SNLI training examples.

Abstract

from arXiv · show

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.

1 Introduction

The paper presents failure-mode contextual bandits for adaptive, annotation-free adversarial data curation in robust NLI and broader NLP tasks. Its framework combines retrieval, generation, validation, failure clustering, and bandit selection, with theory linking sampling to reduced shortcut reliance and bounded drift.

  • Motivation: NLI supports tasks including question answering, summarization, dialogue systems, and fact verification, yet supervised NLP models remain vulnerable to adversarial and out-of-domain examples.The introduction frames robustness limitations as a broad supervised-NLP problem, not only an NLI issue.
  • Framework: The proposed framework learns which validated model-failure modes to sample for retraining without additional human annotation.This casts adversarial data curation as a failure-mode contextual bandit problem.
  • Methodology: The adaptive pipeline combines label-balanced retrieval, LLM candidate generation, target-model failure filtering, automated judge validation, failure-mode clustering, and contextual-bandit selection under a data budget.These components form the paper’s end-to-end curation process.
  • Empirical Evaluation: Empirical evaluation compares failure-mode bandit curation against static adversarial augmentation, reward-threshold filtering, retrieval-only selection, and untargeted synthetic-data baselines.The stated evaluation focuses on robustness and data efficiency.
  • Theoretical Interpretation: Under stated assumptions, failure-mode sampling can reduce shortcut-aligned gradient contributions while preserving core-feature contributions, with bounded distributional drift and policy distortion from bounded utility noise.The theoretical interpretation connects the sampling strategy to controlled optimization and distributional effects.

2 Background and Related Work

The section situates the framework within efforts to improve NLI robustness through adversarial example generation, retrieval-conditioned prompting, and adaptive training-data selection. Existing benchmarks and automated methods expose model weaknesses, while reinforcement-learning and bandit formulations motivate learning which examples to select.

  • Background and Related Work: SNLI and MultiNLI support large-scale supervised NLI training, while ANLI uses human-and-model-in-the-loop collection to obtain harder adversarial examples.These datasets still require substantial annotation effort and leave systematic failures unresolved.
  • Adversarial and Synthetic Example Generation: Automated adversarial pipelines generate constraint-violating or syntactically controlled examples to expose weaknesses and improve robustness or out-of-domain transfer.Prior work includes logical-constraint violations, model-in-the-loop example discovery, and syntactically controlled transformations.
  • Retrieval for Few-Shot Prompting: Retrieved few-shot context controls generated examples’ label distribution and semantic structure, with dense and lexical retrieval providing complementary signals.BGE supplies semantic similarity, whereas BM25 captures surface overlap and exact lexical cues.
  • Reinforcement Learning, Bandits, and Data Selection: Reinforcement-learning and curriculum-learning research motivates policies that select or weight training examples to improve model training.Related work includes learned sample filters and reinforcement-guided curricula for structured prediction.
  • Reinforcement Learning, Bandits, and Data Selection: Recent data-selection methods formulate finetuning as a sequential decision problem in which an agent chooses data subsets using validation rewards.This formulation directly motivates treating adaptive data curation as a decision-making problem.

3 Methodology

The methodology builds a failure-aware adversarial curation loop that retrieves balanced contexts, generates and validates candidates, clusters failures, and uses a contextual bandit to select retraining data. Its policy optimizes validation improvement while penalizing forgetting and computational cost, replacing fixed reward thresholds with learned utility estimation and stochastic selection.

  • Curation Pipeline: Each iteration performs label-balanced retrieval, LLM candidate generation, target-model failure filtering, automated validation, failure-mode clustering, and contextual-bandit selection for retraining.This six-step loop focuses curation on validated examples exposing current model weaknesses.
  • Failure-Mode Construction: Validated failures are clustered without human labels into recurring modes such as lexical shortcuts, negation errors, entity mismatches, numerical reasoning failures, and contradiction confusion.The resulting state summarizes each mode’s difficulty and diversity using loss, entropy, margin, label distribution, retrieval, agreement, novelty, and prior reward statistics.
  • Contextual-Bandit Policy: The trainable stochastic policy selects failure modes from cluster states, sampling during training for exploration and optionally using greedy probabilities at evaluation.An action of 1 selects a cluster for retraining, subject to the adversarial budget.
  • Policy Optimization: The learned utility estimator Rϕ guides policy learning and variance reduction rather than directly keeping or discarding examples, while πθ governs data selection.The critic and policy are updated after each retraining round, eliminating fixed reward-threshold tuning through validation-driven optimization.
  • Stabilization and Theory: Moderate mixing of original and selected adversarial data stabilizes retraining against catastrophic forgetting, with hybrid retrieval performing strongest near λmix = 1/4.The framework’s theoretical interpretation states that failure-focused sampling reduces shortcut-aligned gradient contributions while preserving core-feature contributions under stated assumptions.

4 Evaluation and Results

The evaluation uses automatically generated and validated adversarial examples to test failure-mode contextual-bandit curation on natural language inference and fact verification benchmarks without additional human annotation. The pipeline mines failures from RoBERTa-base-SNLI, validates candidates through unanimous judge agreement, and adaptively selects clustered failure modes.

  • Evaluation setup: Experiments evaluate the failure-mode contextual-bandit curation pipeline on standard natural language inference and fact verification benchmarks without additional human annotation.Both adversarial example generation and validation are automated.
  • Evaluation setup: RoBERTa-base-SNLI, a 125M-parameter model fine-tuned on SNLI, supplies failures for clustering and adaptive adversarial retraining.The target model’s failures are mined, clustered into failure modes, and used for retraining.
  • Evaluation setup: Adversarial hypotheses are generated with LLaMA-4-Scout-17B-16E-Instruct using label-balanced contexts retrieved through semantic, lexical, or hybrid BGE+BM25 retrieval.The retrieved context conditions generation for each input.
  • Evaluation setup: Candidates are retained only when Gemma-3-27B-IT, Phi-4, and Qwen3-32B unanimously agree with the intended gold label.The ensemble validation stage reduces label noise before failure-mode clustering and policy selection.
  • Evaluation setup: The curator uses a contextual-bandit policy over validated failure-mode clusters, while a lightweight MLP critic predicts each selected mode’s expected validation reward.Cluster states include size, target-model loss, entropy, margin, label distribution, retrieval score, judge agreement, novelty, and previous reward.

RoBERTa

RoBERTa performance improves on FEVER and across NLI datasets as the framework transfers beyond NLI and uses stronger retrieved few-shot context. Validated BGE reaches its strongest reported results in the 9-shot setting, while 6-shot retrieval already matches or exceeds prior adversarial mixing.

  • FEVER transfer: 79.86% FEVER score and 82.45% accuracy are achieved by RoBERTa-large, while RoBERTa-base reaches 76.58% FEVER score and 79.42% label accuracy.The gains extend beyond NLI, and lightweight models also benefit from failure-aware curation.
  • Few-shot retrieval: 92.15% on SNLI, 80.26% on ANLI, and 71.15% on MultiNLI are reached by validated BGE in the 9-shot setting.Increasing retrieved few-shot context improves performance across all three datasets.
  • Few-shot retrieval: 6-shot retrieval already matches or exceeds the strongest adversarial mixing results, highlighting the importance of retrieval quality and validated failure selection.BM25 follows similar but slightly weaker trends than validated BGE.

5 Conclusion and Future Work … B.2 Boundedness of Failure-Aware Bandit Updates

The paper presents failure-mode contextual bandits as an adaptive framework for adversarial data curation, selecting validated recurring model failures for retraining while balancing robustness, forgetting, and data cost. Its theoretical analysis shows that failure-mode sampling can reduce shortcut-aligned gradients, while mixture updates and bounded reward noise keep distributional and allocation changes controlled.

  • 5 Conclusion and Future Work: The framework clusters validated model errors into recurring failure modes and learns which modes to sample for retraining instead of using a fixed reward threshold.The curator becomes an adaptive policy receiving validation-based feedback while balancing robustness gains, forgetting, and data cost.
  • B.1 Bias Reduction via Failure-Mode Curation: The method generates adversarial candidates, filters examples that fool the target model, validates them with automated judges, clusters failures, and samples modes with a contextual-bandit policy.
  • B.1 Bias Reduction via Failure-Mode Curation: Failure-focused training is analyzed with a decomposition into task-relevant core information and spurious features that may be unreliable under distribution shift.The decomposition is used for analysis and is not assumed to be explicitly available to the model.
  • B.1 Bias Reduction via Failure-Mode Curation: Under the stated assumptions, failure-mode sampling reduces the shortcut-aligned gradient contribution relative to training only on the original distribution.The reduction is guaranteed when the policy approximation error δs is smaller than the original-distribution expected shortcut gradient contribution.
  • B.2 Boundedness of Failure-Aware Bandit Updates: The effective training distribution mixes original and policy-selected adversarial examples to prevent uncontrolled distributional drift across iterations.The update is analyzed through an abstract mixture controlled by coefficient η.
  • B.2 Boundedness of Failure-Aware Bandit Updates: Proposition A.3 shows that reward noise bounded by ε induces only bounded distortion in the policy-induced allocation over failure modes.Together with the distribution-drift result, this characterizes the failure-aware bandit update as operating in a controlled regime.

B.3 Backbone and Training Strategy Ablation

Across five backbone architectures, reinforcement-guided adversarial training consistently outperforms no fine-tuning and paraphrase augmentation, with benefits preserved across model scales. The results support targeted failure-driven mining as a general and scalable training strategy, while leaving generator-verifier scaling and component contributions for future study.

  • Backbone and Training Strategy: Models without fine-tuning perform worst in every configuration, while paraphrase augmentation provides only moderate gains over No FT.The limitation of paraphrasing is especially apparent on challenging benchmarks such as Adversarial NLI.
  • Backbone and Training Strategy: Reinforcement-guided training achieves the strongest performance across all evaluated backbone architectures and datasets.The experiment spans SmolLM2-360M, Qwen3-0.6B, RoBERTa-base, DeBERTa-v3, and RoBERTa-large, comparing No FT, Paraphrasing, and Ours.
  • Backbone and Training Strategy: 92.60% is RoBERTa-base accuracy with the proposed method on SNLI, versus 90.72% under paraphrasing.The method also improves SmolLM2-360M by +1.81 points and Qwen3-0.6B by +1.91 points on SNLI.
  • Backbone and Training Strategy: The framework preserves its benefits across model scales, helping smaller and medium-sized backbones through targeted adversarial mining rather than excess model capacity.RoBERTa-large achieves higher absolute accuracy, but smaller models also benefit substantially by reshaping training toward informative failure regions.
  • Open Questions: Future work should vary generator and verifier scale and disentangle their individual contributions to robustness gains.Proposed analyses include controlled capacity changes and fixing one component while varying the other to separate generation quality from validation reliability.

B.4 Ablation Study: Effect of the Contextual-Bandit Policy

The ablation study shows that the learned contextual-bandit curator improves SNLI performance over static, frozen, and non-clustered selection strategies. Continual adaptation and recurring failure-mode clustering are important components of the full model.

  • Full Model: 92.60% accuracy is achieved by the full contextual-bandit model on SNLI under the shared backbone, retrieval, validation, budget, and retraining protocol.The full model uses a stochastic policy over failure modes and validation reward balancing robustness gain, forgetting, and data cost.
  • Policy Alternatives: Random and heuristic failure-mode policies underperform the full model, showing that uniform or uncertainty-based static selection is insufficient.The random baseline samples uniformly, while the heuristic ranks failure modes by mean predictive entropy until the adversarial budget is reached.
  • Adaptation: Freezing the policy and critic degrades performance, indicating that continual validation-based adaptation across retraining rounds is important.The frozen variant initializes πθ and Rϕ in the first round and then keeps both fixed.
  • Failure-Mode Clustering: Removing failure-mode clustering further reduces accuracy, confirming that selecting recurring failure types is more effective than selecting isolated validated examples.This baseline retains target-model filtering and automated validation but samples validated failures directly without grouping them into recurring modes.
  • Oracle Comparison: The oracle policy is an upper bound rather than a deployable alternative because it requires separate retraining for each candidate failure mode.It approximates failure-mode utility using observed validation improvement after retraining.

B.5 Ablation with Heuristic Failure-Mode Policies

The ablation isolates failure-mode selection while holding candidate generation, filtering, validation, retrieval, and mixing fixed. The learned contextual-bandit policy outperforms heuristic policies because it uses downstream validation feedback to adapt selection across curation rounds.

  • Ablation setup: The ablation changes only how validated failure modes are selected for retraining, while retaining the same candidates, failure filtering, LLM validation, retrieval weight, and mixing ratio.The full-method settings include α = 0.83 and λmix = 1⁄4.
  • Heuristic policies: Confidence-, loss-, and margin-based policies prioritize low confidence, high supervised loss, and small top-two probability separation, respectively.Each heuristic scores validated failure-mode clusters and selects them in descending order until the adversarial budget is reached.
  • Learned policy: The learned policy uses cluster features including loss, entropy, margin, label distribution, retrieval score, judge agreement, novelty, cluster size, and reward statistics.Its policy is updated using validation reward after retraining, unlike the heuristic policies.
  • Ablation conclusion: Heuristic policies can oversample noisy, redundant, or locally difficult failures, whereas the learned policy adapts to downstream validation feedback across curation rounds.The reported consistent improvement over heuristic policies supports adaptive failure-mode selection over static uncertainty-based selection.

B.6 Component Analysis of Failure-Mode Curation · B.7 Judge Ensemble Configuration

The component analysis shows that the full failure-mode curation pipeline performs best because validation-driven selection, failure filtering, retrieval context, and original-data mixing jointly improve robustness. The judge study selects a three-judge unanimous ensemble, which achieves the highest SNLI accuracy while retaining fewer filtered observations.

  • B.6 Component Analysis of Failure-Mode Curation: 92.60% on SNLI, 80.95% on ANLI, and 71.99% on MultiNLI are achieved by the full method across all benchmarks.The complete pipeline combines retrieval-augmented generation, target-model failure filtering, automated validation, failure-mode clustering, contextual-bandit selection, and controlled original-data mixing.
  • B.6 Component Analysis of Failure-Mode Curation: Random cluster selection performs substantially worse, while top-loss selection improves over random but remains below validation-driven contextual-bandit selection.These results indicate that failure modes differ in retraining usefulness and that simple difficulty-based heuristics are less effective than learned policy selection.
  • B.6 Component Analysis of Failure-Mode Curation: Removing judge validation lowers performance across datasets, while removing target-model failure filtering causes an even larger degradation.The ablations identify automated validation as a safeguard against noisy or incorrectly labeled examples and failure filtering as central to exposing current model errors.
  • B.6 Component Analysis of Failure-Mode Curation: Removing retrieved context reduces performance, and training without original-data mixing also hurts performance.Retrieved few-shot examples provide informative generation context, while original-data mixing supports robustness improvement while limiting forgetting.
  • B.7 Judge Ensemble Configuration: The judge evaluation fixes α = 0.83 and a generated-to-original example ratio of 1:4, then filters examples using unanimous agreement among independent LLM validators.All configurations are evaluated on the SNLI test set after filtering, with accuracy measured on the remaining items.
  • B.7 Judge Ensemble Configuration: 92.13% accuracy is achieved by the three-judge ensemble on 6,438 filtered observations, exceeding the two-judge and single-judge configurations.The two-judge setup reaches 91.49% on 9,312 examples, while the single-judge setup reaches 91.02% on 16,147 examples.
  • B.7 Judge Ensemble Configuration: Gemma-3-27B-IT remains in every configuration, joined by Qwen3-32B for two judges and Phi-4 for three judges.The three-judge configuration is adopted for subsequent experiments.

B.7.1 Evaluation with Small Judge Models

With unanimous agreement required among small judges, the framework retains its robustness benefits while showing lower accuracy and filtering precision than large-model ensembles. Performance improves with ensemble size and degrades gracefully under weaker validation models, supporting low-resource, cost-constrained use.

  • Evaluation protocol: The filtering protocol retained only examples unanimously accepted by all selected judges before measuring classification accuracy on the remaining SNLI test instances.Table 9 evaluates this protocol on SNLI with a 1:4 generated-to-original ratio and α = 0.83.
  • Comparison with large judges: Small judge ensembles produced moderately lower accuracy and weaker filtering precision than large-model ensembles.The comparison is reported against the large-model ensembles in Table 6.
  • Scaling judge ensembles: Performance improved consistently as the small judge ensemble grew, while even one lightweight judge delivered substantial robustness gains.The result indicates that larger small-model ensembles improve evaluation outcomes without eliminating the usefulness of a single judge.
  • Resource-constrained validation: The framework degraded gracefully with weaker validation models, supporting applicability in low-resource and cost-constrained settings.This conclusion follows the observed robustness of the method under small judge models.

B.8 Dataset Comparison

The dataset comparison evaluates lexical and semantic overlap between generated examples and established NLI benchmarks using TF-IDF cosine similarity and BERTScore F1. The hybrid BGE+BM25 set occupies an intermediate position between its single-retriever counterparts while combining benchmark-aligned meaning with varied surface vocabulary.

  • Dataset comparison: The comparison covers seven collections, including SNLI, BGE-generated, BM25-generated, Adversarial NLI, Multi-NLI, and hybrid BGE+BM25-generated data.The analysis first examines frequent non-stopwords, then quantifies similarity across datasets.
  • Lexical overlap: TF-IDF cosine similarity represents each dataset as a TF-IDF vector and computes normalized vector overlap.The similarity is defined as the dot product divided by the product of vector norms.
  • Lexical overlap: 0.0251 with SNLI Train, 0.0188 with SNLI Test, and 0.0150 with Multi-NLI are the hybrid BGE+BM25 set’s TF-IDF similarities, intermediate between BGE-only and BM25-only data.These values come from the 7 × 7 similarity matrix.
  • Semantic overlap: 0.8658 with SNLI Train, 0.8534 with SNLI Test, 0.8458 with Adversarial NLI, and 0.8554 with Multi-NLI are the hybrid set’s BERTScore F1 values.The hybrid scores again fall between the BGE-only and BM25-only pairs.
  • Semantic overlap: Generated datasets show moderate lexical overlap with SNLI Train at cosine similarities around 0.02-0.03 but stronger semantic alignment with SNLI Train and SNLI Test at BERTScore F1 values above 0.85.This indicates varied surface vocabulary alongside well-preserved contextual meaning.

B.9 Generated Dataset Characteristics and Hypothesis Lengths … B.15 Prompt Design for Task-Specific Candidate Generation

The paper characterizes generated candidates, evaluates retrieval and reproducibility choices, and details a failure-mode bandit workflow that remains stable under threshold and reward-noise variation. It also illustrates retrieval and task-specific prompting designs for controlled adversarial candidate generation across NLI and FEVER.

  • B.9 Generated Dataset Characteristics and Hypothesis Lengths: Generated datasets share speculative and gender-related themes across retrieval strategies, with only minor stylistic differences.The most frequent generated terms include “could,” “would,” “woman,” “he,” and “she.”
  • B.9 Generated Dataset Characteristics and Hypothesis Lengths: 98–100 characters (16–17 words) is the average hypothesis length for generated sets, versus ≈37–38 characters (7–8 words) for SNLI annotations.Adversarial NLI instances average ≈64 characters (11 words).
  • B.10 Retrieval Accuracy Across Similarity Metrics: 92.60% (cosine) is the best reported retrieval accuracy, achieved by BGE+BM25, which outperforms BM25 and BGE across all six metrics.BGE+BM25 also achieves 89.85% (dot product), 85.43% (L2), 85.22% (L1), 79.21% (Bray-Curtis), and 79.12% (Canberra).
  • B.12 Illustrative Example: Failure-Mode Bandit Curation for NLI: The curation workflow filters model-induced errors, validates them with an LLM judge ensemble, clusters failure modes, and uses policy rewards after retraining to update selection.The learned policy replaces fixed reward-threshold selection and can be applied beyond NLI to other classification and reasoning tasks.
  • B.13 Sensitivity to Selection Threshold and Reward Noise: Performance remains stable across a broad operating region of selection thresholds and reward-noise levels, with only minor degradation near moderate deviations from τ⋆.The observed plateau around τ⋆ is consistent with the stated bounded-drift analysis under noisy feedback.
  • B.14 Example - Few-Shot Chat Sequence: Few-shot retrieval examples are presented as labeled premise–hypothesis pairs for entailment, neutral, and contradiction, using BGE, optimized BGE+BM25, or BM25 retrieval.The optimized hybrid method linearly interpolates semantic and lexical scores through a tunable alpha parameter.
  • B.15 Prompt Design for Task-Specific Candidate Generation: Task-specific prompts condition generation on target labels, constrain outputs to concise single sentences without narration, and use low-temperature decoding for stable, precise candidates.NLI prompts target entailment, neutral, or contradiction; FEVER prompts target SUPPORTS, REFUTES, or NOT_ENOUGH_INFO and ground claims in evidence.
Loading 2608.18681v1…