1 Introduction↩︎

The rapid proliferation of Large Language Models (LLMs) has made it increasingly challenging for end-users to keep track of advancements and optimally select models for their specific needs. Presently, numerous proprietary LLM providers exist [1]–[3], and thousands more are openly available through repositories such as Hugging Face [4]. This abundance creates a practical systems problem: for a given query, which model should be called? The answer is rarely the most accurate model, since inference incurs nontrivial cost (e.g., pricing, latency, compute). In many applications, the objective is an accuracy–cost trade-off: a slightly less accurate model may be preferable if it is significantly cheaper or faster, and the optimal choice varies with context.

Existing unified interfaces address this challenge through model cascading [5], [6], ensembling [7], and routing [8]–[12]. We focus on routing: selecting one model per query to maximize user-defined utility. A common reduction models LLM routing as a contextual multi-armed bandit, where each LLM is treated as a conditionally independent arm given the incoming query context, and only the selected arm’s reward is observed, requiring costly exploration. In the \(d\)-dimensional linear setting, where the context is represented by a \(d\)-dimensional embedding and each arm’s expected reward is assumed to be a linear function of this embedding, standard methods achieve regret \(R_S=~O(\sqrt{dKT\log(T/d)})\) [13] given \(K\) arms (LLMs) for \(T\) stream of queries, where regret is the cumulative loss relative to the best context-dependent action. However, this conditional-independence abstraction is often too coarse: models exhibit query-dependent correlations due to shared pretraining data, alignment pipelines, and architectures. Exploiting such correlations can reduce uncertainty about unplayed arms and lower the effective exploration burden.

LLM routing also provides a source of side information largely absent from standard bandit settings: offline performance data. Benchmarks or historical logs can train machine learning (ML) predictors that map query contexts to per-arm reward estimates. During online routing, these estimates can serve as surrogate rewards for unplayed arms. Although surrogate rewards may be biased or noisy, they can still reveal relative arm quality. The key challenge is to exploit surrogate rewards when helpful without over-trusting them when misspecified. This motivates our central research question:

Can we design contextual bandit algorithms that systematically exploit both context-dependent inter-model correlations and auxiliary ML-predicted surrogate rewards to accelerate learning, while remaining robust to surrogate misspecification?

To address this question, we study contextual bandits with correlation-aware graph feedback and auxiliary surrogate rewards. Each round, the learner observes the query context and context-dependent side information about inter-arm relationships, which are used to form a feedback graph that identifies a small set of additional arms whose feedback is most informative given the chosen arm. After selecting an arm and observing its realized reward, the learner additionally receives ML-predicted surrogate rewards for the arms in this graph-selected set, providing partial multi-arm side information beyond the chosen arm. We first propose a coupled reward-mixing approach, Correlation-Aware Bandits with Surrogates Coupled (CABS-C), which builds on graph-feedback variants of SquareCB [14], [15] by pooling de-biased true rewards with surrogate rewards to fit a single contextual model. This coupling effectively increases the number of informative observations per round, improving the exploration factor from \(K\) to \(K/(m+1)\), where \(m\) is the number of additional arms for which surrogate rewards are revealed each round. The price of coupling these observations is an added sensitivity to surrogate error. For linear bandits, we show the regret scales as \(R_C=\widetilde{O}\!\left(\sqrt{dKT\log(T/d)/(m+1)}+\varepsilon_n \sqrt{TK^2/(m+1)}\right)\), where \(\varepsilon_n\) quantifies the worst-case magnitude of surrogate noise/misspecification (see Assumption 4). The preceding regret bound makes the trade-off explicit: accurate surrogates reduce the effective exploration burden, while large surrogate error can dominate and make tight coupling brittle.

Motivated by this robustness issue, we propose a decoupled prediction-mixing strategy, Correlation-Aware Bandits with Surrogates Decoupled (CABS-D), that 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 [16], [17]. The resulting meta-bandit promises a best-of-both-worlds guarantee, \(R_{D} = O\!\left(\min \left\{R_{S}, R_{C}\right\}\right)\), 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. We evaluate our LLM routing approaches, observing improved sample efficiency and better accuracy–cost trade-offs compared to online bandit and static baselines.

1.0.0.1 Contributions

In summary, our contributions are:

• We introduce correlation-aware contextual bandits with surrogate rewards, where context-dependent inter-arm relationship signals induce a graph-feedback structure and an auxiliary predictor provides surrogate rewards for graph-specified arms beyond the true reward observed for the chosen arm.

• We propose CABS-C (coupled reward mixing) and CABS-D (decoupled prediction mixing), clarifying when tight coupling is beneficial and when decoupling is necessary for robustness.

• We establish regret bounds that explicitly separate correlation-driven gains (reduced effective exploration) from surrogate noise/misspecification, including a best-of-both-worlds guarantee for CABS-D that matches standard contextual bandits in the worst case, while simultaneously achieving an improved regret that scales inversely with the number of surrogate rewards under mild conditions.

• We evaluate our methods on LLM routing benchmarks and demonstrate improved sample efficiency and stronger accuracy–cost trade-offs relative to online bandit baselines and static routing policies.

Ultimately, these results demonstrate that leveraging predicted surrogate rewards offers a powerful framework for LLM routing, yielding substantial regret reductions. Specifically, our approach effectively interpolates between the full-information and bandit settings, with the performance gains smoothly scaling based on the quality of the surrogate rewards and correlation among LLMs. This finding may be of independent interest beyond the specific setting considered here.

2 Related Work↩︎

In this section, we discuss the works most directly related to our approach and highlight the key distinctions from our method. A more comprehensive overview of the multi-armed bandit and LLM routing literature, including detailed comparisons of regret bounds and theoretical guarantees, is deferred to Appendix 8.

2.0.0.1 Contextual Bandits, Graph Feedback, and Expert Aggregation.

Our work builds on contextual bandits such as SquareCB [13], LinearUCB [18], and NeuralUCB [19], and is closest to contextual bandits with graph feedback [14], [15]. Unlike SquareCB-G [14], which observes true rewards for graph neighbors, and SquareCB-UG [15], which learns an unknown graph after action selection, we observe true reward only for the selected arm and receive surrogate rewards from a context-dependent graph known before action selection. CABS-D is related to Hedge/AdaHedge-style expert aggregation [17], [20] and Corral-style bandit masters [16], but adapts specifically between standard bandit feedback and surrogate graph feedback. Auxiliary-feedback methods [21]–[23] use external signals mainly to improve estimates for the played arm, whereas we propagate surrogate feedback across multiple arms. See [24] and Appendix 8 for a broader overview.

2.0.0.2 LLM routing.

LLM routing selects a model for each query under quality, cost, or latency constraints. Prior work includes offline routers trained from preference or performance data [9]–[11], [25]–[27], benchmarks for cost–performance routing [27], [28], and online routing methods based on contextual bandits or dueling bandits [8], [12], [29]–[32]. Existing online routing strategies largely do not explicitly model intrinsic correlations among LLMs (beyond implicit representation sharing), which can limit sample-efficiency when model behaviors are strongly related. Addressing this gap, our correlation-aware routing framework exploits context-dependent inter-arm relationships and integrates auxiliary surrogate signals while remaining robust to surrogate misspecification, yielding improved routing utility across cost regimes.

3 Problem Setting↩︎

We consider an online LLM routing problem setting operating over a finite horizon of \(T\) discrete rounds, indexed by \(t \in[T]:=\{1,2, \ldots, T\}\). The router (learner) has access to a fixed pool of \(K\) distinct LLMs (arms), denoted by the set \(\mathcal{K}=[K]:=\{1,2, \ldots, K\}\). At each round \(t\in[T]\), the learner observes a prompt and its context embedding vector \(\boldsymbol{x}_t\in\mathcal{X}\subseteq\mathbb{R}^d\), where \(\mathcal{X}\) is a bounded context space, and must select an LLM \(i_t\in[K]\) to process the prompt. The selected model, \(i_t\), produces an outcome \(\boldsymbol{o}_t\) and incurs a cost-aware reward \(r_{t,i_t}(\boldsymbol{x}_t)\); when convenient, we write \(r_{t,i_t}\) and define the corresponding loss \(\ell_{t,i_t}=1-r_{t,i_t}\). The reward typically captures generation quality (e.g., accuracy, helpfulness) and trades off it against operational constraints (e.g., latency, token cost). We explicitly bound the reward and loss to the \([0,1]\) interval, which is standard and necessary for regret analysis. Figure 1 provides an overview of this LLM routing task and the proposed CABS-C/CABS-D online learning pipeline.

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 \(\boldsymbol{x}_t\), the learner observes inter-arm side information from an affinity oracle (), represented either as a correlation matrix \(\boldsymbol{R}_t\in\mathbb{R}^{K\times K}\) or as a time-varying, context-dependent directed feedback graph \(G_t\) over \([K]\) with adjacency matrix \(\boldsymbol{A}_t\in\{0,1\}^{K\times K}\), where \(A_{t,(i,j)}=1\) indicates an edge from arm \(i\) to arm \(j\). By default, \(A_{t,(i,i)} = 1\) for all \(i \in [K]\). The side information among LLMs quantifies how performance scales across different models for a specific context.

Upon selecting arm \(i_t\), the learner observes the true realized reward only for the chosen arm \(i_t\). However, for other arms \(j \neq i_t\) specified by the feedback graph, the learner receives surrogate rewards \(s_{t,j}\) from a reward oracle (). The surrogate rewards predicted by the act as a proxy estimate derived from offline data or cross-model evaluation functions. The conditional bias or variance of the surrogate reward \(s_{t,j}\) relative to the true underlying reward \(r_{t,j}\) is assumed to be bounded.

Similar to [13], [14], we assume that the learner has access to a regression solver () that is invoked at each round \(t\) to update arm models from historical input–output data. The \(\mathsf{RS}\) serves as an abstract optimization primitive that handles the algorithmic complexity of empirical risk minimization and can efficiently compute a predictor from a predefined function class \(\mathcal{F}\) (such as linear models or deep neural networks). Formally, given a sequence of context-dependent importance weights \(w_{\tau, j} \ge 0\), the \(\mathsf{RS}\) processes the accumulated feedback history up to the previous round to solve a weighted empirical risk minimization problem for each arm \(i \in [K]\) to update its model: \[\hat{f}_{t,i} = \arg\min_{f \in \mathcal{F}} \sum_{\tau=1}^{t-1} w_{\tau, i} \left( f(\boldsymbol{x}_\tau, i) - \hat{r}_{\tau,i} \right)^2 + \mathcal{R}(f)\]

where \(\hat{r}_{\tau,i}\) represents the importance-weighted estimate at round \(\tau\) computed based on true or surrogate rewards (equal to \(\hat{r}_{\tau,i_\tau}=r_{\tau,i_\tau}/p_{\tau,i_\tau}\) if \(i = i_\tau\), and \(\hat{r}_{\tau,j}=(s_{\tau,j} - \hat{b}_{\tau,j})/(\boldsymbol{A}_{\tau,(:,j)}^\top\boldsymbol{p}_t)\) if \(i\) is revealed by the feedback graph \(G_\tau\)), with \(r_{\tau,i_\tau}\) being the true observed reward, \(s_{\tau, i}\) being the proxy surrogate reward with estimated bias \(\hat{b}_{\tau,j}\), and \(p_{\tau,i_\tau}\) being the probability of selecting arm \(i\) at round \(\tau\). Here, \(\mathcal{R}(f)\) is an optional regularization term (e.g., \(\ell_2\) penalty) used to prevent overfitting in high-dimensional context spaces. The learner’s objective is to minimize the cumulative regret over \(T\) rounds. The optimal arm at round \(t\) is defined as the arm that maximizes the true reward for the given context: \(i_t^\star = \arg\max_{i \in [K]} \mathbb{E}\left[r_{t,i}(\boldsymbol{x}_t)\right]\). The cumulative regret \(\mathrm{Reg}(T)\) is given by \(\mathrm{Reg}(T) = \sum_{t=1}^T \mathbb{E}\left[ r_{t,i_t^\star}(\boldsymbol{x}_t) \right]- \mathbb{E}\left[r_{t,i_t}(\boldsymbol{x}_t) \right]\), which we use to measure the performance of online routing algorithm.

4 The Proposed Algorithms↩︎

In this section, we propose two algorithms for contextual bandits with surrogate rewards and correlated arms. We first introduce a reward-mixing method, CABS-C, that learns from both true rewards and de-biased surrogate feedback. We then introduce a prediction-mixing method, CABS-D, that dynamically integrates predictions from CABS-C with those from a standard bandit-feedback method (SquareCB [13] or LinUCB [33]) to achieve best-of-both-worlds regret guarantees.

4.1 A coupled algorithm with provable regret bound↩︎

We first propose Correlation-Aware Bandits with Surrogates Coupled (CABS-C), detailed in Algorithm [alg:CABS-C]. CABS-C builds on graph-feedback variants [14], [15] of SquareCB by using the feedback graph to guide exploration, but adapts this framework to the surrogate-feedback setting. Specifically, 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. Each arm is modeled as a function \(f\) which could be linearly parametrized by \(\hat{\boldsymbol{\theta}}_i, i \in[K]\) or implemented as a neural net. At round \(t\), the algorithm observes context \(\boldsymbol{x}_t\) and an affinity matrix \(\boldsymbol{R}_t\), samples a stochastic feedback graph \(G_t\) with adjacency matrix \(\boldsymbol{A}_t\), and obtains reward estimates from the regression solver . The action distribution \(\boldsymbol{p}_t \in \Delta([K])\), where \(\Delta([K])\) is \((K-1)\)-dimensional simplex, is then computed by solving a optimization problem (Line [line:cabsc-predict-end]), which chooses an arm by balancing low immediate regret with high information gain from the feedback graph. Unlike a standard contextual bandit rule, which only trades off reward and uncertainty for the selected arm, this step explicitly accounts for additional side observations revealed through graph neighbors; thus, it may select an arm that is slightly suboptimal in immediate reward if it yields more informative feedback overall. The parameter \(\gamma\) controls this trade-off.

After sampling an arm \(i_t\) according to \(\boldsymbol{p}_t\), the learner observes its true reward and receives surrogate rewards \(\{s_{t,j}\}\) for graph neighbors \(\mathcal{N}_t(i_t):=\{j\in[K]:A_{t,(i_t,j)}=1\}\), where each \(s_{t,j}\in[0,1]\) is produced by the reward oracle RO. CABS-C then de-biases the surrogate feedback, constructs importance-weighted reward estimates, and updates RS. Thus, CABS-C interpolates between bandit and full-information learning by using context-dependent graph structure to propagate imperfect surrogate feedback beyond the selected arm. While graph-based side observations and surrogate feedback have each been studied independently in prior work, CABS-C jointly addresses both in a single correlation-aware bandit framework.

We now turn to establishing the regret bound for the algorithm. Before, we state the assumptions in our setting below.

Note that this assumption holds for any function class. In our analysis, we stick to linear functions, i.e, \(\hat{f}_t(\boldsymbol{x}_t,i) = \boldsymbol{x}_t^{\top} \hat{\boldsymbol{\theta}}_{t,i}\) and \(f^\star(\boldsymbol{x}_t,i) = \boldsymbol{x}_t^{\top}\boldsymbol{\theta}^*_{i}\).

Following [34], since we assume that pulling an arm always reveals its own loss, we satisfy \(i \in \{ j \in [K] ;\; {A}_{t,(j,i)} = 1 \}\), making the graph strongly observable.

Here filtration \(\mathcal{F}_t\) denotes the information available up to time \(t\), and \(\varepsilon_n\) bounds the surrogate noise across rounds and observed arms. Intuitively, \(b_i\) captures systematic surrogate over- or under-estimation for arm \(i\), while larger \(\varepsilon_n\) corresponds to noisier and less informative side observations.

Using the above assumptions, we introduce a lemma for concentrating the noise term of surrogate rewards, and then we introduce a theorem and accompanying corollary for bounding the expected regret of Algorithm [alg:CABS-C]. The proofs are deferred to Appendix 13.2.

For an \(m\)-feedback graph, where each round reveals \(m\) surrogate rewards, Theorem 1 yields the following linear-oracle specialization.

Brief discussion of bias and noise. If we ensure that \(\varepsilon_n = O(1/\sqrt{K})\), we recover the same asymptotic regret (up to polylogs) as the standard contextual bandit with graph feedback setting, \(\widetilde{O}\!\left(\sqrt{\alpha T \, \mathrm{Reg_{Sq}}(T)}\right)\). If the bias and noise are well-behaved (each arm’s respective bias \(b_{i}\) remains relatively constant or changes gradually and noise is small and truly subgaussian), then CABS-C can utilize surrogate rewards to significantly improve the rate of learning (the additional information from the surrogate rewards reduces variance). However, if bias is difficult to estimate or if the magnitude of the subgaussian noise is large, then the surrogate rewards available to CABS-C may hinder its ability to learn (inconsistent surrogate rewards can potentially increase variance).
On bounding \(\boldsymbol{\varepsilon_n}\). We note that the ML reward predictor (), trained offline from a fixed mixture of LLMs, enjoys standard ERM generalization guarantees only with respect to the training distribution [36], [37]. Generalization to unseen contexts reduces to a covariate shift problem [38], [39], and fundamentally depends on support overlap or structural invariance assumptions on the underlying data-generating process. In the next section, we propose an algorithm that is robust to such noise, while retaining the aforementioned regret bound when \(\varepsilon_n\) is small.
Cost-informed utility. To capture the accuracy–cost trade-off inherent in practical LLM deployment, we define a cost-aware utility: \(r_{t,i_t} = y_{t, i_t} - \rho\, c_{t, i_t}\), where \(y_{t, i_t} \in [0,1]\) denotes the accuracy score obtained by using LLM \(i_t\) to answer query \(\boldsymbol{x}_t\), and \(c_{t, i_t} \ge 0\) is the associated inference cost. The parameter \(\rho \ge 0\) controls cost sensitivity: larger values favor cheaper models, while smaller values prioritize accuracy. We additionally clip the utility to \([0,1]\) to enforce bounded rewards.
From linear to neural bandits. While we focus on linear models for clarity of exposition and analysis, our approach extends naturally to nonlinear reward models. In particular, following the machinery (Neural Tangent Kernel matrix analysis) adapted in NeuralUCB [19], the linear predictor \(\boldsymbol{x}_t^\top \boldsymbol{\theta^*}_{i}\) can be replaced by a neural network with uncertainty estimated via its last-layer features. At this point, we empirically adapt our method to neural bandits based estimators and show their performance in Section 5, but note that our regret analysis can be extend to this setting be utilizing the machinery developed in [19].

4.2 A decoupled method with optimal regret↩︎

Theorem 1 shows that surrogate rewards can substantially improve regret for CABS-C algorithm, but can also degrade performance when their bias or noise is poorly controlled. This is a consequence of reward mixing: CABS-C combines true and surrogate rewards at the observation level and trains a single model on the pooled feedback. To obtain robustness to unreliable surrogate rewards, we instead consider a prediction-mixing strategy, where different feedback mechanisms are handled by separate learners and combined only at the decision level.

We introduce Correlation-Aware Bandits with Surrogates Decoupled (CABS-D), shown in Algorithm 2. CABS-D maintains two stateful experts: a standard bandit-feedback expert and the surrogate-aware graph-feedback expert CABS-C from Algorithm [alg:CABS-C]. In our theoretical analysis, the standard expert is SquareCB, described in Appendix 13.1. At each round, the experts output policies \(\boldsymbol{p}_{t,1}\) and \(\boldsymbol{p}_{t,2}\), which are mixed into \(\boldsymbol{q}_t\) using adaptive exponential weights over meta-copies \((m,g)\), in the spirit of Hedge/AdaHedge-style aggregation [17]. After sampling \(i_t\sim\boldsymbol{q}_t\), SquareCB receives only the selected arm’s true reward, while CABS-C also receives surrogate rewards from the sampled feedback graph. The meta-weights are updated using importance-weighted loss estimates and an adaptive second-order learning rate, allowing CABS-D to track the better feedback mechanism without knowing whether surrogates are reliable. Thus, CABS-D retains standard bandit robustness while exploiting informative surrogate feedback. We now state its regret bound.

The proof can be found in Appendix 13.3. While the theorem establishes a general regret guarantee for the algorithm, the following corollary shows that surrogate rewards strictly improve the regret bound.

In a standard linear contextual bandit setting (which corresponds to having no extra surrogate feedback, or \(m = 0\)), the typical optimal expected regret bound over \(T\) rounds is \(\mathbb{E}[\text{Reg}_T] \le \widetilde{O}\left(\sqrt{dKT }\right)\). This baseline matches the second term inside the \(\min\{\cdot\}\) operator of the regret in corollary. Because the overall bound is a minimum between two strategies, the algorithm guarantees performance that is at least as good as the vanilla baseline, but can be significantly better depending on the value of \(m\) and the noise \(\varepsilon_n\). Specifically, for small enough \(\varepsilon_n\) and as long as \(m > 0\), the first term inside the minimum takes effect where the second term in the sum vanishes. We note that the first term in summand remains the dominant factor, as long as \(\varepsilon_n\) satisfies the condition \(\varepsilon_n < \tilde{O}(\sqrt{{d}/{K}})\), which can be easily satisfied by training the offline surrogate model with enough samples. In this case, the regret is entirely dominated by \(\widetilde{O}\left(\sqrt{d\left(\frac{K}{m+1}\right)T }\right)\), yielding a significant improvement over the vanilla \(\widetilde{O}\left(\sqrt{dKT}\right)\) bound as the effective action-space dependence drops from \(K\) to \({K}/({m+1})\). This significantly shrinks both the primary exploration term and the noise propagation term, allowing the algorithm to learn the underlying linear reward structure much faster than standard bandit feedback would allow.

5 Experiments↩︎

In this section, we conduct experiments to answer four research questions:

• How do coupled reward mixing (CABS-C) and decoupled prediction mixing (CABS-D) compare across cost sensitivities?

• Which routing method achieves the highest utility across cost regimes and datasets?

• Do correlation-aware surrogate-reward methods improve the accuracy–cost Pareto frontier relative to standard online bandit baselines?

• Do surrogate rewards accelerate online learning, as measured by reduced cumulative regret over time?

6pt

5.1 Experimental Setup↩︎

5.1.0.1 Datasets.

We evaluate on three LLM routing benchmarks: RouterBench [28], SPROUT [27], and Open LLM Leaderboard v2 [40]. Dataset details and splits are deferred to Appendix 10.

5.1.0.2 Baselines.

We compare against standard online contextual bandit algorithms, including LinUCB [18], LinTS [41], and \(\epsilon\)-greedy. In the LLM routing setting, observing rewards for unplayed arms is generally unrealistic, as only the selected model is executed. Accordingly, we primarily evaluate methods under bandit feedback (single observed reward per round). For reference, we additionally report hypothetical full-feedback variants (LinUCB/NeuralUCB with full feedback) as information upper bounds. We also include strong supervised routing baselines, such as RoBERTa/CARROT [27] and kNN-based routers [28], as well as random selection and single-model baselines.

5.1.0.3 Metrics.

We report relative cumulative regret and average utility over time. Specifically, the relative cumulative regret of an online learner ‘\(\textit{alg}\)’ is defined with respect to ‘\(olg\)’ as \(R_{\textit{alg}}(t) - R_{\text{olg}}(t)\), where \(R_{\textit{alg}}(t)\) denotes the cumulative regret at time \(t\). In addition, we summarize test accuracy and average per-query cost, and characterize accuracy–cost trade-offs via the accuracy–cost Pareto frontier by sweeping the cost-sensitivity parameter for each method.

5.1.0.4 Setup.

We study LLM routing under varying accuracy and cost trade-offs by sweeping the cost-sensitivity parameter over \(\rho\in\{0,50,\ldots,1000\}\), which we group into three cost regimes: low cost sensitivity corresponds to small values of \(\rho\) (\(\le 300\)) where accuracy dominates, medium cost sensitivity (\(300 < \rho < 800\)) represents a balanced trade-off, and high cost sensitivity (\(\rho \ge 800\)) corresponds to large values of \(\rho\) where cost dominates. For correlation-aware methods, the Reward Oracle () imputes surrogate utilities for graph-neighbor arms, while the Affinity Oracle () constructs the context-dependent feedback graph from predicted inter-arm similarity. Both are lightweight MLP heads on top of mDeBERTaV3-base query embeddings [42]. Refer to Appendix 11 for implementation details.

5.2 Reward Mixing vs Prediction Mixing (RQ1)↩︎

Table 1 compares the coupled (CABS-C) and decoupled (CABS-D) variants by reporting average utility (higher is better) and average regret (lower is better) across three cost regimes in SPROUT and RouterBench dataset (see Appendix 12 for Open LLM LB V2 dataset). Across all regimes, CABS-D, i.e, model mixing, consistently improves utility and reduces regret relative to CABS-C, indicating that decoupling reward learning from surrogate learning yields a more reliable use of surrogate signals. This is aligned with our theory: while coupling can amplify surrogate noise or bias, the decoupled construction can exploit informative surrogate predictions early while falling back to reward-based learning as data accumulates. Based on these results, we use CABS-D as the default instantiation in the remainder of our experiments.

6pt

5.3 Utility across cost regimes (RQ2)↩︎

Table 2 reports average utility (higher is better) across three routing benchmarks under Low/Mid/High cost-sensitivity regimes. Across all datasets and regimes, CABS-D is the strongest implementable method: either its linear or neural instantiation attains the best non-oracle utility across cost regimes. Relative to standard online bandit baselines, CABS-D yields consistent gains (\(\approx 8\%\) on Open LLM Leaderboard v2, \(\approx 3\%\) on RouterBench, and \(\approx 1\%\) on SPROUT), with the largest improvements on the more heterogeneous Open LLM Leaderboard v2 (\(\approx 15\%, 9\%, 7\%\) in the Low/Mid/High regimes). On RouterBench and SPROUT, CABS-D continues to dominate online baselines and remains competitive with strong static routers (CARROT/KNN), showing that online correlation-aware learning can match or exceed fixed routers while adapting to the target cost regime. Full-feedback variants provide an information upper bound, but CABS-D closes a substantial portion of this gap under realistic bandit feedback.

5.4 Accuracy vs Cost Trade-off (RQ3)↩︎

Figure 3 shows the trade-off between accuracy and average per-query cost on Open LLM Leaderboard v2, RouterBench, and SPROUT. Each curve traces a routing policy as the cost-sensitivity parameter \(\rho\) varies \(\in [0,1000]\). Across all datasets, CABS-D consistently improve the accuracy versus cost Pareto frontier compared to standard online bandit baselines. In particular, CABS-D achieves higher accuracy at comparable cost levels (or equivalently lower cost for a given accuracy) relative to LinUCB, LinTS, and \(\epsilon\)-greedy. While full-feedback variants provide an optimistic upper bound, CABS-D substantially narrows this gap under realistic bandit feedback, answering RQ3 affirmatively.

5.5 Learning dynamics (RQ4)↩︎

Figure 4 reports relative cumulative regret over time normalized with respect to LinUCB, across low, mid, and high cost-sensitivity regimes for Open LLM Leaderboard V2. Across all regimes, CABS-D consistently achieves lower cumulative regret than standard bandit baselines including LinUCB, LinTS, and \(\epsilon\)-greedy. The improvement is most pronounced in the mid-cost regime, where the routing objective requires balancing accuracy and cost, and surrogate rewards provide the greatest benefit. Notably, the regret gap emerges early and persists throughout the horizon, indicating that surrogate rewards accelerate online learning rather than merely improving asymptotic performance.

6 Discussions↩︎

We discuss the main trade-offs between coupled and decoupled use of surrogate rewards, the computational cost of the proposed algorithms, and practical deployment considerations for LLM routing.

6.0.0.1 Comparison of Coupled and Decoupled Approaches:

CABS-C and CABS-D differ primarily in how tightly they integrate surrogate feedback with true bandit feedback. CABS-C follows a coupled strategy: after debiasing surrogate rewards and applying importance weighting, it feeds both true and surrogate observations into a common regression oracle. This coupling can improve sample efficiency when the surrogate rewards are informative, since each round provides feedback for multiple correlated arms. However, because the two feedback sources update the same predictor, persistent surrogate bias or high surrogate noise can still affect the learned model, even when importance weighting partially controls the contribution of each observation. CABS-D addresses this limitation by maintaining separate learners for the standard bandit-feedback signal and the surrogate-aware feedback signal, and then adaptively mixing their induced policies. As a result, when surrogate rewards are useful, CABS-D can exploit them for faster learning; when they are uninformative or too noisy, it can fall back on the standard contextual bandit expert and retain the corresponding regret guarantee \(\widetilde{O}(\sqrt{KT,\mathrm{Reg}_{\mathrm{Sq}}(T)})\).

6.0.0.2 Computational Complexity:

This robustness comes with only a modest additional computational cost. CABS-C has complexity \(O(K^3 + \text{\sffamily RS})\), while CABS-D has \(O(K^3 + K\ln T + \text{\sffamily RS})\). The \(K^3\) term comes from the optimization in line 8 of Algorithm [alg:CABS-C], and the extra \(K\ln T\) term in CABS-D comes from the dyadic grid in Algorithm 2. For a linear oracle, \(O(\text{\sffamily RS})=O(Kd+d^2)\). In contrast, standard linear bandit methods such as LinUCB/LinTS are typically dominated by matrix operations cubic in the feature dimension \(d\), \(O(Kd^3)\). Therefore, CABS-C and CABS-D may be preferable in settings with large \(d\) or more general nonlinear function classes, whereas their stronger dependence on \(K\) can make them less attractive when the number of actions is very large.

6.0.0.3 Surrogate Reward Prediction Overheads:

Beyond the learner update, the additional cost of producing surrogate rewards is small in our implementation. The reward oracle is a 2-layer MLP that reuses the context embeddings already computed for the bandit, together with the chosen arm embedding and observed reward. In batch-size-1 measurements, embedding generation takes 12.32 ms on average, while surrogate prediction itself takes only 0.22 ms on average. Thus, the incremental overhead from surrogate prediction beyond the encoder is negligible.

6.0.0.4 Extending Utility:

Our current utility objective can already capture latency when latency is embedded in the cost term, for example through a deployment-specific cost that combines monetary price, compute, and response time. More explicitly, one can extend the utility as \(r_{t,a} = s(\mathbf{x}_t,a) - \rho\, c(\mathbf{x}_t,a) - \rho_{1}\, l(\mathbf{x}_t,a),\) where \(s(\mathbf{x}_t,a)\) is the accuracy score, \(c(\mathbf{x}_t,a)\) is the monetary cost, and \(l(\mathbf{x}_t,a)\) is the latency of arm \(a\) for context \(\mathbf{x}_t\). Another promising direction is to combine our framework with delayed-feedback bandits [43], which could model latency more directly.

7 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 bandit-plus-surrogate 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.

Acknowledgment↩︎

This material is based upon work supported by the National Science Foundation (NSF) under Grant Number 2318101. Any opinions, findings, and conclusions or recommendations expressed in this material are those of the author(s) and do not necessarily reflect the views of the National Science Foundation.

Appendix↩︎

The appendix is organized as follows,

• In Appendix 8 we discuss additional related works

• In Appendix 9 we expand on the baselines that are used in our experiments.

• In Appendix 10, we expand on the LLM routing datasets.

• In Appendix 11, implementation details are discussed.

• In Appendix 12, we present additional experiments and results.

• In Appendix 13, we present the proofs for CABS-C and CABS-D algorithms.

8 Additional Related Works↩︎

8.1 Multi-Armed Bandits (MAB)↩︎

Classical stochastic bandits study decision-making with partial feedback and no context. UCB-type methods [44] use optimism under uncertainty, while Thompson sampling [45] samples actions from a posterior distribution. Adversarial algorithms such as Exp3 [46] provide guarantees without stochastic reward assumptions. These methods form the basis for exploration–exploitation trade-offs but do not capture query-dependent model selection.

Contextual bandits incorporate side information into the decision rule. Linear methods such as LinUCB [18], [33] assume rewards are linear in context features, while KernelUCB [47] and NeuralUCB [19] model nonlinear reward functions. More generally, SquareCB [13] reduces contextual bandit learning to online regression by using a reward predictor to construct exploration probabilities, providing a flexible oracle-based template beyond linear models. These methods are natural baselines for LLM routing, but typically observe feedback only for the selected arm and do not explicitly exploit inter-arm relationships.

Bandits with graph feedback study settings where playing one arm reveals additional observations from neighboring arms. Earlier work considers non-contextual graph-feedback models [48], [49], while SquareCB-G [14] and SquareCB-UG [15] extend this idea to contextual settings. SquareCB-G assumes graph-structured observations of true rewards, whereas SquareCB-UG handles graphs that are not known before action selection. Our setting differs because the graph is context-dependent and available before action selection, while additional observations are potentially noisy/misspecified surrogate rewards rather than true rewards.

Auxiliary-feedback bandits use external predictors or side observations to improve learning. MIN-UCB [21], OFUL-AF [22], and MLA-UCB [23] are representative examples. These methods typically use auxiliary information to improve reward estimation for the played arm or to regularize learning. Our setting instead uses context-dependent graph structure to decide which unplayed arms receive surrogate feedback, requiring robustness to surrogate noise and bias.

Expert-aggregation methods provide another relevant perspective for our decoupled algorithm. Hedge and multiplicative-weight algorithms aggregate expert predictions with regret guarantees against the best fixed expert in hindsight [20], while adaptive variants such as AdaHedge tune the learning rate online based on observed losses [17]. Related master algorithms, such as Corral [16], combine multiple bandit algorithms and compete with the best base learner under partial feedback. CABS-D follows this meta-algorithmic view, but the experts correspond to different feedback mechanisms: standard bandit feedback and correlation-aware surrogate graph feedback.

Table 3 summarizes representative bandit settings, feedback models, and regret rates. We report worst-case order bounds, hiding polylogarithmic factors in \(\widetilde{O}(\cdot)\), under standard assumptions; some algorithms admit sharper gap-dependent bounds, which we omit for brevity.

8.2 LLM Routing↩︎

Here we give a high level overview of selected LLM routing methods,

• RouterBench [28] introduces a standardized benchmark/dataset and evaluation protocol for multi-LLM routers under explicit cost–performance trade-offs, enabling apples-to-apples comparison across routing strategies.

• RouterDC [25] proposes a query-based router trained via dual contrastive learning, aligning query embeddings with top-performing LLM embeddings (and separating from poorly performing ones) using sample–LLM and sample–sample contrastive losses to improve routing stability and robustness.

• CARROT / SPROUT [27] propose a cost-aware router based on estimating model cost/performance, and introduce the SPROUT dataset to evaluate routing over a wide query distribution and modern model pools.

• RouteLLM [9] trains a router from preference supervision to decide between cheaper and stronger LLMs, targeting cost savings at fixed quality levels and demonstrating transfer across model pairs.

• Eagle [10] proposes a training-free router that uses ELO-style ranking (global and local) to efficiently update routing decisions, aiming at scalable online serving without heavy retraining.

• TensorOpera / TO-Router [11] studies multi-model orchestration as a system problem and routes queries to expert LLMs to improve the throughput–cost–quality “trilemma” in practical deployments.

• GraphRouter [26] models query–task–LLM interactions as a heterogeneous graph and performs edge/attribute prediction to recommend an LLM, emphasizing generalization to new LLMs via relational inductive bias.

• MetaLLM [8] frames routing as a contextual bandit and learns an online selection policy from partial feedback to improve cost-efficiency on targeted tasks (e.g., classification / MCQ).

• LLM Bandit [12] formulates LLM selection as a bandit problem with preference-conditioned routing, allowing users to specify trade-offs at inference while learning under partial-feedback constraints.

• MixLLM [29] extends contextual-bandit routing to dynamic objectives (quality/cost/latency) with continual learning and changing model pools, using lightweight predictors plus a meta decision module.

• BaRP [30] trains routers to align with the deployment partial-feedback setting while conditioning on a user preference vector, enabling a single learned policy to realize multiple cost–quality trade-offs at test time.

• LLM Routing with Dueling Feedback [32] treats routing as contextual dueling bandits and learns from pairwise preferences between two candidate model outputs, using CCFT representations and FGTS.CDB for online learning.

9 Baselines↩︎

• LinUCB [18] / MetaLLM [8]: the standard linear upper-confidence bound algorithm, which selects the arm with the highest optimistic estimate of the linear reward.

• LinUCB [18] + Full Feedback: an oracle variant of LinUCB that observes the full reward vector for all arms at every round, rather than only the reward of the chosen arm. This baseline is not implementable in practice but provides a strong upper bound on what is achievable with linear models.

• LinTS [41]: Linear Thompson sampling with a Gaussian posterior over the regression parameters.

• Epsilon Greedy: a linear regression model updated online together with an \(\epsilon\)-greedy exploration policy.

• NeuralUCB [19]: a neural contextual bandit algorithm that uses a neural network reward predictor and an upper-confidence bonus based on the network’s linearized features (e.g., via the last-layer/NTK approximation), selecting the arm with the highest optimistic predicted reward.

• NeuralUCB [19] + Full-Feedback: an oracle variant of NeuralUCB that observes the full reward vector for all arms at each round and updates the neural predictor using all arms’ rewards, providing an upper bound on what is achievable with neural reward models under full information.

All online learning baselines use exactly the same feature representation as CABS-C/CABS-D and differ only in their exploration strategy and feedback model. We additionally compare against static routers that do not adapt during test time:

• CARROT [27] (RoBERTa): we adapt the current SOTA CARROT router to our utility objective. Instead of predicting the probability of success, we fine-tune a RoBERTa encoder as a regressor to predict, for each (query, model) pair, both the correctness reward and the cost. At test time, the router selects the model with the largest predicted utility, \(\hat{r}_{t, i_t} = \hat{y}_{t,i_t} - \rho \cdot \hat{c}_{t, i_t}\).

• KNN (RouterBench [28]): a non-parametric router operating on text embeddings. For each test query, we retrieve the \(k=3\) nearest training queries (in embedding space) for each candidate model and estimate reward and cost by averaging over their observed outcomes. The chosen model maximizes the estimated utility \(\hat{r}_{t, i_t}\).

• Random Selection: For each query, this method randomly selects an arm from the K possible arms.

• Finally, we report the performance of individual LLMs used by the benchmarks as static baselines that always route to a single model. This highlights the benefit of routing compared to selecting a single fixed model.

10 Dataset Details↩︎

• RouterBench [28] contains roughly 35k text queries spanning 8 reasoning, coding, and commonsense datasets, with responses from 11 LLMs annotated using GPT-4 LLM-as-judge binary correctness and token-based cost. We randomly split the dataset into 70% train, 10% validation, and 20% test.

• SPROUT [27] contains roughly 44k queries drawn from six reasoning datasets, with responses from 15 LLMs annotated with cost and numeric correctness scores using LLama 3.1 70B LLM-as-judge. We use the default train/val/test split.

• Open LLM Leaderboard v2 [40] provides responses from 18 contemporary LLMs for routing evaluation. The scores are generated using exact match-based accuracy. We randomly split the dataset into 70% train, 10% validation, and 20% test.

We also show the average performance of each LLMs in the three LLM Routing datasets (Table 4, Table 5 and Table 6).

11 Implementation Details↩︎

All experiments are conducted on a machine equipped with 1xNVIDIA H100 GPU (80 GB). For all linear online learning methods, we set the regularization parameter to \(\lambda = 10\) and \(\gamma = 4\) whenever applicable. For neural online learning methods, we set regularization parameter to \(\lambda = 0.001\) and the hidden dimension to 64, with learning rate of \(0.001\). Reported results for all online learning methods are averaged over 10 independent trials to account for randomness in data ordering. For \(\varepsilon\)-greedy, we use the decay schedule, \(\epsilon_t=\max\!\left(\epsilon_{\min}, \frac{\epsilon_0}{1+\text{decay}\cdot t}\right)\), tune \(\epsilon_0 \in \{1.0, 0.5, 0.2, 0.1\}\) per dataset, and fix \(\text{decay}=0.001\) and \(\epsilon_{\min}=0.01\). For LinUCB/NeuralUCB with full feedback, we update using the ground-truth rewards of all arms at each round; since all rewards are observed, we set the exploration parameter to zero.

11.0.0.1 Affinity and Reward Oracles.

We implement the Reward Oracle (RO) and Affinity Oracle (AO) using a lightweight multi-head network on top of query embeddings \(\boldsymbol{x}_t\). Each query is encoded once, using mDeBERTaV3-base [42] in our experiments, and passed to small MLP-style heads (using Tanh/GELU/ReLU nonlinearities with normalization/dropout). The framework is otherwise encoder-agnostic.

The RO predicts the accuracy and cost components separately and combines them into the cost-informed utility \(r_{t,j}=y_{t,j}-\rho c_{t,j}\). For query embedding \(\boldsymbol{x}_t\) and anchor arm \(i_t\), we define the accuracy delta \(\delta_{t,j}(i_t)=y_{t,j}-y_{t,i_t}\) for all \(j\in[K]\). A correlation-aware propagation head predicts \(\widehat{\delta}_{t,j}(i_t)=g(\boldsymbol{x}_t,i_t,y_{t,i_t})\), yielding surrogate accuracies \(\tilde{y}_{t,j}=y_{t,i_t}+\widehat{\delta}_{t,j}(i_t)\) after observing \(y_{t,i_t}\). In parallel, a cost head predicts \(\widehat c_{t,j}=h(\boldsymbol{x}_t,j)\) (trained with MSE on \(z\)-scored costs), and the two are combined as \(s_{t,j}=\tilde{y}_{t,j}-\rho \widehat c_{t,j}\). These surrogate utilities are then filtered to generate additional feedback for correlation-aware updates.

The AO is constructed from a context-only accuracy prediction head, which outputs per-arm scores \(\bar y_{t,j}\) for all \(j\in[K]\). Since estimating a true context-conditional arm correlation matrix would require repeated joint observations for the same query, we instead use the proxy \(R_{t,(i,j)}=1-2|\bar y_{t,i}-\bar y_{t,j}|\), which is high for arms predicted to behave similarly and low, possibly negative, for arms predicted to differ sharply. This proxy defines the feedback graph used by our correlation-aware methods, although our algorithms and theory allow any context-dependent choice of \(\boldsymbol{R}_t\). To enforce self loops, we set \(R_{t,(i,i)} = -\infty\).

In addition to the cost, propagation, and context-only accuracy heads, the network includes a bandit feature head that maps \(\boldsymbol{x}_t\) to a 128-dimensional representation for downstream bandit algorithms. The heads are trained offline using Adam with learning rate \(10^{-4}\), MSE for cost prediction, and \(L_1\)/Huber-style losses for accuracy prediction and propagation. Since we sweep large cost sensitivities up to \(\rho=1000\), we feed the learner a bounded utility signal \(\bar r_{t,i_t}=\mathrm{clip}(r_{t,i_t},0,1)\) for numerical stability. Finally, although RO can produce surrogate utilities for all graph-neighbor arms, we only incorporate pseudo-observations for a pruned subset selected by an IQR rule on predicted delta magnitudes, retaining arm \(j\) only if \(|\widehat{\delta}_{t,j}|>Q_3+\kappa\,\mathrm{IQR}\) with \(\kappa=1.5\). This acts as a conservative sparsification guard without changing the underlying feedback model.

11.0.0.2 Efficient Computation of \(\boldsymbol{p}_t\) (Line 8 of Algorithm [alg:CABS-C]):

Line 8 of Algorithm [alg:CABS-C] requires solving for the sampling distribution \(\boldsymbol{p}_t\in\Delta(K)\) that minimizes the objective \(\phi(\boldsymbol{p};\hat{\boldsymbol{\theta}}_t,\boldsymbol{x}_t,\boldsymbol{A}_t)\). Rather than optimizing this objective in its original form, we follow the convex reformulation of [14], which converts the problem into an equivalent convex program over \((\boldsymbol{p},z)\) that can be solved efficiently with standard solvers at each round, as stated in the following theorem.

We refer readers to [14] for the proof. Although Theorem 3 is stated in terms of loss scores \(\hat{f}(\boldsymbol{x}_t,j)\), our bandit models produce reward predictions \(\hat{r}_{t,j}\in[0,1]\). We therefore apply the theorem with the transformation, \(\hat{f}(\boldsymbol{x}_t,j)\;=\;1-\hat{r}_{t,j}\). In particular, in the linear case \(\hat{r}_{t,j}=\boldsymbol{x}_t^\top \hat{\boldsymbol{\theta}}_{t,j}\), while in the neural case \(\hat{r}_{t,j}=f_{\hat{\boldsymbol{\theta}}_t}(\boldsymbol{x}_t,j)\) denotes the prediction of the corresponding neural reward model. Thus, \(\hat{\boldsymbol{f}}(\boldsymbol{x}_t,\cdot)\in\mathbb{R}^K\) is obtained component-wise from \(\hat{\boldsymbol{r}}_t(\boldsymbol{x}_t,\cdot)\in\mathbb{R}^K\) via \(\hat{\boldsymbol{f}}_t=\mathbf{1}-\hat{\boldsymbol{r}}_t\). Since \(\boldsymbol{p}\in\Delta(K)\) implies \(\boldsymbol{p}^\top \mathbf{1}=1\), this substitution only changes the objective by an additive constant (absorbed into the slack variable \(z\)) and leaves the optimizer \(\boldsymbol{p}_t\) unchanged. We use Theorem 3 to compute \(\boldsymbol{p}_t=\arg\min_{\boldsymbol{p}\in\Delta(K)}\phi(\boldsymbol{p};\hat{\boldsymbol{\theta}}_{t},\boldsymbol{x}_t,\boldsymbol{A}_t)\) in Algorithm [alg:CABS-C].

12 Additional Graphs and Results↩︎

12.1 Reward Mixing vs Prediction Mixing (CABS-C vs CABS-D)↩︎

Table 7 presents the same ablation study on Open LLM LB V2 [40]. The overall pattern is consistent with the findings on RouterBench [28] and SPROUT [27] as shown in Section 5.2. Across most cost regimes, the decoupled variant (CABS-D) achieves comparable or improved utility relative to the coupled variant (CABS-C), and generally attains lower regret in the linear setting. On Open LLM LB V2, the differences are smaller but remain directionally similar, with CABS-D typically matching or slightly outperforming CABS-C. Overall, these results support the conclusion drawn in the main text that decoupling reward learning from surrogate prediction provides a stable and effective alternative to direct reward mixing.

6pt

6pt

12.2 Encoder Robustness↩︎

We repeat our experiments with a second, widely used text embedding model (all-MiniLM-L6-v2 [52]) in addition to mDeBERTaV3-base [42]. Across datasets, method rankings and improvements over bandit baselines with respect to utility is consistent (see Table 8), indicating that our conclusions do not depend on a specific encoder choice. Absolute utilities shift slightly, with the largest change on RouterBench in the high-cost regime; notably, the full-feedback upper bound decreases similarly, suggesting the shift is driven by representation-limited separability rather than algorithmic instability.

12.3 Learning Dynamics & Accuracy vs Cost Trade-offs↩︎

Figure 5 extends the learning-dynamics analysis of Section 5.5 to RouterBench and SPROUT, reporting relative cumulative regret normalized by LinUCB across low (\(\rho \le 300\)), mid (\(300<\rho<800\)), and high (\(\rho \ge 800\)) cost regimes. Consistent with the Open LLM Leaderboard v2 results in the main text, CABS-D achieves uniformly lower relative cumulative regret than standard bandit baselines in most regimes, and the separation typically appears early in the horizon, indicating faster learning from surrogate feedback rather than solely an asymptotic gain. The trends are especially stable in the mid-cost regime, where the objective most strongly couples accuracy and cost and thus benefits from exploiting inter-arm correlations. As expected, full-feedback variants remain an optimistic upper bound and provide a reference for the achievable regret reduction under complete feedback. Further, Figure 7 shows relative cumulative regret graphs for Neural variants where we observe similar trends as that of linear variants.

Figure 6 provides the analogous accuracy–cost Pareto frontiers for neural online-learning variants. The qualitative behavior mirrors the linear case: correlation-aware surrogate methods shift the frontier outward, achieving higher accuracy at comparable cost or reducing cost for a fixed accuracy, with the largest advantages occurring in regimes where cost sensitivity meaningfully constrains the policy.

13 The Technical Proofs↩︎

For clarity, we decompose the analysis of regret of proposed algorithms into three components, which progressively build toward our final guarantee. In particular, for the ease of reading, we split our algorithmic analysis into three parts as follows:

1. Baseline: Contextual Bandit In Subsection 13.1, we first present a classical contextual bandit with regret bound \(\widetilde{O}\!\left(\sqrt{KT\mathrm{Reg}_{\mathrm{Sq}}(T)}\right)\). In particular for the linear case, we get \(\widetilde{O}\!\left(\sqrt{dKT \log \left(\frac{T}{d}\right)}\right)\).

2. Correlation-Aware Bandits with Surrogates — Coupled (CABS-C). We next analyze our proposed graph-feedback algorithm with noisy surrogate rewards in Subsection 13.2. When feedback graph \(G_t\) has independence number no greater than \(\alpha\), the regret satisfies \[\widetilde{O}\!\left( \sqrt{\alpha T\, \mathrm{Reg}_{\mathrm{Sq}}(T)} + \varepsilon_{n} \sqrt{\alpha K T} \right).\] In particular, for the linear case, when exactly \(m\) surrogate rewards are observed each round, this becomes \[\widetilde{O}\!\left( \sqrt{d\frac{K}{m+1}T\log\left(\frac{T}{d}\right)} + \varepsilon_{n} \sqrt{\frac{K^2}{m+1}T} \right).\] demonstrating a significant improvement by scale of roughly \(1/m\) as long as the noise from ML-predicted rewards \(\varepsilon_n\) is well-bounded.

3. Correlation-Aware Bandits with Surrogates — Decoupled (CABS-D) Finally, in Subsection 13.3 we establish the regret bound for the proposed prediction-mixing algorithm, a meta-bandit algorithm with two arms ((1) and (2) above) with regret bounded by min of them, \[\widetilde{O}\!\left( \min \left\{\sqrt{K T \mathrm{Reg}_{\mathrm{Sq}}(T)}, \sqrt{\alpha T\, \mathrm{Reg}_{\mathrm{Sq}}(T)} + \varepsilon_{n} \sqrt{\alpha K T} \right\}\right)\] For the linear case, when exactly \(m\) surrogate rewards are observed each round, this becomes \[\widetilde{O}\!\left( \min \left\{\sqrt{dKT \log \left(\frac{T}{d}\right) }, \sqrt{d\frac{K}{m+1}T\log \left(\frac{T}{d}\right)}+ \varepsilon_{n}\sqrt{\frac{K^2}{m+1}T}\right\}\right)\]

13.1 Contextual bandits Regret Bounds↩︎

We begin with the SquareCB [13] algorithm, which reduces contextual bandits to online regression. We state the modified algorithm that works on rewards instead of losses in Algorithm 8. We provide its regret guarantee for completeness and refer the reader to [13] for the proof.

13.2 Correlation Aware Bandits with Surrogates Coupled (CABS-C)↩︎

In this section, we present algorithmic analysis for Contextual Bandits with Noisy Graph Feedback with regret bound depending on a noise factor, \(\widetilde{O}\!\left(\sqrt{\alpha T\, \mathrm{Reg}_{\mathrm{Sq}}(T)}+\varepsilon_{n} \sqrt{\alpha K T}\right)\). In our main paper, we present this as Correlation Aware Bandits with Surrogates Coupled (CABS-C), Algorithm [alg:CABS-C]. We include here all intermediate lemmas and derivations to ensure completeness and reproducibility of the analysis.

13.2.1 High Level Proof Sketch↩︎

The proof of Theorem 1 proceeds in three main steps. First, we decompose the total expected regret into an algorithmic optimization term (sometimes referred to as the non-bias term), the cumulative regret of the regression oracle \(\mathrm{Reg_{Sq}}(T)\), and an expected bias term \(\mathbb{E}[\mathcal{B}_T]\) that accounts for the misspecification of the surrogate rewards. Next, we further decompose this expected bias term into two distinct components: a regression-propagation error and an observation noise term. We continue by systematically bounding each piece of this decomposition: the optimization term is bounded using a minimax argument over the action probabilities, the regression-propagation error is bounded using the Cauchy-Schwarz inequality and properties of harmonic sums, and the observation noise is bounded using the Azuma-Hoeffding inequality applied to the bounded martingale difference sequence. Combining these individual bounds and tuning the learning rate parameter \(\gamma=\sqrt{\alpha T/Reg_{Sq}(T)}\) to balance the components yields the final stated regret bound of \(\tilde{O}(\sqrt{\alpha T~Reg_{Sq}(T)}+\epsilon_n\sqrt{\alpha KT})\).

Regret Bound for Algorithm [alg:CABS-C] ()

13.2.1.1 Notation

Let \(n_{t,j}\) be the number of times arm \(j\) has been observed up to (and including) round \(t\), and \(\mathcal{T}_t(j)=\{\tau\le t: \mathcal{A}_{\tau,(i_\tau, j)}=1\}\). Finally denote \(\Delta_{\tau,j}:= |\boldsymbol{x}_\tau^{\top} \hat{\boldsymbol{\theta}}_{\tau,k} - \boldsymbol{x}_\tau^{\top}\boldsymbol{\theta}^*_{k} |\).

To prove Theorem 1, we first introduce the following auxiliary lemmas regarding the regret and bias decomposition. We continue by introducing lemmas to individually bound each term of the decomposition. Finally, we combine the results of all these lemmas in the final proof of Theorem 1.

Before introducing the lemmas to bound the noise terms of our regret decomposition, we introduce the following useful graph-theoretic result, which follows from  [35].

13.2.1.2 Proof of Theorem 1.

We now utilize the preceding lemmas to prove that under the assumptions outlined in Section 4.1, if the feedback graph \(G_t\) is deterministic with independence number no more than \(\alpha\). Then, with probability \(1-\delta\), Algorithm [alg:CABS-C] with choice \(\gamma = \max\left\{4, \sqrt{\alpha T / \mathrm{Reg_{Sq}}(T)}\right\}\) guarantees that \[\mathbb{E}[\mathrm{Reg}_T] \le \widetilde{O}\!\left( \sqrt{\alpha T\,\mathrm{Reg}_{\mathrm{Sq}}(T)} + \varepsilon_n\,\sqrt{\alpha K T}\; \right).\]

13.2.2 Bounds for \(m\)-Transitive Feedback Graphs↩︎

For completeness, we now state a theorem relating the independence number \(\alpha\) of the feedback graph \(G_t\) with a constant number of additional arms \(m\) observed at each round. This result enables us to tighten existing regret bounds for both standard and contextual bandits by a factor of roughly \(1/\sqrt{m}\), quantifying the direct benefit of the auxiliary feedback.

13.2.2.1 Deriving Corollary 1.

Corollary 1 follows by substituting \(\alpha=\frac{K}{m+1}\) from Theorem 5 into the general regret bound of Theorem 1, and applying the standard square-loss regret bound for linear regression oracles [13], \(\mathrm{Reg}_{\mathrm{Sq}}(T) = O(d\log \left(\frac{T}{d}\right))\).

13.3 Correlation Aware Bandits with Surrogates Decoupled (CABS-D)↩︎

We now present a meta bandit algorithm that achieves the minimum regret of the above two sections, i.e, \(\widetilde{O}\Big( \min \{\sqrt{dKT \log \left(\frac{T}{d}\right)}, \sqrt{d\frac{K}{m+1}T\log \left(\frac{T}{d}\right)}+ \varepsilon_n\sqrt{\frac{K^2}{m+1}T}\Big)\). In our main section, we present this as CABS-D (Algorithm 2). We start by presenting the Algorithm 2 for Correlation-Aware Bandits with Surrogates Decoupled (CABS-D) and then prove the regret bound of it.

13.3.1 High Level Proof Sketch↩︎

The proof of Theorem 2 establishes a best-of-both-worlds guarantee by analyzing CABS-D as an aggregation of two base experts: a bandit-feedback expert, SquareCB (Algorithm 8) and a graph-feedback expert, CABS-C (Algorithm [alg:CABS-C]), each instantiated over a geometric grid of learning rates to form \(M = 2L\) meta-copies. The analysis proceeds by first establishing the optimistic bias, uniform boundedness, and second-moment properties of the Implicit eXploration (IX) loss estimators utilized by both experts. Next, we decompose the total expected regret into the regret of the best-performing base expert (\(R_m\)), an additive bias term, and a meta-learning overhead term. By bounding the second-order variance terms independently for the regimes where either the bandit or graph expert dominates, and demonstrating that the geometric grid contains an optimal scaling parameter \(\gamma^{(g)}\) for each respective regime, we show that the aggregation overhead is tightly controlled. Consequently, the total expected regret asymptotically matches the minimum regret of either the pure bandit or correlation-aware surrogate expert, yielding \(\tilde{O}\left(\min\{\sqrt{KT \text{Reg}_{Sq}(T)}, \sqrt{\alpha T \text{Reg}_{Sq}(T)} + \epsilon_n\sqrt{\alpha KT}\}\right)\).
Regret Bound for Algorithm 2 ()
Overview and setup. We follow the same algorithmic structure as in Algorithm 2 (geometric gridded \(\gamma\) per base expert). For clarity we restate the essential components and set the notation used throughout the proof. Note that we utilize losses in the following analysis. There are 2 base experts (bandit-style (\(m=1\)) and graph-style (\(m=2\))). Each base expert \(m\in\{1,2\}\) is instantiated with a geometric grid of \(L\) parameter values \(\Gamma=\{\gamma^{(1)},\dots,\gamma^{(L)}\}\), with dyadic spacing \(\gamma^{(g+1)}=2\gamma^{(g)}\) (or any constant ratio \(>1\)). Each pair \((m,g)\) is a distinct meta-copy. The total number of meta-copies is \(M \;=\; 2L\).

To prove Theorem 2, we first introduce the following auxiliary lemma regarding the properties of the Implicit eXploration (IX) estimators. We continue by introducing auxiliary lemmas regarding the regret decomposition and bounding second-order terms. We combine and utilize the results of these lemmas in the final proof of Theorem 2.

13.3.1.1 Proof of Theorem 2.

We now use the preceding lemmas to prove that if there are \(M=2\) experts: one bandit-feedback expert and one graph-feedback expert (tuned for known independence number \(\alpha\)), then Algorithm 2 satisfies \[\begin{align} &\mathbb{E}\left[\mathrm{Reg_{Sq}}(T)\right] \\ &\le O\!\left(\min\left\{R_1 + \ln\!\left(KT\ln T\right)\sqrt{KT\ln M \ln T}, \right.\right. \left.\left. R_2 + \ln\!\left( \alpha T \ln(T/\alpha\epsilon)\ln T\right)\sqrt{\alpha T \ln(T/\alpha\epsilon) \ln M \ln T}\right\}\right) \end{align}\]

Using \(R_1 = \widetilde{O}(\sqrt{KT \, \mathrm{Reg_{Sq}}(T)})\) from Theorem 4 and \(R_2 = \widetilde{O}(\sqrt{\alpha T \, \mathrm{Reg_{Sq}}(T)} + \sqrt{\alpha KT}\,\varepsilon_n)\) from Theorem 1, we obtain the regret bound \[\widetilde{O}\!\left( \min \{\sqrt{KT \, \mathrm{Reg_{Sq}}(T)}, \sqrt{\alpha T \, \mathrm{Reg_{Sq}}(T)} + \sqrt{\alpha KT}\,\varepsilon_n\}\right).\]

13.3.1.2 Deriving Corollary 2.

Corollary 2 follows by substituting \(\alpha=\frac{K}{m+1}\) from Theorem 5 into the general regret bound of Theorem 2, and applying the standard square-loss regret bound for linear regression oracles [13], \(\mathrm{Reg}_{\mathrm{Sq}}(T) = O(d\log \left(\frac{T}{d}\right))\).

References↩︎