# Correlation-Aware Contextual Bandits with Surrogate Rewards for LLM Routing

arXiv: https://arxiv.org/abs/2607.09015 · HTML: https://mlsanity.com/paper/2607.09015/

## Key facts

- Cited by 3 papers in the map (data through 14 Oct 2026).
- Builds on 27 of its 55 references that are in the map.
- First posted on arXiv 10 Jul 2026.

## Builds on (most cited first, 27 of 27)

- [DeBERTa: Decoding-enhanced BERT with Disentangled Attention](https://mlsanity.com/paper/2006.03654/) (2020)
- [FrugalGPT: How to Use Large Language Models While Reducing Cost and Improving Performance](https://mlsanity.com/paper/2305.05176/) (2023)
- [RouteLLM: Learning to Route LLMs with Preference Data](https://mlsanity.com/paper/2406.18665/) (2024)
- [LLM-Blender: Ensembling Large Language Models with Pairwise Ranking and Generative Fusion](https://mlsanity.com/paper/2306.02561/) (2023)
- [RouterBench: A Benchmark for Multi-LLM Routing System](https://mlsanity.com/paper/2403.12031/) (2024)
- [The Statistical Complexity of Interactive Decision Making](https://mlsanity.com/paper/2112.13487/) (2021)
- [Beyond UCB: Optimal and Efficient Contextual Bandits with Regression Oracles](https://mlsanity.com/paper/2002.04926/) (2020)
- [Finite-Time Analysis of Kernelised Contextual Bandits](https://mlsanity.com/paper/1309.6869/) (2013)
- [GraphRouter: A Graph-based Router for LLM Selections](https://mlsanity.com/paper/2410.03834/) (2024)
- [Corralling a Band of Bandit Algorithms](https://mlsanity.com/paper/1612.06246/) (2016)
- [RouterDC: Query-Based Router by Dual Contrastive Learning for Assembling Large Language Models](https://mlsanity.com/paper/2409.19886/) (2024)
- [Online Learning with Feedback Graphs: Beyond Bandits](https://mlsanity.com/paper/1502.07617/) (2015)
- [Nonstochastic Multi-Armed Bandits with Graph-Structured Feedback](https://mlsanity.com/paper/1409.8428/) (2014)
- [Relative Loss Bounds for On-line Density Estimation with the Exponential Family of Distributions](https://mlsanity.com/paper/1301.6677/) (2013)
- [Efficient and Robust Algorithms for Adversarial Linear Contextual Bandits](https://mlsanity.com/paper/2002.00287/) (2020)
- [MetaLLM: A High-performant and Cost-efficient Dynamic Framework for Wrapping LLMs](https://mlsanity.com/paper/2407.10834/) (2024)
- [Cost-Effective Online Multi-LLM Selection with Versatile Reward Models](https://mlsanity.com/paper/2405.16587/) (2024)
- [CARROT: A Cost Aware Rate Optimal Router](https://mlsanity.com/paper/2502.03261/) (2025)
- [LLM Bandit: Cost-Efficient LLM Generation via Preference-Conditioned Dynamic Routing](https://mlsanity.com/paper/2502.02743/) (2025)
- [Learning to Route LLMs from Bandit Feedback: One Policy, Many Trade-offs](https://mlsanity.com/paper/2510.07429/) (2025)
- [Understanding Bandits with Graph Feedback](https://mlsanity.com/paper/2105.14260/) (2021)
- [Multi-Armed Bandits With Machine Learning-Generated Surrogate Rewards](https://mlsanity.com/paper/2506.16658/) (2025)
- [Practical Contextual Bandits with Feedback Graphs](https://mlsanity.com/paper/2302.08631/) (2023)
- [Exploiting Correlated Auxiliary Feedback in Parameterized Bandits](https://mlsanity.com/paper/2311.02715/) (2023)
- [Contextual Linear Bandits with Delay as Payoff](https://mlsanity.com/paper/2502.12528/) (2025)
- [LLM Routing with Dueling Feedback](https://mlsanity.com/paper/2510.00841/) (2025)
- [Efficient Contextual Bandits with Uninformed Feedback Graphs](https://mlsanity.com/paper/2402.08127/) (2024)

## Built on by (most cited first, 3 of 3)

- [The Surrogate Is Not the Reward: Post-Surrogate Primary-Outcome Acquisition in Contextual Bandits](https://mlsanity.com/paper/2610.06610/) (2026)
- [CopRoute: Dependence-Aware Conformal Routing for Language Models](https://mlsanity.com/paper/2610.13276/) (2026)
- [Gaussian Process Bandits for AI System Optimization with Selectively Verified Feedback](https://mlsanity.com/paper/2610.17760/) (2026)

## Abstract (arXiv)

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.