Correlation-Aware Contextual Bandits with Surrogate Rewards for LLM Routing

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.