cs.CLFeb 23, 2026

No One Size Fits All: QueryBandits for Hallucination Mitigation

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

Organizations: JPMorgan AI Research, New York, NY, USA

Read PDF
Ask AI
Cite

Copy citation

DOINot available for this paper

Abstract

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.

Figures & tables

Figure 1: QueryBandits selects a rewrite that fixes a counting error. The original query xtx_{t} elicits a hallucinatory count ( 88 integers) due to an ambiguous lower bound ( 66 ). Conditioned on the query’s 17-dimensional feature vector, QueryBandits selects Expand and rewrites the query to xt′x^{\prime}_{t} with explicit bounds; the LLM then returns the correct cardinality ( 99 ). Noticeably, the feature vector also shifts: subordination (more complex clauses) appears while specialization (domain-specific knowledge required) disappears-illustrating how rewriting alters the salient semantics of xtx_{t} .
Figure 2: ( 2(a) ) Our chosen (α,β,γ)(\alpha,\beta,\gamma) lies deep in the 1% optimal frontier. ( 2(b) ) Breakdown of per‐dataset arm performance: different datasets consistently favor different rewrite strategies
Figure 3: Cumulative Reward (averaged across all runs). Sorted by final performance, highlighting gains achieved by contextual bandits over non‐contextual learners and static rewrites.
Table 2: Left: Rewrite-policy Performance: final cumulative exploration-adjusted reward, mean cumulative regret, and win rate vs. no-rewrite. Right: Who Wins Where: best accuracy per dataset and gain over No-Rewrite baseline (pp). TS = Thompson Sampling; (C) = Contextual.
Figure 5
Appendix figures & tables22 assets

Supplementary material from the paper’s appendix.

Appendix

DatasetScenarioDomainLicenseCountCitation
SQuADv2E, AWikipediaCC BY-SA 4.086KRajpurkar et al. (2016) ; Rajpurkar et al. (2018)
TruthfulQAM, AGeneral KnowledgeApache-2.0807Lin et al. (2022)
SciQM, AScienceCC BY-NC 3.013KJohannes Welbl (2017)
MMLUMVariousMIT15KHendrycks et al. (2021)
PIQAMPhysical CommonsenseAFL-3.017KBisk et al. (2020)
BoolQMYes/No QuestionsCC BY-SA 3.013KClark et al. (2019) ; Wang et al. (2019)
Appendix
Table 3: Datasets. Overview of datasets, including domain, license, number of examples, associated scenarios, etc. These datasets span a diverse range of question types, domains, and reasoning skills, supporting robust evaluation. E = Extractive, M = Multiple Choice, A = Abstractive.
Figure 6: Summary of reward validity. Left: (a) numerical ROC–AUC and CIs across sample sizes. Right: (b) power curve; (c) class-conditional reward histogram of rtr_{t} vs. human labels.
Table 5: TruthfulQA MC1 comparison. Δ\Delta reports absolute percentage–point change vs our No–Rewrite baseline (80.7%). QueryBandits achieves the best score (LinUCB 88.8%, +8.1 pp) and strong Non–Contextual TS (88.7%, +8.0 pp); Contextual TS also improves (+4.5 pp). Closed GPT baselines cluster near ∼\sim 81%, while open–model interventions reported on Llama–2–7B remain far below the GPT–4o baseline (e.g., TruthX 54.22%, −26.5-26.5 pp). Results across families highlight that context–aware linear bandits (LinUCB) are most effective on MC1, with TS (Non–Contextual) close but lacking per–query adaptation.
Figure 7: For Non-Contextual bandits, almost every dataset is dominated by a single arm with the highest global reward (typically 40%-60% of the trials). The remaining 40-60% is split among the other four arms as noise, the non-contextual policy has no way to ”know” when within a dataset a different arm might do better. In contrast, Contextual bandits show a more even mix: the top arm is only ∼\sim 25-30%, with two or three other arms contributing sizable shares (15-25% each). The contextual policy reads the features and diversifies its choices within each dataset.
Figure 8: Impact of the No‐Rewrite Arm. Note that these experiments are conducted on the original query ”as-is” in the benchmark dataset, with no perturbations. Upon enabling the No Rewrite option, our contextual bandit rapidly converges to this arm, which then achieves the highest reward on several datasets. We attribute this behavior to the LLM’s tendency to memorize benchmark questions.
Figure 9: Comparison of Feature Variance between ( 9 ) our contextual bandits and ( 9 ) its non-contextual counterparts. Polysemy , Constraints and Entities show the most variation. Presupposition , Excessive Details , and Grounding have the least.
Figure 10: Comparison of Inter-Arm Context Distances (Symmetric KL) between ( 10 ) our contextual bandits and ( 10 ) its non-contextual counterparts. Arm pairs such as Expand and Paraphrase in the non-contextual bandit setting exhibit high KL distances at 1.01. One interpretation is that the context-clouds barely overlap from dataset to dataset (Figure 7 ).
Figure 11: Comparison of Raw feature‐level regression coefficients between ( 12 ) our contextual bandits and ( 12 ) its non-contextual counterparts. Each cell shows how enables a raw view into how specific linguistic feature changes the expected reward under each rewrite strategy.
Model AModel B% AgreeCohen’s κ\kappaMCC
gpt-5-2025-08-07gpt-5-mini-2025-08-070.9600.9160.916
gpt-4.1-2025-04-14gpt-4o-2024-11-200.9250.8260.826
gpt-4o-2024-11-20gpt-5-2025-08-070.9090.8020.810
gpt-4.1-2025-04-14gpt-5-2025-08-070.9060.7940.807
gpt-4o-2024-11-20gpt-5-mini-2025-08-070.9030.7900.801
gpt-4.1-2025-04-14gpt-5-mini-2025-08-070.9000.7820.798
Appendix
Table 6: Inter-model agreement on the LLM-as-judge labels over 1,000 validation queries. Values reported are fraction of exact label agreement, Cohen’s κ\kappa , and Matthews correlation coefficient (MCC).
StageMedian TokensMean Tokens
Original query (input)1619.3
Feature-tagger output110110.0
Rewrite input2629.3
Rewrite output1828.1
Answer input6491.3
Answer output70157.8
Appendix
Table 7: Token-level breakdown per query for QueryBandits. The total corresponds to a per-query cost of approximately $0.00035 at gpt-40-2024-11-20 pricing.
Figure 12: Comparison of Min-Max Normalized feature‐level regression coefficients between ( 12 ) our contextual bandits and ( 12 ) its non-contextual counterparts. Each cell shows how enables a relative view into how specific linguistic feature changes the expected reward under each rewrite strategy. Table 8 highlights contextual bandit trends.
Arm aafmax⁡+f^{+}_{\max}Interpretationfmax⁡−f^{-}_{\max}Interpretation
DisambiguateSubordination (100 %)Long or nested clauses benefit from targeted disambiguation, which isolates and clarifies the core semantic relation.Polysemy (0 %)Highly polysemous terms lead disambiguation to pick the wrong sense, degrading downstream reward.
SimplifyPragmatics (100 %)Pragmatic cues (e.g. discourse markers, politeness) guide safe simplification without loss of meaning.Superlative (0 %)Stripping superlative constructions removes essential comparative context, hurting reward.
ExpandConstraints (100 %)Queries already rich in constraints (time, location, numeric bounds) gain precision when expanded with further qualifiers.Ambiguity (0 %)Underspecified queries offer no detail to expand, so further addition of terms only introduces noise.
ParaphraseAnswerability (100 %)Paraphrasing queries that are already answerable refreshes wording while preserving solvability, boosting LLM performance.Presupposition (0 %)Altering queries with strong presuppositions can break implied assumptions, reducing effective reward.
Clarify TermsRarity (100 %)Defining rare or domain-specific terms anchors the LLM’s understanding of technical queries.Subordination (0 %)Clarifications in convoluted sentences can introduce further parsing difficulty, impeding reward.
Appendix
Table 8: Top Drivers ( fmax⁡+f^{+}_{\max} ) and Reducers ( fmax⁡−f^{-}_{\max} ) of Rewrite Strategies per Linguistic Features For each rewrite arm, we list the feature whose normalized coefficient was highest (100 %) and lowest (0 %), alongside a brief rationale for its positive or negative impact on downstream reward.
Figure 13: Reward Uplift by Contextual Feature and Strategy. Feature Uplift measures how much the presence of a binary feature changes the expected reward for a given rewrite arm, formally Δ(fi,a)=E[rt∣arm=a,fi=1]  −  E[rt∣arm=a,fi=0].\Delta(f_{i},a)=\mathbb{E}[r_{t}\mid\text{arm}=a,f_{i}=1]\;-\;\mathbb{E}[r_{t}\mid\text{arm}=a,f_{i}=0]. ( 13 ) Under the contextual bandit, the strongest positive uplifts come from Answerability ( ≈\approx +17 uniformly) and Grounding (+15–18), while Ambiguity ( ≈\approx –15 to –18) and Subjectivity ( ≈\approx –10 to –14) impose the largest hits across all arms. Mid-range features like Presupposition and Constraints deliver modest boosts ( ≈\approx 5), and Excessive Details and Anaphora slightly hurt performance ( ≈\approx –5 to –7). ( 13 ) The non-contextual bandit amplifies these trends: Answerability and Grounding remain the top drivers ( ≈\approx +18–20), but Ambiguity becomes even more detrimental ( ≈\approx –17 to –18), and Mismatch drops to nearly –16 under some arms. Notably, the non-contextual model shows a stronger negative effect for Excessive Details (up to –12) and Entities ( ≈\approx –6) than the linear one, suggesting it more sharply penalizes noisy contexts. Together, these heatmaps reveal which linguistic signals each rewrite strategy leverages (or struggles with), and how context vs. context-blind policies weigh them differently.
Figure 14: Per-dataset distributions of rtr_{t} (normalized density).
Figure 15: Anchored reward–difference distributions (left column) and dominance CDFs (right column) for the baseline policy ( Baseline ) and core static rewrite strategies ( Paraphrase , Simplify , Disambiguate ). Each row fixes an anchor policy and compares its per-query reward against all competitors.
Figure 16: Anchored reward–difference distributions (left) and dominance CDFs (right) for the remaining static rewrite strategies ( Clarify , Expand ) and simple non-contextual bandits ( EXP3 , FTPL ).
Figure 17: Anchored reward–difference distributions (left) and dominance CDFs (right) for advanced non-contextual bandits ( ϵ\epsilon -FTRL , TS (NC) ) and core contextual bandits ( LinUCB , LinUCB+KL ).
Figure 18: Anchored reward–difference distributions (left) and dominance CDFs (right) for the remaining contextual bandit policies ( LinEXP3 , LinFTPL , TS (Contextual) ).
Figure 19: Pairwise Normalized Coefficient Differences for Contextual Bandits. Each cell shows the min–max–normalized difference in regression weight for a given linguistic feature (rows) between two rewrite arms (columns), e.g. “Paraphrase vs Disambiguate,” “Simplify vs Expand,” etc. Cells labeled “Win” (blue) indicate the feature favors the first arm in the matchup, while “Loss” (red) indicates it favors the second. Values are expressed as a percentage of the feature’s full coefficient range.
IDArmSystem prompt template
a0a_{0}ParaphraseYou are a rewriting module. You will be given a user query: {original_query} . Rephrase it to improve clarity and introduce lexical diversity while strictly preserving semantic meaning, entities (including casing/accents), numbers, units, and constraints. Do not add or remove information. Output only the rewritten query.
a1a_{1}SimplifyYou are a rewriting module. You will be given a user query: {original_query} . Simplify it by removing nested clauses and complex syntax. Use short, concrete phrasing (S–V–O order), keep all entities, numbers, units, and constraints, and avoid changing intent. Do not invent details. Output only the simplified query.
a2a_{2}DisambiguateYou are a rewriting module. You will be given a user query: {original_query} . Resolve vague references by replacing ambiguous pronouns (e.g., it/they/this) and temporal expressions with explicit, context-grounded referents and normalized dates. If a referent cannot be determined from the query alone, insert a bracketed placeholder (e.g., [ENTITY], [DATE]) rather than guessing. Preserve the original intent. Output only the disambiguated query.
a3a_{3}ExpandYou are a rewriting module. You will be given a user query: {original_query} . Expand it by making implicit context explicit and adding salient, non-speculative attributes (e.g., scope, timeframe, location, units) that are entailed by the query. If crucial specifics are missing, insert neutral bracketed placeholders (e.g., [TIMEFRAME], [LOCATION]) instead of fabricating facts. Preserve the original intent and constraints. Output only the expanded query.
a4a_{4}Clarify TermsYou are a rewriting module. You will be given a user query: {original_query} . Identify domain-specific jargon or terms of art and add concise parenthetical glosses (e.g., “term (brief definition)”) where the meaning is standard and unambiguous. If uncertain, use a bracketed clarification placeholder (e.g., [DEFINE: TERM]) rather than guessing. Do not alter intent, entities, or constraints. Output only the clarified query.
Appendix
Table 9: System prompt templates for rewrite arms. Replace {original_query} with the input at runtime. Each template must output only the rewritten query (no explanations).
FeatureDescriptionCitation
Structural Features
AnaphoraContains anaphoric references (e.g., it , this )Schuster (1988) ; Chen et al. (2018)
SubordinationContains multiple subordinate clauses (multi-clause structure)Jeong et al. (2024) ; Blevins et al. (2023)
Scenario-Based Features
MismatchQuestion–task mismatch (e.g., open-ended query against retrieval-style task)Gao et al. (2024) ; Kamath et al. (2024)
PresuppositionAssumptions within the query are implicitly regarded as truthfulKarttunen (2016) ; Levinson (1983)
Appendix
Table 10: Binary linguistic feature vector f∈{0,1}17\mathbf{f}\in\{0,1\}^{17} identified as challenging from a linguistics and LLM perspective. Features are grouped by type and grounded in prior work. For more specific examples, see Table 11 .
 FeatureDefinitionExample
StructuralAnaphoraPresence of pronouns or references requiring external context.”What about that one?” (Unclear reference)
SubordinationMeasures the presence of multiple subordinate clauses”While I was walking home, I saw a cat that looked just like my friend’s.”
Scenario-BasedMismatchMismatch between the query’s intended output and its actual structure.”Find me this paragraph in this document” (When document isn’t given, this query cannot be answered)
PresuppositionUnstated assumptions embedded in the query.”Who is the musician that developed neural networks?” (Assumes such a musician exists)
PragmaticsCaptures context-dependent meanings beyond literal interpretation.”Can you pass the salt?” (A request, not a literal ability)
LexicalRarityUse of rare or niche terminology.”What are the ramifications of quantum decoherence?” (Uses low-frequency terms)
Appendix
Table 11: Detailed Summary and Examples of Feature Categories, Definitions, and Examples (See Table 10 ). These definitions and examples become the prompts to create the binary context vector for our bandits.
More in this topic

LLM Hallucination Mitigation

Large language model hallucination mitigation develops ways to detect and reduce unsupported, fabricated, or inconsistent outputs.

See all 178 papers

Explore similar work

CardsList
May 26, 2026

Chain-based Adaptive Reconfiguration Over Lattices for Hallucination Reduction

We introduce CAROL (Chain-based Adaptive Reconfiguration Over Lattices), a probabilistic framework for test-time hallucination reduction in large language models. Rather than relying on token-level uncertainty, CAROL defines a semantic uncertainty measure based on the consistency between generated responses and a trusted context, inducing a string-submodular objective over a lattice of textual sequences. This formulation enables hallucination mitigation to be cast as a Markov chain accept-reject process with provable convergence and near-optimality guarantees, allowing the model to iteratively refine outputs toward semantic consistency. By operating at the level of meaning, CAROL unifies hallucination detection and mitigation within a single framework. Empirical results on question answering and multi-agent reasoning benchmarks show that CAROL significantly reduces hallucinations and improves reliability and interpretability compared to likelihood-based and retrieval-augmented baselines, while maintaining competitive computational efficiency.

Joan Vendrell Gallart, Solmaz Kia, Russell Bent +1

Department of Mechanical and Aerospace University of California Irvine Irvine, CA 92617-4322, USA · T-5 Los Alamos National Laboratory Los Alamos, NM 88220, USA · CAI-4 Los Alamos National Laboratory Los Alamos, NM 88220, USA

LLM Hallucination MitigationHallucination in Language ModelsLLM ReliabilityLLM Hallucination Detection
AskPDF
Jun 14, 2026

BALTO: Balanced Token-Level Policy Optimization for Hallucination Mitigation

Hallucinations remain a major obstacle to deploying large language models (LLMs) in knowledge-intensive settings, where generated responses must be faithfully grounded in provided evidence. Reinforcement learning (RL) is a promising direction for hallucination mitigation, but response-level faithfulness rewards suffer from a granularity mismatch: localized hallucinations can cause supported content to receive spurious penalties. Although recent work introduces fine-grained feedback such as claim-level verification and token-level rewards, unbalanced credit assignment can still induce length, verbosity, or optimization-noise biases. We propose BALTO, a Balanced Token-level Policy Optimization framework for hallucination mitigation. BALTO extracts checkable factual claims, verifies them against the reference context, and projects claim-level judgments to token-level labels. A balanced token-level credit assignment mechanism is introduced into the framework. This design redistributes probability mass from unsupported content toward faithful content, rather than suppressing the entire response. We systematically analyze the limitations of response-level rewards from a theoretical standpoint, and prove BALTO's advantages in training stability and optimization efficiency for hallucination mitigation. Experiments on ConFiQA, RAGTruth, and FinLLM-Eval show that BALTO achieves the highest faithfulness across all six model--benchmark settings and consistently outperforms existing post-training baselines in Q-Score, demonstrating a stronger faithfulness--informativeness trade-off.

Ning Li, Zixuan Guo, Yan Xu +7

1Shanghai Jiao Tong University · 2Tencent · 3The

RL for Language ModelsLLM Hallucination MitigationToken-Level Credit Assignment
AskPDF
Jul 12, 2026

To Answer or to Abstain: Mitigating Search-Agent Hallucinations via Abstention-Aware Reinforcement Learning

Recent advances in equipping Large Language Models (LLMs) with search tools and outcome-reward reinforcement learning (RL) have achieved new state-of-the-art results on open-domain QA tasks. However, we argue that current training paradigms harbor a critical vulnerability: they predominantly reward correct answers but fail to penalize fabricated ones when retrieval fails, thereby implicitly exacerbating hallucinations. To address this, we propose Abstention-Aware Reinforcement Learning (AWA-RL), which dynamically shapes the abstention reward utilizing the model's query-specific prior capabilities and continuous on-policy training observations. We also introduce a novel metric, RA-F1, to measure the capability-reliability trade-off. Compared to non-abstaining baselines, AWA-RL boosts absolute precision by up to 10.3% and overall RA-F1 by 2.9%, with only marginal sacrifice in raw accuracy. These results confirm that AWA-RL successfully yields highly capable and reliable search agents. The code, data, and model weights are publicly available at https://github.com/zfj1998/AWA-RL.

Fengji Zhang, Tianyu Fan, Yuxiang Zheng +4

City University of Hong Kong · Alibaba Group

Question AnsweringRL for LLM AgentsLLM Hallucination MitigationAI Agent ReliabilityAgentic RAGSelective QA+1
AskPDF