PaperlayerView exact revision on arXiv

Source-linked AI summary

No One Size Fits All: QueryBandits for Hallucination Mitigation

Nicole Cho, William Watson, Alec Koppel, Sumitra Ganesh, Manuela Veloso

arXiv:2602.20332v1cs.CLcs.AIcs.LG

TL;DR

Most hallucination-mitigation research targets open-weight models, leaving limited evidence for closed-source systems. QueryBandits adaptively selects query-rewrite strategies online, achieving an 87.5% win rate over no rewriting across 16 QA scenarios.

  • Problem

    Most hallucination-mitigation approaches target open-weight models, while evidence on hallucinations in closed-source models remains limited despite their institutional prevalence.

  • Method

    QueryBandits uses linguistic query features and a calibrated composite reward to learn online which rewrite strategy to select per query.

  • Results

    87.5% win rate over NO-REWRITE across 16 QA scenarios, with Thompson Sampling outperforming PARAPHRASE and EXPAND by 42.6% and 60.3%, respectively.

  • Takeaways & Limitations

    Query utility varies by linguistic features, so no single fixed rewrite policy is universally beneficial across queries.

  • Takeaways & Limitations

    The fuzzy-match reward term must be tuned carefully to avoid degrading overall accuracy.

Abstract

from arXiv · show

Advanced reasoning capabilities in Large Language Models (LLMs) have led to more frequent hallucinations; yet most mitigation work focuses on open-source models for post-hoc detection and parameter editing. The dearth of studies focusing on hallucinations in closed-source models is especially concerning, as they constitute the vast majority of models in institutional deployments. We introduce QueryBandits, a model-agnostic contextual bandit framework that adaptively learns online to select the optimal query-rewrite strategy by leveraging an empirically validated and calibrated reward function. Across 16 QA scenarios, our top QueryBandit (Thompson Sampling) achieves an 87.5% win rate over a No-Rewrite baseline and outperforms zero-shot static policies (e.g., Paraphrase or Expand) by 42.6% and 60.3%, respectively. Moreover, all contextual bandits outperform vanilla bandits across all datasets, with higher feature variance coinciding with greater variance in arm selection. This substantiates our finding that there is no single rewrite policy optimal for all queries. We also discover that certain static policies incur higher cumulative regret than No-Rewrite, indicating that an inflexible query-rewriting policy can worsen hallucinations. Thus, learning an online policy over semantic features with QueryBandits can shift model behavior purely through forward-pass mechanisms, enabling its use with closed-source models and bypassing the need for retraining or gradient-based adaptation.

1 INTRODUCTION

QueryBandits introduces a model-agnostic contextual bandit that uses calibrated factuality rewards and query-semantic features to select rewrites online. Its results show contextual adaptation outperforms fixed strategies, while rewrite utility varies by query.

  • Framework: QueryBandits selects among five rewrite strategies using a 17-dimensional linguistic feature vector that maps query properties to online arm choices.The policy balances exploration under uncertainty with exploitation when features are predictive.
  • Reward Modeling for Factuality: 0.973 average ROC–AUC validates the composite reward’s ability to separate correct from incorrect answers across resampling settings.The reward combines LLM-judge, fuzzy-match, and BLEU-1 metrics, with α, β, γ = (0.6, 0.3, 0.1).
  • Contextual Adaptation Wins: 87.5% win rate over NO-REWRITE is achieved by Thompson Sampling across 13 QA benchmarks and 16 scenarios.It also outperforms PARAPHRASE by 42.6% and EXPAND by 60.3%.
  • Interpretable Decision Weights: No single rewrite strategy maximizes reward across all query types because each arm’s effectiveness depends on semantic features.EXPAND is especially effective for queries exhibiting Domain Specialization, unlike SIMPLIFY.
  • Scope & Utility: QueryBandits operates at the input layer as a model-agnostic online policy for closed-source LLMs without modifying model weights.This distinguishes it from mitigation methods that modify internal representations or decoding in open-source models.

2 RELATED WORKS

Related work frames hallucinations as a societal and epistemic failure requiring dedicated mitigation, while mitigation research is moving from post-hoc detection toward preemptive query shaping. Linguistic features provide actionable context for adapting query-rewriting policies.

  • Societal Stakes and Gap of Closed-Source Models: Hallucinations erode societal trust and are framed as an epistemic failure linked to training and evaluation procedures that reward guessing over uncertainty acknowledgment.
  • From Post-hoc Detection to Preemptive Query Shaping: Mitigation research has expanded from post-hoc detection and iterative correction to preemptive grounding, query restructuring, risk estimation, and Rewrite-Retrieve-Read pipelines.
  • Linguistic Features as Actionable Context: 17 key linguistic features identified per query serve as context for the bandit policy, drawing on LLM literature and traditional linguistics.Feature selection prioritizes properties known to affect comprehension for humans and models.

3 METHODOLOGY AND EVALUATION METRICS

QueryBandits formulates query rewriting as a contextual multi-armed bandit, using semantic query features to select among rewrite strategies and a calibrated composite reward. Evaluation measures exploration-adjusted performance, convergence through cumulative regret, and accuracy gains over no rewriting.

  • Context and actions: Five rewrite arms span PARAPHRASE, SIMPLIFY, DISAMBIGUATE, EXPAND, and CLARIFY TERMS strategies.Each arm is implemented through prompt instructions to an LLM.
  • Context and actions: 17 binary features provide query context, enabling contextual bandits to learn when to apply each rewrite strategy.The feature vector captures linguistically motivated properties affecting human and LLM comprehension.
  • Reward model: The bounded reward rt ∈[0, 1] combines GPT-4o correctness, RapidFuzz token-set similarity, and unit-capped BLEU-1 as rt = αsllm + βsfuzz + γsbleu.The weights satisfy α + β + γ = 1 and α, β, γ ≥ 0.
  • Reward validation: 0.9729 macro-average ROC–AUC validates the reward proxy across sample sizes of 5–1000, with the 95% CI lower bound exceeding 0.97 by 150 samples.The selected weights (0.6, 0.3, 0.1) lie within the top 1% Pareto frontier; varying α by ±0.2 shifts AUC by < 0.5%.
  • Evaluation metrics: The evaluation reports exploration-adjusted reward, mean cumulative regret, and win rate versus the NO-REWRITE baseline.These metrics respectively assess exploration–exploitation balance, convergence toward good answers, and how often policy reward strictly exceeds the baseline.

4 EXPERIMENTS

The experiments evaluate QueryBandits across 16 QA scenarios using semantically invariant lexical perturbations, comparing contextual and non-contextual bandits with prompting baselines. Each round extracts query features, selects and applies a rewrite arm, queries gpt-4o-2024-08-06, evaluates reward, and updates the bandit.

  • Pipeline: Each decision round extracts features, selects a rewrite arm, issues the rewritten query to gpt-4o-2024-08-06, computes reward, and updates the bandit.The reward is a scalar rt ∈[0, 1], and the bandit update uses the selected arm and reward.
  • Dataset and Query Construction: 16 QA scenarios are constructed from 13 diverse benchmarks using lexically perturbed but semantically invariant queries.Queries must be answerable by gpt-4o-2024-08-06 before perturbation, with perturbation validity assessed by an LLM-as-judge and n-gram metrics.
  • Dataset and Query Construction: Perturbed queries prevent contextual policies from collapsing to NO-REWRITE, unlike canonical queries where rewriting rarely improved accuracy.The setup targets benchmark memorization arising from pre-training and fine-tuning exposure.

5 RESULTS

QueryBandits improves hallucination-related performance over no-rewrite and static strategies, with contextual methods benefiting from query features. Results also show that rewrite effectiveness varies by query type, while input-layer adaptation improves closed-source GPT-4o performance.

  • Hypothesis 1: QueryBandits and hallucination reduction: 87.5% query-level win rate: contextual Thompson Sampling achieves this against NO REWRITE, with 819.04 exploration-adjusted reward versus 729.20 (∆= −89.84).The comparison spans 13 QA benchmarks, 16 scenarios, and 1,050 queries per dataset.
  • Hypothesis 2: Static rewriting: 819.04 aggregate reward: Contextual TS exceeds PARAPHRASE and EXPAND at 732.39 and 639.25, respectively, while static rewriting never tops a dataset on accuracy.Contextual methods typically gain +6–12 pp over the baseline, including +12.6 on BoolQA and +10.6 on HotpotQA.
  • Hypothesis 3: Contextual versus non-contextual bandits: −5.8 pp, –64.38 reward: Removing the 17-dimensional feature vector lowers Thompson Sampling from 87.5% to 81.7% query-level win rate and from 819.04 to 754.66 reward.On accuracy, Thompson Sampling wins 8/16 scenarios, while LinUCB/LinUCB+KL take the remainder with tie-split scores of 4.5 and 3.5.
  • Hypothesis 4: Query features and reward: 17 linguistic features: Arms show distinct sensitivities, and the same feature can reverse importance across arms; Domain Specialization predicts EXPAND strongly but SIMPLIFY weakly.The passage proposes that EXPAND adds qualifiers or entities for domain-specific questions, whereas SIMPLIFY may remove critical semantics.
  • Hypothesis 5: One strategy for all queries: No single rewrite strategy: Learned feature fingerprints and diverse winning arms show that optimal rewrites vary across query types; SIMPLIFY helps with pragmatic cues but struggles on superlatives.Contextual TS and the LinUCB family split the contextual wins across scenarios.
  • Hypothesis 6: Closed-source model performance: 88.8% (+8.1 pp): QueryBandits lifts GPT-4o performance through input-layer operation, compared with GPT-4o’s 80.7% MC1 and lower reported MC1 values for TruthX (54.2%) and DoLa (32.2%).The passage contrasts QueryBandits’ input-layer approach with methods that modify internal representations or decoding.

6 CONCLUSION · S. M Towhidul Islam Tonmoy, S M Mehedi Zaman, Vinija Jain, Anku Rani, Vipula Rawte, Aman · A APPENDIX

QueryBandits is a plug-and-play contextual online-learning policy that selects among rewrite strategies using lightweight linguistic features. Across 13 QA benchmarks and 16 scenarios, contextual learners substantially outperform No-Rewrite, while non-contextual bandits generally exceed static prompts.

  • 6 CONCLUSION: 0.766 macro-average accuracy versus 0.681 for No-Rewrite across 13 QA benchmarks and 16 scenarios.Contextual TS and the LinUCB family win nearly all benchmarks.
  • 6 CONCLUSION: ∼6–12 pp typical gains over No-Rewrite are reported for contextual learners.The comparison is summarized as a macro-average result in Table 1.
  • 6 CONCLUSION: Contextual TS and the LinUCB family win nearly all benchmarks.These learners use query-level linguistic features as context when selecting rewrite strategies.
  • 6 CONCLUSION: QueryBandits selects among K rewrite strategies to minimize a query’s hallucinatory trajectory.The policy is plug-and-play and operates through lightweight linguistic features as context.
  • 6 CONCLUSION: The framework applies online learning to query rewriting rather than requiring model retraining.Its policy is described as plug-and-play and context-driven.
  • 6 CONCLUSION: Non-contextual bandits generally beat static prompts.The supplied conclusion passage states this comparison without specifying an aggregate metric.

A.1 LIMITATIONS · A.2 ETHICS & SOCIETAL IMPACT.

The paper identifies modeling and causal-interpretation limits in its contextual bandit framework, while emphasizing that query rewriting reduces but does not eliminate hallucinations. It also notes reward misspecification, domain shift, and auditing considerations as practical and societal concerns.

  • A.1 LIMITATIONS: A.1 LIMITATIONS: The framework treats all 17 features as independent and does not model higher-order feature interactions.Future work could measure whether feature combinations jointly exacerbate hallucination.
  • A.1 LIMITATIONS: A.1 LIMITATIONS: Measuring joint feature effects is identified as an avenue for future research.The proposed direction concerns whether combinations of features jointly exacerbate hallucination.
  • A.1 LIMITATIONS: A.1 LIMITATIONS: Feature-arm regression weights do not stipulate a causal relationship.The passage notes that the underlying causal relationships may be highly sophisticated.
  • A.2 ETHICS & SOCIETAL IMPACT.: A.2 ETHICS & SOCIETAL IMPACT.: The method alters inputs rather than model weights and can reduce factually incorrect outputs without eliminating them.Its intervention operates through input modification rather than parameter changes.
  • A.2 ETHICS & SOCIETAL IMPACT.: A.2 ETHICS & SOCIETAL IMPACT.: The stated failure modes are reward misspecification and domain shift.These conditions can undermine the method’s behavior outside the assumptions reflected in its reward or data domain.
  • A.2 ETHICS & SOCIETAL IMPACT.: A.2 ETHICS & SOCIETAL IMPACT.: The authors report error analyses and release prompts to support auditing and replication.They also discuss the societal impact of hallucinations in related works.

A.3 REPRODUCIBILITY STATEMENT.

The paper provides the learning setup, rewrite-arm prompts, contextual-policy features, reward specification, and evaluation-dataset details needed for reproducibility.

  • A.3 REPRODUCIBILITY STATEMENT.: The reproducibility materials specify five rewrite arms, exact system-prompt templates, contextual-policy features, reward components and weights, and dataset protocols.They also document evaluation datasets, splits, preprocessing, dataset-specific details, and licenses.

B DISCUSSION ON RL AND BANDIT METHODS

The discussion motivates contextual bandits over full RL because hallucinations arise as a per-query generation phenomenon, and relates QueryBandits algorithms to established RL methods. It also positions contextual bandits along the spectrum from model-based to model-free approaches.

  • Why bandits vs. full RL?: Hallucinations are characterized as per-query phenomena arising during generation, motivating bandit methods rather than full reinforcement learning for query-level adaptation.For each query, the transformer attends over a fixed context window and computes token probabilities during generation.
  • Algorithm Choices and RL Methods: Thompson sampling has posterior sampling as an RL analogue, while FTRL variants originate from proximal-gradient methods related to PPO.The passage explicitly links posterior sampling to Thompson sampling and FTRL’s methodological origins to PPO.
  • Algorithm Choices and RL Methods: Contextual bandits sit within a spectrum spanning model-based and model-free approaches, alongside CoT training with an explicit value model.This comparison situates contextual bandits relative to reinforcement-learning paradigms rather than presenting them as isolated methods.

C TRUTHFULQA METRICS AND EVALUATION SETUP · D SUMMARY OF BANDITS · D.1 LINUCB

The evaluation centers on TruthfulQA’s MC1 accuracy while distinguishing it from multi-true and generation-based metrics. QueryBandits compares non-contextual and contextual bandit strategies that select query rewrites online from observed rewards, including LinUCB’s confidence-based arm selection.

  • C TRUTHFULQA METRICS AND EVALUATION SETUP: MC1 selects the single true option using the highest completion log probability, with accuracy computed over questions.This is the paper’s primary TruthfulQA regime for comparability across methods.
  • C TRUTHFULQA METRICS AND EVALUATION SETUP: MC2 assigns credit through the normalized total probability given to the set of true answers in multi-true questions.Unlike MC1, MC2 evaluates probability mass over multiple reference answers.
  • C TRUTHFULQA METRICS AND EVALUATION SETUP: Generation evaluates 1–2-sentence free-form answers with GPT-judge, GPT-info, BLEURT, ROUGE, BLEU, and a maxtrue sim −maxfalse sim score.GPT-judge and GPT-info are fine-tuned binary classifiers for truthfulness and informativeness, so their labels are not accuracy and are not directly comparable to generative evaluation.
  • D SUMMARY OF BANDITS: EXP3, FTPL, KL-UCB, and Thompson Sampling form the non-contextual bandit methods, using weights, Gumbel perturbations, KL-divergence bounds, or Gaussian posteriors.The summary also lists contextual extensions including LinearEXP3, LinearFTPL, and LinUCB-related methods.
  • D SUMMARY OF BANDITS: LinearEXP3 samples arms from exponentiated linear scores, while LinearFTPL applies Gumbel perturbations to linear reward estimates.These methods provide contextual adaptations of EXP3 and FTPL.
  • D.1 LINUCB: LinUCB estimates arm parameters from query features and selects the arm with the highest upper confidence bound.The parameter estimate and UCB are defined for each arm; α controls the exploration–exploitation trade-off.
  • D.1 LINUCB: Each round observes query features, scores every arm, applies the selected rewrite, observes reward, and updates the algorithm using the chosen arm, context, and reward.The general loop supports EXP3, FTPL, LinUCB, KL, FTRL, and Thompson.

D.2 LINUCB+KL BANDIT STRATEGY

The LINUCB+KL strategy initializes per-arm linear models, selects arms using context-dependent exploration scores, and updates the selected arm from observed rewards. Its exploration-adjusted reward is validated as a correctness proxy, while benchmark results show strong gains over the No-Rewrite baseline.

  • Initialization: Each arm initializes its matrix Aa to λId and vector ba to 0d, using shared parameters for arms, context dimension, regularization, exploration, noise, and KL bounds.The algorithm requires narms, d, λ, α, σnoise, and c.
  • Arm selection: The policy scores every arm from the context vector x, time step t, and selection counts, then selects the highest-scoring arm for exploration.Selection uses an exploration-adjusted score rather than a fixed global policy.
  • Reward validity: By ∼150 groups, the ROC–AUC confidence-interval lower bound exceeds 0.97, supporting the exploration-adjusted reward radj as a correctness proxy.The result reports mean ROC–AUC with 95% confidence intervals based on 10 resamples per n.
  • Model update: The update method revises Aa and ba for the selected arm using the received reward rt, while the KL bound dynamically adjusts exploration bonuses.This mechanism is intended to balance exploration and exploitation in a contextual setting.
  • Benchmark results: 88.8% is the LinUCB TruthfulQA MC1 score, improving 8.1 pp over the No–Rewrite baseline of 80.7%.Non–Contextual TS reaches 88.7%, +8.0 pp, while Contextual TS improves +4.5 pp.

D.3 FTRL … D.7 THOMPSON SAMPLING

The paper presents several linear contextual and non-contextual bandit policies, culminating in Thompson Sampling, which selects arms from sampled Bayesian parameter vectors and updates them using observed rewards. Across these methods, exploration is balanced with exploitation through regularization, randomization, exponential weighting, perturbations, or posterior sampling.

  • D.3 FTRL: FTRL maintains per-arm cumulative gradients and squared gradients, then selects the arm whose weight-context dot product is highest.Its updates incorporate cumulative reward and uncertainty through regularization scaled by α and β.
  • D.4 ε-GREEDY FOLLOW-THE-REGULARIZED-LEADER (FTRL) BANDIT POLICY: ε = 0.10 is used for exploration, with ε-greedy FTRL choosing a uniformly random arm with probability ε and otherwise maximizing the linear score.The selected arm alone is updated using ridge-regularized FTRL, while other arms remain unchanged.
  • D.5 LINEAR EXP3: Linear EXP3 computes an arm probability distribution from the context and updates the selected arm using an estimated reward.Exponential weighting and exploration bonuses balance exploration and exploitation in the linear contextual setting.
  • D.6 LINEAR FTPL: Linear FTPL selects the arm with the highest context-based perturbed score.Random Gumbel perturbations encourage exploration of suboptimal arms while retaining accumulated performance information.
  • D.7 THOMPSON SAMPLING: Thompson Sampling samples ˜θa ∼N(µa, Σa) for each arm and selects the arm maximizing the sampled score.Standard Bayesian linear regression updates µa and Σa using the observed reward r.
  • D.7 THOMPSON SAMPLING: Table 8 identifies, for each rewrite arm, the linguistic feature with the highest normalized coefficient at 100 % and the lowest at 0 %.The table also provides brief rationales for the features’ positive or negative effects on downstream reward.
  • D.7 THOMPSON SAMPLING: Answerability produces the strongest positive reward uplift at ≈+17 uniformly, while Ambiguity yields the largest negative uplift at ≈–15 to –18 across arms.Grounding also contributes +15–18, whereas Subjectivity imposes hits of ≈–10 to –14.
  • D.7 THOMPSON SAMPLING: Pairwise normalized coefficient differences indicate whether a linguistic feature favors the first or second rewrite arm in each matchup.“Win” denotes support for the first arm and “Loss” denotes support for the second, measured as a percentage of the feature’s full coefficient range.
Loading 2602.20332v1…