cs.CLFeb 23, 2026
SaveNo One Size Fits All: QueryBandits for Hallucination Mitigation
Organizations: JPMorgan AI Research, New York, NY, USA
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 elicits a hallucinatory count ( integers) due to an ambiguous lower bound ( ). Conditioned on the query’s 17-dimensional feature vector, QueryBandits selects Expand and rewrites the query to with explicit bounds; the LLM then returns the correct cardinality ( ). 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 .
Figure 2: ( 2(a) ) Our chosen 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.
Appendix figures & tables22 assets
Supplementary material from the paper’s appendix.
Appendix
| Dataset | Scenario | Domain | License | Count | Citation |
| SQuADv2 | E, A | Wikipedia | CC BY-SA 4.0 | 86K | Rajpurkar et al. (2016) ; Rajpurkar et al. (2018) |
| TruthfulQA | M, A | General Knowledge | Apache-2.0 | 807 | Lin et al. (2022) |
| SciQ | M, A | Science | CC BY-NC 3.0 | 13K | Johannes Welbl (2017) |
| MMLU | M | Various | MIT | 15K | Hendrycks et al. (2021) |
| PIQA | M | Physical Commonsense | AFL-3.0 | 17K | Bisk et al. (2020) |
| BoolQ | M | Yes/No Questions | CC BY-SA 3.0 | 13K | Clark et al. (2019) ; Wang et al. (2019) |
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 vs. human labels.
Table 5: TruthfulQA MC1 comparison. 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 81%, while open–model interventions reported on Llama–2–7B remain far below the GPT–4o baseline (e.g., TruthX 54.22%, 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 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 A | Model B | % Agree | Cohen’s | MCC |
| gpt-5-2025-08-07 | gpt-5-mini-2025-08-07 | 0.960 | 0.916 | 0.916 |
| gpt-4.1-2025-04-14 | gpt-4o-2024-11-20 | 0.925 | 0.826 | 0.826 |
| gpt-4o-2024-11-20 | gpt-5-2025-08-07 | 0.909 | 0.802 | 0.810 |
| gpt-4.1-2025-04-14 | gpt-5-2025-08-07 | 0.906 | 0.794 | 0.807 |
| gpt-4o-2024-11-20 | gpt-5-mini-2025-08-07 | 0.903 | 0.790 | 0.801 |
| gpt-4.1-2025-04-14 | gpt-5-mini-2025-08-07 | 0.900 | 0.782 | 0.798 |
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 , and Matthews correlation coefficient (MCC).
| Stage | Median Tokens | Mean Tokens |
| Original query (input) | 16 | 19.3 |
| Feature-tagger output | 110 | 110.0 |
| Rewrite input | 26 | 29.3 |
| Rewrite output | 18 | 28.1 |
| Answer input | 64 | 91.3 |
| Answer output | 70 | 157.8 |
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 | Interpretation | Interpretation | ||
| Disambiguate | Subordination (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. |
| Simplify | Pragmatics (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. |
| Expand | Constraints (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. |
| Paraphrase | Answerability (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 Terms | Rarity (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. |
Table 8: Top Drivers ( ) and Reducers ( ) 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 ( 13 ) Under the contextual bandit, the strongest positive uplifts come from Answerability ( +17 uniformly) and Grounding (+15–18), while Ambiguity ( –15 to –18) and Subjectivity ( –10 to –14) impose the largest hits across all arms. Mid-range features like Presupposition and Constraints deliver modest boosts ( 5), and Excessive Details and Anaphora slightly hurt performance ( –5 to –7). ( 13 ) The non-contextual bandit amplifies these trends: Answerability and Grounding remain the top drivers ( +18–20), but Ambiguity becomes even more detrimental ( –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 ( –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 (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 ( -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.
| ID | Arm | System prompt template |
| Paraphrase | You 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. | |
| Simplify | You 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. | |
| Disambiguate | You 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. | |
| Expand | You 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. | |
| Clarify Terms | You 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. |
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).
| Feature | Description | Citation |
| Structural Features | ||
| Anaphora | Contains anaphoric references (e.g., it , this ) | Schuster (1988) ; Chen et al. (2018) |
| Subordination | Contains multiple subordinate clauses (multi-clause structure) | Jeong et al. (2024) ; Blevins et al. (2023) |
| Scenario-Based Features | ||
| Mismatch | Question–task mismatch (e.g., open-ended query against retrieval-style task) | Gao et al. (2024) ; Kamath et al. (2024) |
| Presupposition | Assumptions within the query are implicitly regarded as truthful | Karttunen (2016) ; Levinson (1983) |
Table 10: Binary linguistic feature vector 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 .
| Feature | Definition | Example | |
| Structural | Anaphora | Presence of pronouns or references requiring external context. | ”What about that one?” (Unclear reference) |
| Subordination | Measures the presence of multiple subordinate clauses | ”While I was walking home, I saw a cat that looked just like my friend’s.” | |
| Scenario-Based | Mismatch | Mismatch 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) |
| Presupposition | Unstated assumptions embedded in the query. | ”Who is the musician that developed neural networks?” (Assumes such a musician exists) | |
| Pragmatics | Captures context-dependent meanings beyond literal interpretation. | ”Can you pass the salt?” (A request, not a literal ability) | |
| Lexical | Rarity | Use of rare or niche terminology. | ”What are the ramifications of quantum decoherence?” (Uses low-frequency terms) |
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