arXiv:cs.LG· Ajay Narayanan Sridhar, Ronak Singh, Mehrdad Mahdavi, Vijaykrishnan Narayanan·· 2 小时前
面向 LLM 路由的相关性感知上下文老虎机与代理奖励
Correlation-Aware Contextual Bandits with Surrogate Rewards for LLM Routing
AI 导读
研究针对 LLM 路由等场景中臂间存在相关性、且可获取机器学习模型产生的代理奖励信号的上下文老虎机问题,提出耦合奖励混合与解耦预测混合两类算法。解耦设计在代理信号失配时仍能保持与纯老虎机方法相当的 regret 界,代理预测足够有效时 regret 更优。在 LLM 路由基准上,该方法相比标准上下文老虎机基线和强静态路由方法展现出更好的样本效率与准确率-成本权衡。
正文
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.
| Subjects: | Machine Learning (cs.LG); Artificial Intelligence (cs.AI) |
| Cite as: | arXiv:2607.09015 [cs.LG] |
| (or arXiv:2607.09015v2 [cs.LG] for this version) | |
| https://doi.org/10.48550/arXiv.2607.09015 arXiv-issued DOI via DataCite |
Submission history
From: Ajay Narayanan Sridhar [view email]
[v1]
Fri, 10 Jul 2026 00:42:10 UTC (3,194 KB)
[v2]
Tue, 13 Oct 2026 20:50:09 UTC (4,388 KB)
来源:arXiv:cs.LG · arxiv.org