Learning to Route LLMs from Bandit Feedback: One Policy, Many Trade-offs

TL;DR

This work tackles the problem of selecting the most suitable LLM for each prompt while balancing accuracy and cost, addressing the gap between training-time full supervision and deployment-time bandit feedback. BaRP reframes routing as a preference-conditioned contextual bandit trained with REINFORCE, using a frozen prompt encoder and a small learnable preference encoder to map prompts and user trade-offs to model selections. The method eliminates the need for full-information labels and enables test-time control over performance-cost preferences, demonstrated by superior in-distribution and out-of-distribution performance and favorable cost-efficiency compared with strong offline routers and large models. These results suggest that partial-feedback training combined with per-request preference control provides a practical, scalable approach for deploying multi-LLM routing in real-world systems. BaRP thus offers a principled, tunable, and resource-conscious solution for scalable LLM routing with meaningful implications for deployment efficiency and user experience.

Abstract

Efficient use of large language models (LLMs) is critical for deployment at scale: without adaptive routing, systems either overpay for strong models or risk poor performance from weaker ones. Selecting the right LLM for each query is fundamentally an online decision problem: models differ in strengths, prices fluctuate, and users value accuracy and cost differently. Yet most routers are trained offline with labels for all candidate models, an assumption that breaks in deployment, where only the outcome of the chosen model is observed. We bridge this gap with BaRP, a Bandit-feedback Routing with Preferences approach that trains under the same partial-feedback restriction as deployment, while supporting preference-tunable inference: operators can dial the performance/cost trade-off at test time without retraining. Framed as a contextual bandit over prompt features and a user preference vector, our method simulates an online feedback setting during training and adapts its routing decisions to each new prompt, rather than depending on full-information offline supervision. Comprehensive experiments show that our method consistently outperforms strong offline routers by at least 12.46% and the largest LLM by at least 2.45%, and generalizes robustly for unseen tasks.

Figures (3)

• Figure 1: Testing score of baselines and BaRP on in-distribution and out-of-distribution tasks.
• Figure 2: The training pipeline of BaRP. The router takes the context (query $x_t$ and preference $w_t$) and selects an LLM. It then receives bandit feedback (the score and cost of the chosen LLM only) to calculate a reward $r_t$. This reward drives a learning algorithm to update the router's parameters, including policy gradient methods like REINFORCE (Sec. \ref{['sec:learning']}) and classic bandit algorithms such as LinUCB, Thompson Sampling, and $\epsilon$-greedy (Sec. \ref{['bandit_algo']}).
• Figure 3: Effects of $\omega^c$