Cost-Aware LLM Routing with NeuralUCB: Balancing Quality and Efficiency

Paper: 2603.30035 Authors: Ming-Hua Tsai, Phat Tran Categories: cs.LG, cs.CL

The Gap

Existing LLM routing methods fall into two camps. Supervised routing trains a classifier on labeled data to predict which model to use, but it’s static—can’t adapt when model performance shifts or new models arrive. Partial-feedback methods like reinforcement learning can adapt online, but they’re sample-inefficient and struggle with the exploration-exploitation tradeoff when you have many models to choose from.

The core tension: you want to route queries to cheaper models when possible (save money), but you need to know which queries each model can handle well (requires exploration). Do it wrong and you either waste money on expensive models or tank quality with cheap ones.

Problem: Many LLMs, different costs/quality
         Need to route queries efficiently
              |
              v
    Supervised methods     Partial-feedback methods
    (static, no adapt) <-> (adaptive, sample-hungry)
              |
              v
    Gap: Need online adaptation + sample efficiency
              |
              v
    Method: NeuralUCB (contextual bandit)
              |
              v
    Evidence: RouterBench experiments
              |
              v
    Conclusion: Competitive quality at lower cost

The Increment

One sentence: Before this paper, LLM routing was either static-but-efficient or adaptive-but-wasteful; after, we have a contextual bandit approach that adapts online while maintaining sample efficiency through neural uncertainty estimation.

Core Mechanism

NeuralUCB treats LLM routing as a contextual bandit problem. Each query is a context, each LLM is an arm, and the reward combines quality (how good the answer is) and cost (how much inference costs). The algorithm maintains a neural network that predicts expected reward for each query-model pair, plus an uncertainty estimate. At decision time, it picks the model with highest upper confidence bound: predicted reward plus uncertainty bonus.

The neural network gets trained on accumulated feedback. When you route a query to a model and observe the outcome, that becomes a training example. The uncertainty comes from an ensemble of networks or a Bayesian approximation—regions of input space with sparse data get high uncertainty, encouraging exploration there.

The key insight: use the neural network’s uncertainty to guide exploration. If you’re confident a cheap model will work (low uncertainty, high predicted reward), use it. If you’re unsure (high uncertainty), the UCB bonus makes you try it to gather information. If you’re confident it won’t work, avoid it.

Query arrives
     |
     v
[Neural Net] --predicts--> Reward estimates for each LLM
     |                            |
     +--estimates--> Uncertainty  |
                          |       |
                          v       v
                    UCB = reward + bonus
                          |
                          v
                    Pick max UCB
                          |
                          v
                    Route to LLM
                          |
                          v
                    Observe outcome
                          |
                          v
                    Update neural net

Think of it like a restaurant manager assigning dishes to cooks of different skill levels and wages. You have a star chef (expensive, reliable) and several line cooks (cheap, variable). Each order comes with context—complexity, cuisine type, time pressure. The manager maintains mental models of each cook’s strengths and uncertainties. For a simple burger, confidently assign the cheapest cook. For a complex soufflé with uncertain difficulty, maybe try a mid-tier cook to learn their limits—the uncertainty bonus justifies the risk. Over time, the manager’s mental model sharpens, and assignments get more efficient. The neural network is that mental model, UCB is the decision rule that balances “use what works” with “learn the boundaries.”

Key Concepts

• Contextual Bandit: You face a sequence of decisions. Each time, you see a context (query features), choose an action (which LLM), and get a reward (quality minus cost). The goal is to maximize cumulative reward. Unlike supervised learning, you only see the reward for the action you took, not the counterfactual “what if I’d chosen differently.” Unlike full reinforcement learning, your action doesn’t change the next context—each query is independent. The challenge is balancing exploitation (pick the action you think is best) with exploration (try uncertain actions to improve your model). Imagine a doctor choosing treatments: each patient is a context, each treatment is an action, and you only learn the outcome of the treatment you prescribed, not the alternatives.

• Upper Confidence Bound (UCB): A principled way to handle exploration-exploitation. For each action, compute an upper confidence bound: your best guess of the reward plus a bonus proportional to your uncertainty. Pick the action with highest UCB. If you’re certain an action is bad, the bonus is small and it won’t get picked. If you’re uncertain, the bonus is large—maybe it’s actually great and you should find out. Over time, as you gather data, uncertainties shrink and the algorithm converges to exploiting the best actions. It’s like bidding in an auction where you add a “information value” premium to items you know little about.

• Neural Uncertainty Estimation: Standard neural networks give point predictions but no confidence intervals. To get uncertainty, you can train an ensemble (multiple networks with different initializations—their disagreement indicates uncertainty) or use Bayesian approximations (treat network weights as distributions, not fixed values). In regions with lots of training data, predictions converge and uncertainty is low. In sparse regions, predictions vary and uncertainty is high. This lets the UCB algorithm know where it’s ignorant and needs to explore. Think of weather forecasting: a single model gives a temperature prediction, but an ensemble of models gives a range—if they all agree, you’re confident; if they diverge, you’re uncertain.

Framework Shift

Before (mainstream approach):        After (this paper):

Train classifier on labeled data     Start with no labels
         |                                   |
         v                                   v
  Query -> Classifier -> LLM          Query -> NeuralUCB -> LLM
         (fixed policy)                      |
                                             v
                                      Observe reward
                                             |
                                             v
                                      Update neural net
                                      (adaptive policy)

Static: Can't adapt to new models    Dynamic: Learns online
Supervised: Needs labels upfront     Bandit: Learns from rewards

From offline training to online adaptation, the core shift is treating routing as sequential decision-making under uncertainty rather than a one-time classification problem.

Expert Assessment

Problem choice: Real gap. LLM inference costs are skyrocketing, and the zoo of models (GPT-4, Claude, Llama, Gemini, specialized fine-tunes) keeps growing. Routing is a practical bottleneck. The paper correctly identifies that existing methods are either too rigid (supervised) or too sample-hungry (naive RL). This sits at the intersection of practical ML systems and bandit theory—timely and grounded.

Method maturity: Solid application of existing theory, not a novel algorithm. NeuralUCB has been around since 2020 (Zhou et al.). The contribution here is recognizing that LLM routing fits the contextual bandit framing and implementing it on RouterBench. The execution is competent but not groundbreaking—no new theoretical guarantees, no architectural innovations. A simpler linear UCB might have been a useful ablation to show where the neural component adds value.

Experimental integrity: Baselines are reasonable (random, min-cost, max-quality) but not exhaustive. Missing comparisons to recent supervised routers (RouteLLM, FrugalGPT) and other bandit algorithms (Thompson Sampling, LinUCB). The “simulated online setting” is a bit hand-wavy—how exactly are queries sequenced? Are there distribution shifts? The paper shows NeuralUCB beats random and min-cost, which is a low bar. The gap to max-quality is large, suggesting the method still struggles with hard queries. More error analysis would help: which query types does it route poorly?

Writing quality: The abstract and intro are clear, but the method section is thin—barely two pages. The reader has to fill in gaps: what neural architecture? How is the ensemble trained? What’s the exact UCB formula? The experiments section jumps straight to results without describing the setup in detail. Rewriting the method section with pseudocode and architectural diagrams would elevate the paper from “application note” to “reproducible contribution.”

Verdict: Weak accept—the problem is real, the approach is sensible, and the results are promising, but the execution feels rushed and the evaluation is incomplete. This is a good workshop paper or a starting point for a more thorough study, not a landmark contribution.

Takeaways

Practitioners can steal the framing: if you’re routing between multiple models (LLMs, vision models, speech recognizers), treat it as a contextual bandit, not a classification problem. This lets you adapt online without labeled data. The UCB principle is simple to implement: maintain a model of expected reward and uncertainty, then pick the action with highest reward + uncertainty bonus. You can start with a linear model (cheap, interpretable) and upgrade to neural if you have complex context features. The key insight transfers beyond LLMs: anytime you’re choosing between options with different cost-quality tradeoffs and you can observe outcomes, contextual bandits are a natural fit.