We investigate contextual bandit problems with correlated arms and utilize surrogate reward signals produced by a machine learning model, particularly motivated by applications like large language model (LLM) routing. Unlike classical contextual bandits that rely solely on bandit feedback and assume conditional independence across arms, our setting allows for context-dependent inter-arm correlations and auxiliary reward information that may be noisy or misspecified.
We propose two complementary algorithmic designs to leverage these surrogate rewards. The first, a coupled reward-mixing approach, pools true and surrogate rewards to accelerate learning when surrogate signals are reliable. The second, a decoupled prediction-mixing approach, maintains separate estimators for bandit feedback and surrogate rewards, adaptively combining their predictions. This decoupling enhances 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. Results demonstrate improved sample efficiency and consistently better accuracy-cost trade-offs compared to standard contextual bandit baselines and strong static routing methods.
Blogger's Review: This study offers a fresh perspective on contextual bandit problems by incorporating surrogate reward signals, significantly enhancing routing efficiency for large language models. The coupled and decoupled strategy designs effectively balance learning speed and robustness, providing strong theoretical support and empirical results for practical applications, warranting further exploration in related fields.