No One Size Fits All: QueryBandits for Hallucination Mitigation

1. Problem Statement

Large Language Models (LLMs), particularly advanced reasoning models, frequently suffer from hallucinations (generating factually incorrect outputs). While significant research exists on mitigating hallucinations in open-source models (via parameter editing, internal representation manipulation, or logit-based detection), there is a critical gap in addressing closed-source models (e.g., GPT-4o), which dominate institutional deployments.

• Limitation of Current Methods: Existing solutions often require access to model weights or internal states, making them inapplicable to black-box APIs. Furthermore, many approaches rely on "one-size-fits-all" static prompting strategies (e.g., always paraphrasing or always expanding), which fail to account for the specific linguistic nuances of individual queries.
• Core Hypothesis: The optimal strategy to mitigate hallucinations varies per query based on its linguistic "fingerprint." A static policy is suboptimal; an adaptive, online learning mechanism is required.

2. Methodology: QueryBandits

The authors propose QueryBandits, a model-agnostic, contextual bandit framework that operates entirely at the input layer. It adaptively selects a query-rewrite strategy for each incoming query to steer the LLM away from hallucinations without retraining the model.

A. Framework Formulation

QueryBandits frames the problem as a Contextual Multi-Armed Bandit:

• Context ($x_t$): A 17-dimensional binary vector of linguistically motivated features extracted from the query (e.g., presence of anaphora, subordination, domain specialization, ambiguity, negation).
• Actions ($A$): Five distinct rewrite strategies (Arms):
  1. PARAPHRASE: Introduce lexical diversity while preserving semantics.
  2. SIMPLIFY: Remove nested clauses and complex syntax.
  3. DISAMBIGUATE: Resolve vague references (pronouns, temporal expressions).
  4. EXPAND: Add salient entities and attributes to enrich context.
  5. CLARIFY TERMS: Define domain-specific jargon.
• Reward ($r_t$): A composite score calculated after the LLM generates a response. It is a weighted sum of three metrics:
  $$r_t = \alpha s_{llm} + \beta s_{fuzz} + \gamma s_{bleu}$$
  • $s_{llm}$: Binary correctness label from an LLM-as-a-judge (GPT-4o).
  • $s_{fuzz}$: Fuzzy string similarity (RapidFuzz).
  • $s_{bleu}$: BLEU-1 score for lexical overlap.
  • Optimal Weights: Through Pareto optimization on a human-labeled set, the authors determined $(\alpha, \beta, \gamma) = (0.6, 0.3, 0.1)$. This proxy achieves a ROC-AUC of 0.973 in distinguishing correct from incorrect answers.

B. Algorithms Evaluated

The paper compares several bandit algorithms against static prompts and a "No-Rewrite" baseline:

• Contextual Bandits: Thompson Sampling (TS), LinUCB, LinUCB+KL, Linear FTRL, Linear EXP3, Linear FTPL.
• Non-Contextual (Vanilla) Bandits: TS, EXP3, FTPL (treats all queries as identical).
• Baselines: Zero-shot static prompts (applying one rewrite strategy to all queries) and No-Rewrite.

3. Key Contributions

1. Reward Modeling for Factuality: Introduction of a calibrated, composite reward function ($r_t$) that effectively proxies hallucination. The authors rigorously validate that the LLM-judge term ($\alpha=0.6$) combined with fuzzy and lexical metrics provides a stable, high-discriminative signal for online learning.
2. Contextual Adaptation Superiority: Demonstration that Contextual Thompson Sampling significantly outperforms static policies and non-contextual bandits. The system learns to map specific linguistic features to the most effective rewrite strategy (e.g., using "EXPAND" for queries with high "Domain Specialization" but avoiding it for ambiguous queries).
3. No Single Optimal Policy: Empirical evidence proving that no single rewrite strategy is universally optimal. The effectiveness of an arm depends entirely on the query's semantic features. Static policies often incur higher cumulative regret than doing nothing (No-Rewrite), indicating that inflexible rewriting can worsen hallucinations.
4. Black-Box Applicability: The method operates purely as a forward-pass input filter, requiring no access to model weights, gradients, or internal representations. This makes it uniquely suitable for closed-source models where other mitigation techniques fail.

4. Experimental Results

The method was evaluated on 13 QA benchmarks (16 scenarios, ~1,050 queries per dataset) using GPT-4o in a black-box setting.

• Performance vs. Baseline:
  • The top-performing algorithm, Contextual Thompson Sampling, achieved an 87.5% win rate over the "No-Rewrite" baseline.
  • It improved macro-average accuracy from 68.1% (No-Rewrite) to 76.6% (+8.5 percentage points).
  • It outperformed zero-shot static policies (e.g., PARAPHRASE, EXPAND) by 42.6% and 60.3% respectively in win rate.
• Contextual vs. Non-Contextual:
  • Contextual bandits consistently outperformed non-contextual (vanilla) bandits. Ablating the 17-dimensional feature vector dropped the win rate from 87.5% to 81.7%.
  • Non-contextual bandits tended to converge to a single strategy per dataset, whereas contextual bandits diversified choices based on feature patterns.
• Regret Analysis:
  • Contextual learners incurred slightly higher instantaneous regret early in training (due to exploration) but achieved significantly lower cumulative regret and higher final accuracy compared to static policies.
  • Notably, some static strategies (like "Expand" on certain datasets) performed worse than the No-Rewrite baseline, confirming that blind rewriting is harmful.
• Closed-Source Impact:
  • On the TruthfulQA benchmark, QueryBandits lifted GPT-4o's MC1 score from 81.4% to 88.8% (+7.4 pp).
  • This contrasts with methods like DoLa or TruthX, which improve open-source models (e.g., Llama-2) but cannot be applied to closed models and show diminishing returns on stronger backbones.

5. Significance and Implications

• Bridging the Gap: QueryBandits addresses the critical lack of hallucination mitigation tools for the vast majority of enterprise-deployed closed-source LLMs.
• Paradigm Shift: It moves the field from "one-size-fits-all" prompting to adaptive, feature-conditioned policy learning. The paper establishes that linguistic features are strong predictors of which rewrite strategy will succeed.
• Efficiency: The approach is computationally lightweight (input-layer only) and requires no fine-tuning, making it a "plug-and-play" solution for reducing hallucinations in production systems.
• Insight on Memorization: The authors observed that on unperturbed benchmark queries, bandits often converged to "No-Rewrite," suggesting that standard benchmarks suffer from prompt memorization. The study highlights the need for lexically perturbed, semantically invariant queries to truly evaluate rewriting capabilities.

In conclusion, QueryBandits demonstrates that by treating query rewriting as an online contextual decision problem, it is possible to significantly reduce hallucinations in powerful, closed-source LLMs by dynamically selecting the optimal input transformation based on the query's linguistic structure.