Correlation-Aware Contextual Bandits with Surrogate Rewards for LLM Routing

• URL: http://arxiv.org/abs/2607.09015v1
• Date: Fri, 10 Jul 2026 00:42:10 GMT
• Title: Correlation-Aware Contextual Bandits with Surrogate Rewards for LLM Routing
• Abstract summary: We study contextual bandit problems with correlated arms and access to surrogate reward signals produced by a machine learning model.<n>We propose algorithms that leverage such surrogate rewards through two complementary designs.
• Score: 15.152913607564876
• License: http://arxiv.org/licenses/nonexclusive-distrib/1.0/
• Abstract: We study contextual bandit problems with correlated arms and access to surrogate reward signals produced by a machine learning model, motivated by applications such as large language model (LLM) routing. Unlike classical contextual bandits that rely solely on bandit feedback and assume conditional independence across arms, our setting allows context-dependent inter-arm correlations and auxiliary reward information that may be noisy or misspecified. We propose algorithms that leverage such surrogate rewards through two complementary designs. A coupled reward-mixing approach pools true and surrogate rewards to accelerate learning when surrogate signals are reliable, while a decoupled prediction-mixing approach maintains separate estimators for bandit feedback and surrogate rewards and adaptively combines their predictions. This decoupling yields robustness to surrogate misspecification, recovering regret guarantees comparable to reward-only bandit methods in the worst case, while achieving improved regret when surrogate predictions are sufficiently informative. We provide theoretical regret analyses for both approaches and evaluate them on LLM routing benchmarks under varying accuracy versus cost trade-offs. The results demonstrate improved sample efficiency and consistently better accuracy-cost trade-offs compared to standard contextual bandit baselines and strong static routing methods.

Related papers

• Progressive Content Refinement with Decaying Reward Joint LinUCB [2.1290878226779877]
  We propose a novel contextual bandit algorithm that explicitly incorporates reward decay modeling.<n>We show that our method achieves significant performance gains over strong baselines.<n>Our ablation study confirms that the integration of reward decay modeling within the bandit framework is crucial for mitigating over-exploitation.
  arXiv  Detail & Related papers  (2026-08-07T03:17:22Z)
• TRACE: A Unified Rollout Budget Allocation Framework for Efficient Agentic Reinforcement Learning [52.67756371231985]
  Reinforcement learning with verifiable rewards (RLVR) is a promising approach for enhancing reasoning and agentic behavior in large language models.<n>We introduce Tree Rollout Allocation for Contrastive Exploration (TRACE), a unified rollout allocation framework.<n>Technically, TRACE allocates rollout budget to both prompt roots and intermediate prefixes that are most likely to yield mixed terminal rewards.
  arXiv  Detail & Related papers  (2026-06-09T17:16:03Z)
• Reasoning Is Not Free: Robust Adaptive Cost-Efficient Routing for LLM-as-a-Judge [4.511996087821266]
  Reasoning-capable large language models (LLMs) have recently been adopted as automated judges.<n>We show that explicit reasoning substantially improves judgment accuracy on tasks requiring structured verification.<n>We propose a Robust Adaptive Cost-Efficient Routing (RACER) which selects between reasoning and non-reasoning judges.
  arXiv  Detail & Related papers  (2026-05-11T16:30:20Z)
• Beyond Static Bias: Adaptive Multi-Fidelity Bandits with Improving Proxies [16.078266766067838]
  Multi-fidelity multiarmed bandits (MF-MAB) enable individual arms to be evaluated using diverse feedback sources that vary in both cost and accuracy.<n>We investigate adaptive MF-MAB with improving proxy sources, and focus on the canonical two-fidelity case in which the low-fidelity source becomes more informative with repeated use.<n>We propose the Threshold-Based Adaptive Continuation Companion (TACC), an optimistic algorithm that uses a bounded continuation rule to decide when low-fidelity sampling remains cost-effective and when to escalate.
  arXiv  Detail & Related papers  (2026-05-08T23:36:53Z)
• Best Arm Identification in Generalized Linear Bandits via Hybrid Feedback [23.37199172197756]
  We study fixed-confidence best arm identification in generalized linear bandits under a hybrid feedback model.<n>We introduce a likelihood-ratio-based confidence sequence that unifies heterogeneous generalized linear observations.<n>We propose a hybrid Track-and-Stop algorithm that adaptively allocates queries by tracking a minimax-optimal design over a joint action space of arms and pairs.
  arXiv  Detail & Related papers  (2026-05-07T06:40:42Z)
• Semi-Parametric Batched Global Multi-Armed Bandits with Covariates [0.48342038441006807]
  The multi-armed bandits (MAB) framework is a widely used approach for sequential decision-making.<n>We propose a novel semi-parametric framework for batched bandits with coparametrics and a shared parameter across arms.<n>Our algorithm, Batched single-Index Dynamic binning and Successive arm elimination (BIDS), employs a batched successive arm elimination strategy.
  arXiv  Detail & Related papers  (2025-03-01T17:23:55Z)
• Jailbreaking as a Reward Misspecification Problem [80.52431374743998]
  We propose a novel perspective that attributes this vulnerability to reward misspecification during the alignment process.<n>We introduce a metric ReGap to quantify the extent of reward misspecification and demonstrate its effectiveness.<n>We present ReMiss, a system for automated red teaming that generates adversarial prompts in a reward-misspecified space.
  arXiv  Detail & Related papers  (2024-06-20T15:12:27Z)
• Noisy Correspondence Learning with Self-Reinforcing Errors Mitigation [63.180725016463974]
  Cross-modal retrieval relies on well-matched large-scale datasets that are laborious in practice. We introduce a novel noisy correspondence learning framework, namely textbfSelf-textbfReinforcing textbfErrors textbfMitigation (SREM)
  arXiv  Detail & Related papers  (2023-12-27T09:03:43Z)
• Robust Lottery Tickets for Pre-trained Language Models [57.14316619360376]
  We propose a novel method based on learning binary weight masks to identify robust tickets hidden in the original language models. Experimental results show the significant improvement of the proposed method over previous work on adversarial robustness evaluation.
  arXiv  Detail & Related papers  (2022-11-06T02:59:27Z)
• Bias-Robust Bayesian Optimization via Dueling Bandit [57.82422045437126]
  We consider Bayesian optimization in settings where observations can be adversarially biased. We propose a novel approach for dueling bandits based on information-directed sampling (IDS) Thereby, we obtain the first efficient kernelized algorithm for dueling bandits that comes with cumulative regret guarantees.
  arXiv  Detail & Related papers  (2021-05-25T10:08:41Z)
• Output-Weighted Sampling for Multi-Armed Bandits with Extreme Payoffs [11.1546439770774]
  We present a new type of acquisition functions for online decision making in bandit problems with extreme payoffs. We formulate a novel type of upper confidence bound (UCB) acquisition function that guides exploration towards the bandits that are deemed most relevant.
  arXiv  Detail & Related papers  (2021-02-19T18:36:03Z)
• Robustness Guarantees for Mode Estimation with an Application to Bandits [131.21717367564963]
  We introduce a theory for multi-armed bandits where the values are the modes of the reward distributions instead of the mean. We show in simulations that our algorithms are robust to perturbation of the arms by adversarial noise sequences.
  arXiv  Detail & Related papers  (2020-03-05T21:29:27Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.