Adaptive LLM Routing under Budget Constraints (2025)

Abstract: Large Language Models (LLMs) have revolutionized natural language processing, but their varying capabilities and costs pose challenges in practical applications. LLM routing addresses this by dynamically selecting the most suitable LLM for each query/task. Previous approaches treat this as a supervised learning problem, assuming complete knowledge of optimal query-LLM pairings. However, real-world scenarios lack such comprehensive mappings and face evolving user queries. We thus propose to study LLM routing as a contextual bandit problem, enabling adaptive decision-making using bandit feedback without requiring exhaustive inference across all LLMs for all queries (in contrast to supervised routing). To address this problem, we develop a shared embedding space for queries and LLMs, where query and LLM embeddings are aligned to reflect their affinity. This space is initially learned from offline human preference data and refined through online bandit feedback. We instantiate this idea through Preference-prior Informed Linucb fOr adaptive rouTing (PILOT), a novel extension of LinUCB. To handle diverse user budgets for model routing, we introduce an online cost policy modeled as a multi-choice knapsack problem, ensuring resource-efficient routing.

## introduction deploying large language models (llms) in realworld systems faces a critical challenge: balancing performance with cost-effectiveness @xcite . while larger models offer superior performance, their high costs makes their universal de-ployment impractical. this challenge is particularly acute given the varying pricing structures of proprietary models and the resource requirements of deploying the open-source alternatives. to understand the need for varying resource requirements, consider a customer service chatbot handling diverse queries. for simple queries like "what are your business hours?", a smaller, costeffective model might suffice. however, for complex inquiries, such as "i'm torn between two of your smartphone models: the x200 and the z300. i need a phone with excellent battery life, a highquality camera, and robust performance for multitasking. can you provide a detailed comparison, including any potential drawbacks of each model?", a more powerful (and costly) model may be necessary to ensure better planning and reasoning capabilities, which a smaller model might lack. this scenario illustrated the need for dynamic query routing -the ability to dynamically route queries to the most appropriate model based on complexity and cost considerations, known as the modelrouting problem @xcite . existing approaches model the routing problem as a supervised learning task, requiring large-scale labeled datasets mapping queries to their optimal llm pairings @xcite . this problem formulation faces two limitations: (1) gathering such labeled datasets is very expensive, as it requires responses from each model in the model pool for every query to discover the optimal query-llm pairing. (2) lack of adaptability to change in query distribution. thus, to find a more practical problem setting, we draw parallels with news recommendation systems, where models can only learn from user feedback, such as clicks on a single article @xcite . these systems must predict the best article for a user without showing all possible articles , receiving feedback solely on the selected article. similarly, llm routing in-volves choosing the best model for a query, where obtaining user feedback on all models is laborintensive and costly. building on this insight, we reformulate llm routing as a contextual bandit learning problem -a formulation commonly used in news and ad recommendation scenarios @xcite . thus, instead of requiring outputs from every llm to identify the best match, our problem setting relies only on a binary bandit feedback, i.e., whether the chosen llm's response is good or not. this approach is practical, as simple feedback mechanisms, like thumbs up/down ratings (i.e. like/dislike feedback), are now common in chat interfaces @xcite , allowing effective learning from user interactions without need for extensive annotation across llms. to address this newly formulated problem of llm routing with bandit feedback, we propose to develop an evolving shared embedding space for queries and llms, where distances represent routing affinity. initially pretrained on human preference data @xcite , it is refined through online user feedback. furthermore, to enforce cost constraints in an online setting, we introduce a novel policy modeled as an online multi-choice knapsack problem @xcite , dynamically allocating resources to balance budget adherence and performance. while authors in @xcite have approached the routing problem from a bandit perspective, their work was focused on the limited setting of classification tasks and did not incorporate explicit online cost constraints. our key contributions are as follows, (i) we formulate llm routing as a budget constrained contextual bandit problem for adaptive decision-making with limited supervision (ii) we propose a preference-prior informed lin-ucb algorithm (pilot) that combines offline human preference data with online bandit feedback to route queries (sec 2.2.2). we also show our preference prior helps achieve a lower regret bound than standard algorithm. this algorithm is further coupled with an online cost policy to dynamically allocate cost budget to queries (sec 2.2.3). (iii) in sec 4 we show that our method outperforms existing bandit baselines across datasets, achieving lowest regrets and highest performance across different cost budgets. ## results and analysis we evaluate pilot's performance across various facets of llm routing, considering diverse, multitask applications and specific use cases. our experiments utilize two data types: (i) single-task data source: the mmlu benchmark from routerbench, focusing on multi-choice question answering. (ii) multi-task data source: full routerbench dataset, encompassing tasks like code generation, math problems, & multi-turn conversations, simulating a broad range of user queries. using these datasets, we investigate several aspects of our bandit-based llm routing algorithm, detailed under main experiments in ## discussion and ablations this section provides qualitative and quantitative analyses of pilot's routing behavior, computational efficiency, cost policy, and sensitivity, offering a holistic view. these analyses, summarized under 'analysis' in ## conclusion we address llm routing with budget constraints using bandit feedback in this work. we propose pilot, a human-preference prior based contextual bandit algorithm, coupled with a novel online cost policy that optimizes budget allocation across queries. our approach achieves 93% of gpt-4's performance at 25% of its cost on routerbench.