Research map / Bandits and online learning / Contextual bandits and Thompson sampling

Correlation-Aware Contextual Bandits with Surrogate Rewards for LLM Routing

NeurIPS 2026 · Bandits and online learning › Contextual bandits and Thompson sampling · approach: Oracle-efficient contextual bandits · 0 citations (OpenAlex)

Read the paperCheck your idea against 108,967 papers

AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.

Proposes CABS-C (coupled reward mixing) and CABS-D (decoupled prediction mixing) contextual bandits that exploit context-dependent arm correlations and ML surrogate rewards for LLM routing, with regret bounds and experiments on RouterBench, SPROUT, and Open LLM LB v2.

Problem

Online LLM routing is usually modeled as a contextual bandit with conditionally independent arms and only bandit feedback, which is sample-inefficient; surrogate reward predictions from ML models could help but may be biased or noisy.

Approach

Use a context-dependent feedback graph known before action selection, and ML-predicted surrogate rewards for graph-neighbor arms; CABS-C pools de-biased surrogates with true rewards in one regression oracle, while CABS-D runs a bandit expert (SquareCB) and CABS-C as experts and aggregates them with an adaptive Hedge-style meta-learner for best-of-both-worlds regret.

Key results

• On SPROUT low-cost regime, CABS-D (linear) achieves utility 0.7715 vs 0.7579 for LinUCB and 0.7628 for LinTS.
• On RouterBench medium-cost regime, CABS-D (linear) reaches utility 0.5817 vs 0.5725 for LinTS (best online baseline).
• On Open LLM LB v2, CABS-D is reported to gain about 8% over standard online bandit baselines on average, with ~15% in the low-cost regime.

Surprising findings

• CABS-C (coupled) is consistently worse than CABS-D (decoupled), indicating that coupling surrogate and true feedback can hurt when surrogates are noisy.
• CABS-D with CABS-D surrogate-driven gains over online baselines are modest on SPROUT (~1%), and CABS-D underperforms the full-feedback oracle by a noticeable margin.

Limitations

• Surrogate reward quality depends on offline ML model generalization; theory assumes bounded noise epsilon_n and bias that the paper does not validate empirically on unseen contexts.
• Affinity graph uses a proxy correlation (1-2|y_i - y_j|) rather than true context-conditional correlations, and surrogate pseudo-observations are pruned via an IQR heuristic.
• Regret bounds are for linear oracles; neural extension is empirical only, and the K^3 cost limits scaling to large model pools.

Experimental setup

Datasets: RouterBench SPROUT Open LLM Leaderboard v2

Benchmarks: RouterBench SPROUT Open LLM Leaderboard v2

Metrics: average utility cumulative regret accuracy-cost Pareto frontier average cost

Baselines: LinUCB LinTS Epsilon-Greedy NeuralUCB CARROT RouterBench KNN Random Selection LinUCB + Full-Feedback NeuralUCB + Full-Feedback

Builds on: SquareCB LinUCB NeuralUCB Hedge/AdaHedge Corral

Closest papers by meaning

• Active Context Selection Improves Simple Regret in Contextual Bandits (NeurIPS 2026 · similarity 0.87)
• Mixed-Effects Contextual Bandits (AAAI 2024 · similarity 0.87)
• Causal LLM Routing: End-to-End Regret Minimization from Observational Data (NeurIPS 2025 · similarity 0.87)
• Exploration via Feature Perturbation in Contextual Bandits (NeurIPS 2025 · similarity 0.87)
• Local Anti-Concentration Class: Logarithmic Regret for Greedy Linear Contextual Bandit (NeurIPS 2024 · similarity 0.87)
• Feel-Good Thompson Sampling for Contextual Bandits: a Markov Chain Monte Carlo Showdown (NeurIPS 2025 · similarity 0.87)