Correlation-Aware Contextual Bandits with Surrogate Rewards for LLM Routing

arXiv:2607.09015 · cs.LG, cs.AI · Submitted 2026-07-10 · Read on arXiv

Listen

Radio episode about this paper

Transcript

Introduction to the show: ident: AI Radio. Generated commentary on the latest Artificial Intelligence papers.

Tom: I'm Tom.

Jane: And I'm Jane, here with Lu, Meng and our in-house large language model, Lalam.

Tom: Today's paper: "Correlation-Aware Contextual Bandits with Surrogate Rewards for LLM Routing".

Jane: The gist: We study contextual bandit problems with correlated arms and access to surrogate reward signals produced by a machine learning model,

Tom: First, who's behind it and why it matters.

Title and authors: Tom: So, to recap, this paper is about using context-aware feedback graphs and ML predictions as surrogate rewards in contextual bandit problems with correlated arms. The two main approaches they propose are a coupled reward-mixing method called CABS-C and a decoupled prediction-mixing approach called CABS-D.

Jane: Let's unpack that a bit for the listener. Think of it like this: one method, coupled, pools the real rewards you get with the predicted surrogate rewards to train one big model. The other one, decoupled, keeps two separate experts—one just for standard bandit feedback and another for handling those correlation-aware surrogate signals—and then they combine them smartly using an expert aggregation master.

Lu: That decoupling is interesting because it suggests you can track which learning mechanism is better without having to fully trust the surrogate model across the board. They want to maintain robustness even if the ML prediction for the surrogate reward is wrong.

Meng: If we're deploying this, I'm thinking about complexity. CABS-C seems like it might be training one complex model, while CABS-D maintains two stateful experts, which sounds more like managing two distinct pipelines.

Lalam: The paper shows that this decoupling allows them to separate the noise from the true signal better when the surrogate feedback is unreliable. They are aiming for a guarantee comparable to just a standard reward-only bandit method in those tricky situations.

Jane: So, they aren't just trying to get high accuracy; they’re trying to get reliable learning even when you don't trust all your auxiliary information. This is important because in real applications like LLM routing, those proxy signals aren't always perfect.

Tom: It really boils down to finding the sweet spot between leveraging that predictive power and sticking to what the true rewards actually tell us. Now, we need to see how they quantify these improvements in the next segment.

The paper's summary: Tom: The authors summarize their work by focusing on how they tackle the central research question: Can we design contextual bandit algorithms that systematically exploit both context-dependent intermodel correlations and auxiliary ML-predicted surrogate rewards to accelerate learning, while remaining robust to surrogate misspecification?

Jane: So, essentially, they are trying to figure out if you can use the extra information—the AO and the RO signals—to get faster learning without letting a bad prediction ruin your entire strategy.

Lu: They spend a lot of time on defining the feedback graph structure that identifies which arms need attention based on what has already been chosen, which is their way of handling those context-dependent correlations.

Meng: I noticed they mention observing the true reward only for the selected arm, and then getting these surrogates for others in that graph set. That's a specific setup, not just a general idea.

Lalam: It’s about that partial multi-arm feedback structure they set up there. They are using the ML prediction to give you information about neighbors in this graph structure, even though you only got the real reward for your chosen arm.

Jane: So, it's not just getting rewards; it’s getting a structured way of seeing how those models relate to each other through context and prediction. It’s like they are building a map of relationships to guide their exploration.

Tom: And they explicitly state that their setup is different from existing works, like SquareCB-G or SquareCBUG, because they observe the true reward only for the chosen arm and receive surrogates for others in this graph known before action selection.

The paper's improvements: Tom: They present two main algorithmic improvements: CABS-C, the coupled approach, and CABS-D, the decoupled approach. CABS-C pools true and surrogate rewards to train a single contextual model, which they claim accelerates learning when the surrogate signals are reliable.

Jane: The improved regret bounds show that for CABS-C, they get a scaling of RC equal to Oe(p dKT log(T /d)/(m + one) plus epsilon n p TK2/(m + one)) <ref:2607.09015#pg1>. That’s getting tighter on the exploration part when you have reliable surrogates.

Lu: For CABS-D, the decoupling is what gives them the best guarantee overall. They show a result RD equal to O(min min RS, RC). This means they get the better of both worlds: they can match standard contextual bandit rates in the worst case while still getting those correlation-driven gains if the surrogate feedback is trustworthy.

Meng: From a deployment standpoint, I’m looking at that CABS-D result RD equal to O(min min RS, RC). That sounds like it gives you a reliable performance ceiling across different scenarios. It’s not all about the theoretical best case for one specific setup.

Lalam: And the experimental results back this up by showing that CABS-D consistently improves utility and reduces regret compared to CABS-C across all regimes, suggesting that separating the learning process is generally better for reliable use of those surrogate signals.

Jane: The experimental evidence suggests that CABS-D is the strongest implementable method they have, achieving the best non-oracle utility across different cost regimes for both linear and neural implementations.

Conclusion: Tom: So to wrap up, the paper introduces Correlation-Aware Contextual Bandits with Surrogate Rewards for LLM Routing, proposing CABS-C for coupling and CABS-D for decoupling. The main implication is that you can use context correlations and ML predictions to speed up learning, but you have to choose your strategy based on whether you trust those surrogate signals.

Jane: The big picture here is that in real applications like LLM routing, where the true rewards are expensive or stochastic, having a framework that lets you adapt—either by combining everything or by keeping experts separate—is really powerful for managing risk.

Lu: What they showed is that the correlation-aware graph feedback helps improve exploration by reducing the effective action-space dependence from K to K/(m + one) when those surrogates are reliable, which significantly shrinks both the primary exploration term and the noise propagation term <ref:2607.09015#pg1>.

Meng: I just see it as a way to get more information per round. The coupled approach does that by pooling true and surrogate rewards, even if it means a slightly higher complexity of training one large model.

Lalam: For me, the most important part is that CABS-D allows you to track the better feedback mechanism without needing to know if the surrogates are reliable; it lets you fall back on standard bandit learning when things get noisy.

Tom: We’ve covered how they handle correlations and surrogate rewards with CABS-C and CABS-D, showing that decoupling is often more robust for practical use cases where the ML predictions might not be perfect.

Jane: That’s the essence of it. This framework gives us a concrete way to think about how we can leverage predicted surrogate rewards in complex decision-making environments like routing models.

Tom: It’s solid work on context, correlation, and reward signals with LLM routing problems. We'll catch another paper soon.

Ajay Narayanan Sridhar, Ronak Singh, Mehrdad Mahdavi, Vijaykrishnan Narayanan

The Pennsylvania State University

cs.LG, cs.AI

Submitted: 2026-07-10

Updated: 2026-10-13

License: http://arxiv.org/licenses/nonexclusive-distrib/1.0/

Importance score: 86/100

The gist: The gist: 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

Key concepts

Contextual Bandit with Correlated Arms
This is an online learning problem where the learner must choose one of several options (arms) based on a given context (like a text embedding). The key difficulty is that the choices are not independent; some LLMs might be correlated in their performance, meaning choosing one influences the expected outcome for others.
Surrogate Reward Signals
These are predicted rewards generated by an auxiliary machine learning model. They serve as cheap proxies for the true, expensive rewards. The paper explores how to use these noisy or potentially inaccurate predictions to guide decision-making in a real-world scenario like LLM routing.
Coupled Reward-Mixing Approach (CABS-C)
This algorithm combines the true realized reward with the surrogate reward into a single model. It aims to leverage the reliability of the surrogate signal when it is good, effectively de-biasing it using both types of feedback to train one unified, more robust prediction model.
Decoupled Prediction-Mixing Approach (CABS-D)
This strategy uses two separate experts: one for standard bandit feedback and another for correlation-aware surrogate feedback. It adaptively combines their predictions, offering robustness against errors in the surrogate signal while still gaining benefits when the surrogate information is accurate.

Terminology

Summary

The gist: 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.

Problem Setting

The setting considers an online LLM routing problem operating over a finite horizon of T discrete rounds, where the learner must select an LLM from a fixed pool of K distinct LLMs given a context embedding vector xt in R d The selected model, it, produces an outcome ot and incurs a cost-aware reward rt,it(xt); when convenient, we write rt,it and define the corresponding loss lt,it = 1 − rt,it To exploit query-dependent correlations among LLMs and auxiliary reward predictions from offline data, we formulate LLM routing as a correlation-aware contextual bandit with surrogate rewards At each round t, in addition to xt, the learner observes inter-arm side information from an affinity oracle (AO), represented either as a correlation matrix Rt in R K×K or as a time-varying, context-dependent directed feedback graph Gt over [K] with adjacency matrix At in K×K Upon selecting arm it, the learner observes the true realized reward only for the chosen arm it. However, for other arms j!= it specified by the feedback graph, the learner receives surrogate rewards st,j from a reward oracle (RO)

Proposed Algorithms

The paper proposes two complementary designs to leverage these signals

  1. Coupled reward-mixing approach (CABS-C): This algorithm pools true and surrogate rewards to train a single contextual model CABS-C de-biases surrogate rewards from correlated arms and pools them with the selected arm’s true reward to train a single reward model

  2. Decoupled prediction-mixing approach (CABS-D): This strategy treats (i) a reward-only contextual bandit; and (ii) a correlation-aware (surrogates induced graph feedback) bandit (CABS-C) as two experts and combines them using an adaptive expert-aggregation master, in the spirit of adaptive Hedge methods CABS-D maintains two stateful experts: a standard bandit-feedback expert and the surrogate-aware graph-feedback expert CABS-C

Theoretical Guarantees

The paper establishes regret bounds that explicitly separate correlation-driven gains (reduced effective exploration) from surrogate noise/misspecification

  1. For CABS-C, the regret scales as RC = Oe(p dKT log(T /d))/(m + 1) + εn p TK2/(m + 1)

  2. For CABS-D, the resulting meta-bandit promises a best-of-both-worlds guarantee, RD = O(min min RS, RC), allowing us to match the standard contextual bandit rate (plus mild overhead) in the worst case while inheriting correlation-driven gains when surrogate feedback is reliable

  3. For a linear oracle and m-feedback graph, Corollary 4.8 shows that CABS-D guarantees E[RegT] ≤ Oe(min (sdK m + 1 T log T d + εn rK2 m + 1 T!), sdKT log T d)

Experimental Results

Experiments on LLM routing benchmarks demonstrate improved sample efficiency and better accuracy–cost trade-offs compared to online bandit baselines and static routing policies

  1. CABS-D consistently improves utility and reduces regret relative to CABS-C across all regimes, indicating that decoupling reward learning from surrogate learning yields a more reliable use of surrogate signals

  2. CABS-D is the strongest implementable method, attaining the best non-oracle utility across cost regimes for both linear and neural instantiations

  3. Surrogate rewards accelerate online learning as measured by reduced cumulative regret over time, with the improvement most pronounced in the mid-cost regime

Discussion

The comparison between coupled (CABS-C) and decoupled (CABS-D) approaches shows that CABS-D can exploit informative surrogate predictions early while falling back to reward-based learning as data accumulates CABS-D retains standard bandit robustness while exploiting informative surrogate feedback The computational complexity for CABS-C is O(K cubed + RS), while CABS-D has O(K cubed + K ln T + RS) The incremental overhead from surrogate prediction beyond the encoder is negligible

Conclusion

We introduce a framework for contextual bandits with correlated arms and noisy surrogate rewards, motivated by LLM routing settings where true rewards are expensive or stochastic but cheaper proxy signals are available We propose two complementary algorithms for leveraging these signals: a coupled reward-mixing method that pools true and surrogate feedback, and a decoupled prediction-mixing method that maintains separate learners for bandit-only feedback and the surrogate-aware feedback before adaptively combining their policies Our regret analysis clarifies when tight coupling is beneficial and when decoupling is needed for robustness Experiments on LLM routing benchmarks demonstrate that correlation-aware surrogate integration significantly improves sample efficiency and achieves superior accuracy–cost trade-offs compared to existing baselines, validating both the theoretical and practical benefits of the proposed framework The results demonstrate that leveraging predicted surrogate rewards offers a powerful framework for LLM routing, yielding substantial regret reductions

References

[1] Alekh Agarwal, Haipeng Luo, Behnam Neyshabur, and Robert E Schapire. Corralling a band of bandit algorithms. In Conference on Learning Theory, pages 12–38. PMLR, 2017

[4] Noga Alon, Nicolò Cesa-Bianchi, Claudio Gentile, Shie Mannor, Yishay Mansour, and Ohad Shamir. Nonstochastic multi-armed bandits with graph-structured feedback, 2014

[5] Noga Alon, Nicolò Cesa-Bianchi, Claudio Gentile, Shie Mannor, Yishay Mansour, and Ohad Shamir. Nonstochastic multi-armed bandits with graph-structured feedback, 2014

[9] Yair Caro. New results on the independence number. Technical report, Technical Report, Tel-Aviv University, 1979

[15] Wei Chu, Lihong Li, Lev Reyzin, and Robert Schapire. A contextual-bandit approach to personalized news article recommendation. In Proceedings of the 19th international conference on World wide web, pages 661–670, 2010

[25] Qitian Jason Hu, Jacob Bieker, Xiuyu Li, Nan Jiang, Benjamin Keigwin, Gaurav Ranganath, Kurt Keutzer, and Shriyash Kaustubh Upadhyay. Routerbench: A benchmark for multi-llm routing system, 2024

[31] Tor Lattimore and Csaba Szepesvári. Bandit algorithms. Cambridge University Press, 2020

[54] Mengxiao Zhang, Yuheng Zhang, Olga Vrousgou, Haipeng Luo, and Paul Mineiro. Practical contextual bandits with feedback graphs. Advances in Neural Information Processing Systems, 36:30592–30617

[17] Tim Erven, Wouter M Koolen, Steven Rooij, and Peter Grünwald. Adaptive hedge. Advances in Neural Information Processing Systems, 24, 2011

[37] Isaac Ong, Amjad Almahairi, Vincent Wu, Wei-Lin Chiang, Tianhao Wu, Joseph E Gonzalez, M Waleed Kadous, and Ion Stoica. Routellm: Learning to route llms from preference data. In The Thirteenth International Conference on Learning Representations, 2025

[41] Seamus Somerstep, Felipe Maia Polo, Allysson Flavio Melo de Oliveira, Prattyush Mangal, Mírian Silva, Onkar Bhardwaj, Mikhail Yurochkin, and Subha Maity. Carrot: A cost aware rate optimal router. arXiv preprint arXiv:2502.03261, 2025

[56] Dongruo Zhou, Lihong Li, and Quanquan Gu. Neural contextual bandits with ucb-based exploration. In International conference on machine learning, pages 11492–11502. PMLR, 2020

[33] Yang Li. Llm bandit: Cost-efficient llm generation via preference-conditioned dynamic routing, 2025

[48] Xinyuan Wang, Yanchi Liu, Wei Cheng, Xujiang Zhao, Zhengzhang Chen, Wenchao Yu, Yanjie Fu, and Haifeng Chen.

Improvements for AI systems

  1. textbfRewarding LLM Routing via Correlation-Aware Surrogate Rewards (CABS-D): The improved system can select models that maximize a cost-aware utility function defined as rt,it = yt,it − ρ ct,it, effectively balancing accuracy and operational constraints based on the context embedding xt. This allows the router to dynamically adapt its choice between models like cheaper or faster ones when operating in high cost sensitivity regimes (ρ ≥ 800), leading to a better accuracy–cost trade-off as demonstrated by achieving higher utility across cost regimes in Table 2 of Section 5.3.

  2. textbfRobust Routing Against Surrogate Misspecification via Decoupled Prediction Mixing (CABS-D): The system maintains two separate experts—a standard bandit-feedback expert and a correlation-aware graph-feedback expert—and combines their policies adaptively using an adaptive second-order learning rate. This strategy allows the system to track the better feedback mechanism without knowing whether surrogates are reliable, ensuring that when surrogate rewards are uninformative or too noisy, it fall back on the standard contextual bandit expert and retain the corresponding regret guarantee Oe(qKT, RegSq(T)).

  3. textbfIncreased Sample Efficiency via Graph-Feedback Exploration (CABS-C): The system utilizes context-dependent inter-arm correlations to form a feedback graph that identifies a small set of additional arms whose feedback is most informative given the chosen arm. This allows for an improved exploration factor from K to K/(m + 1), which, when surrogate rewards are reliable, results in an improvement where the effective action-space dependence drops from K to K/(m + 1), significantly shrinking both the primary exploration term and the noise propagation term.

  4. textbfAdaptive Exploration Strategy via Coupled Reward Mixing (CABS-C): The coupled approach pools de-biased true rewards with surrogate rewards to fit a single contextual model, which effectively increases the number of informative observations per round, improving the exploration factor from K to K/(m + 1). This coupling is beneficial when surrogates are reliable but necessitates an explicit trade-off, as the regret scales as RC = Oe(p dKT log(T /d)/(m + 1) + εn p TK2/(m + 1)).

Sources

Related papers