Andrii Bidochko

• Updated: July 23, 2026
• 6 min read

Correlation-Aware Contextual Bandits with Surrogate Rewards for LLM Routing

Direct Answer

The paper introduces Correlation‑Aware Contextual Bandits with Surrogate Rewards, a framework that jointly exploits inter‑arm correlations and noisy auxiliary reward signals to route queries to the most suitable large language model (LLM). This matters because it dramatically reduces the number of real‑world interactions needed to learn effective routing policies, cutting cost while preserving or improving accuracy.

Background: Why This Problem Is Hard

Enterprises that expose multiple LLMs—each tuned for different domains, latency budgets, or pricing tiers—must decide, for every incoming request, which model to invoke. The decision problem is a classic contextual bandit: the system observes a request (the context) and selects an arm (an LLM) to receive a reward (e.g., user satisfaction, cost efficiency). In practice, three intertwined challenges make this problem especially difficult:

• Correlated arms: LLMs are not independent; they share training data, architecture families, or fine‑tuning objectives, which creates statistical dependencies that classic bandit theory ignores.
• Expensive feedback: True reward signals (e.g., human ratings) are costly to obtain at scale, leading to sparse and delayed feedback.
• Surrogate signals: Organizations often have cheap, model‑generated proxies—such as confidence scores, token‑level likelihoods, or downstream classifier outputs—that are noisy or even biased.

Existing contextual bandit algorithms assume conditional independence across arms and rely exclusively on bandit feedback. When arms are correlated, those methods waste information that could otherwise accelerate learning. Moreover, they cannot safely incorporate surrogate rewards because a misspecified proxy can mislead the learner, inflating regret.

What the Researchers Propose

To address these gaps, the authors propose a two‑pronged algorithmic family that treats surrogate rewards as a first‑class citizen while explicitly modeling arm correlations:

1. Coupled Reward‑Mixing (CRM)

CRM blends the observed bandit reward with the surrogate prediction in a single estimator. When the surrogate is trustworthy, the mixed signal provides a richer learning target, shrinking the exploration horizon.

2. Decoupled Prediction‑Mixing (DPM)

DPM maintains two parallel estimators—one for true bandit feedback, another for surrogate predictions. An adaptive gating mechanism then combines their outputs based on an online estimate of surrogate reliability. This design safeguards against surrogate misspecification, reverting to pure bandit learning when the proxy proves unhelpful.

Both approaches embed a correlation‑aware model (e.g., a Gaussian process or a low‑rank linear factorization) that captures how the performance of one LLM informs expectations about the others under the same context.

How It Works in Practice

The end‑to‑end workflow can be visualized as a pipeline of four logical components:

1. Context Encoder: Transforms the incoming request (text, metadata, user profile) into a dense feature vector.
2. Correlation Module: Uses a learned covariance structure to predict the expected reward of each arm given the encoded context.
3. Surrogate Processor: Queries auxiliary models (e.g., a lightweight confidence estimator) to produce a surrogate reward for every candidate LLM.
4. Decision Engine: Applies either CRM or DPM to merge the bandit and surrogate signals, then selects the arm with the highest mixed score.

After the chosen LLM generates its response, the system may collect a delayed true reward (e.g., a click‑through or human rating). This feedback updates the bandit estimator, while the surrogate processor is continuously retrained on the growing pool of labeled data.

The key differentiator is the adaptive mixing step. In DPM, a confidence weight α_t is computed on the fly based on recent prediction errors of the surrogate. If α_t is low, the decision engine leans heavily on the pure bandit estimate, effectively neutralizing a noisy proxy. Conversely, a high α_t lets the surrogate dominate, accelerating convergence when the proxy is accurate.

This architecture can be deployed as a microservice within any LLM orchestration platform, requiring only standard REST endpoints for context encoding and reward reporting.

Evaluation & Results

The authors benchmarked both CRM and DPM on three realistic LLM routing datasets:

• Multi‑Domain QA: Queries span legal, medical, and technical domains, each best served by a specialized LLM.
• Cost‑Sensitive Chat: A mix of high‑performance (expensive) and low‑latency (cheap) models, with a surrogate cost estimator.
• Dynamic Availability: Simulated outages force the system to re‑route in real time, testing robustness to changing arm sets.

Key findings include:

• Sample Efficiency: DPM achieved a 35 % reduction in cumulative regret compared to a vanilla LinUCB baseline, reaching comparable performance with roughly half the number of true reward observations.
• Robustness to Misspecification: When surrogate signals were deliberately corrupted (Gaussian noise added), DPM’s adaptive gating limited regret inflation to under 5 %, whereas CRM’s regret grew by more than 20 %.
• Cost‑Accuracy Trade‑off: In the cost‑sensitive scenario, both algorithms delivered higher accuracy per dollar spent, with CRM slightly edging out DPM when the surrogate was highly reliable.

These results demonstrate that the proposed methods not only accelerate learning but also gracefully degrade when auxiliary information is unreliable—a critical property for production systems that cannot guarantee perfect surrogate models.

Why This Matters for AI Systems and Agents

From a systems‑engineering perspective, the paper offers a concrete recipe for building self‑optimizing LLM routers that can be integrated into existing AI stacks. The practical implications are threefold:

1. Reduced Operational Costs: By leveraging cheap surrogate signals, organizations can cut the number of expensive human‑in‑the‑loop evaluations required to fine‑tune routing policies.
2. Improved User Experience: Faster convergence means users see higher‑quality responses sooner, especially in domains where model selection dramatically impacts answer correctness.
3. Scalable Orchestration: Correlation‑aware learning naturally extends to heterogeneous fleets of models, enabling seamless addition or removal of arms without retraining from scratch.

For AI agents that must call external LLM services—such as autonomous assistants, customer‑support bots, or data‑analysis pipelines—embedding a DPM‑style router can become a “decision‑making layer” that balances latency, cost, and accuracy in real time. Teams building such agents can prototype the approach using the UBOS platform overview, which already supports plug‑in reward models and custom correlation matrices.

What Comes Next

While the study makes a strong case for correlation‑aware bandits, several open challenges remain:

• Non‑Stationary Correlations: Real‑world model performance can drift as data distributions shift. Future work should explore online covariance adaptation or meta‑learning techniques.
• Multi‑Objective Rewards: Many deployments care about a vector of metrics (e.g., latency, privacy, compliance). Extending the framework to Pareto‑optimal bandits is an exciting direction.
• Scalable Surrogate Training: Training high‑fidelity surrogate models at scale may itself become a bottleneck; lightweight distillation or few‑shot learning could mitigate this.

Addressing these gaps will broaden the applicability of the approach to edge deployments, federated settings, and large‑scale enterprise AI platforms. Companies interested in experimenting with correlation‑aware routing can start by integrating the methodology into their existing workflow automation pipelines, such as the Workflow automation studio, and then iterate on surrogate design using the Chroma DB integration for efficient similarity search.

For a deeper dive into the theoretical foundations and to reproduce the experiments, consult the original arXiv paper. The authors also release code and data under an open‑source license, making it straightforward for researchers and engineers to benchmark against their own routing scenarios.

Andrii Bidochko

CTO UBOS

Andrii Bidochko is an AI entrepreneur and researcher focused on AI agents, reinforcement learning, and autonomous systems. He writes about the technologies shaping the future of machine intelligence, from frontier models and agent architectures to real-world AI applications.