Correlation-Aware Contextual Bandits with Surrogate Rewards for LLM Routing

2026-07-09

AI Summary

This work addresses the limitations of traditional contextual bandits in large language model (LLM) routing, which typically ignore action correlations and rely solely on sparse feedback. To overcome these issues, the paper introduces a novel framework that jointly models inter-action dependencies and noisy proxy rewards within the contextual bandit setting. Two new algorithms are proposed: a coupled reward-mixing approach to accelerate learning and a decoupled prediction-mixing method to enhance robustness against proxy signal mismatch. An adaptive combination mechanism seamlessly integrates the strengths of both strategies. Experimental results demonstrate that the proposed method significantly improves sample efficiency on standard LLM routing benchmarks and achieves a superior trade-off between accuracy and inference cost compared to conventional bandit algorithms and static routing policies.

Abstract

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. Unlike classical contextual bandits that rely solely on bandit feedback and assume conditional independence across arms, our setting allows context-dependent inter-arm correlations and auxiliary reward information that may be noisy or misspecified. We propose algorithms that leverage such surrogate rewards through two complementary designs. A coupled reward-mixing approach pools true and surrogate rewards to accelerate learning when surrogate signals are reliable, while a decoupled prediction-mixing approach maintains separate estimators for bandit feedback and surrogate rewards and adaptively combines their predictions. This decoupling yields robustness to surrogate misspecification, recovering regret guarantees comparable to reward-only bandit methods in the worst case, while achieving improved regret when surrogate predictions are sufficiently informative. We provide theoretical regret analyses for both approaches and evaluate them on LLM routing benchmarks under varying accuracy versus cost trade-offs. The results demonstrate improved sample efficiency and consistently better accuracy-cost trade-offs compared to standard contextual bandit baselines and strong static routing methods.

Problem

• contextual bandits
• correlated arms
• surrogate rewards
• LLM routing
• reward misspecification

Innovation

• Contextual Bandits
• Surrogate Rewards
• Arm Correlation
• LLM Routing
• Robust Learning