Reward-Based Online LLM Routing via NeuralUCB

1. Problem Definition

The paper addresses the challenge of cost-aware Large Language Model (LLM) routing. In real-world deployments, LLMs vary significantly in both capability (quality) and inference cost. High-end models offer superior performance but are expensive, while cheaper models may fail on complex tasks.

The goal is to design a dynamic routing policy that selects the most appropriate model ($a_t$) for a given user query ($x_t$) to maximize long-term utility. This is framed as a Contextual Bandit problem with the following characteristics:

• Partial Feedback: The system only observes the reward (quality and cost) of the selected model, not the counterfactual outcomes of unselected models.
• Non-linearity: The mapping from query features to the optimal model is highly non-linear and difficult to model with simple linear assumptions.
• Objective: Maximize the expected cumulative utility, defined as a tradeoff between answer quality and inference cost, rather than maximizing quality alone.

2. Methodology

The proposed solution, NeuralUCB-based Routing, consists of three core components: a utility function, a neural predictor (UtilityNet), and an exploration-exploitation policy.

A. Utility Function

To balance quality and cost, the authors define a utility reward $r(x, a)$:
$$r(x, a) = q(x, a) \cdot \exp(-\lambda \tilde{c}(x, a))$$
Where:

• $q(x, a)$ is the model quality.
• $\tilde{c}(x, a)$ is the log-normalized cost, mapping raw costs to a $[0, 1]$ range to stabilize training.
• $\lambda$ is a hyperparameter controlling the penalty strength for cost.

B. UtilityNet Architecture

UtilityNet is a neural network designed to predict the utility reward and provide feature representations for uncertainty estimation. It processes:

1. Inputs: Text embeddings ($x_{emb}$), auxiliary features ($x_{feat}$), and domain IDs ($d$).
2. Encoders: Separate MLPs for text, features, and domain embeddings.
3. Fusion: The text, feature, and action embeddings are concatenated to form a joint representation.
4. Dual Heads:
   • Utility Head: Regresses the predicted utility $\mu(x, a)$ using Huber loss.
   • Gating Head: Predicts a probability $p(x)$ indicating whether exploration (UCB bonus) is beneficial for the current context. This allows the system to switch between aggressive exploration and conservative exploitation.

C. NeuralUCB Policy

The routing policy uses the Neural Upper Confidence Bound (NeuralUCB) algorithm for action selection:

• Feature Representation: The last hidden layer of UtilityNet, $h(x, a)$, serves as the feature vector for the UCB calculation.
• UCB Score: $s(x, a) = \mu(x, a) + \beta \sqrt{g(x, a)^\top A^{-1} g(x, a)}$, where $A^{-1}$ is a shared inverse covariance matrix estimating uncertainty in the feature space, and $\beta$ controls the exploration bonus.
• Gated Selection Rule:
  • If the gating probability $p(x) \geq \tau_g$ (threshold), the system selects the action with the highest UCB score (exploration enabled).
  • Otherwise, it falls back to a safe action ($a_{safe}$) that maximizes the predicted mean utility $\mu(x, a)$ (exploitation).

D. Training Protocol

The model is evaluated under a simulated online setting using the RouterBench dataset:

• Data is split into 20 sequential slices.
• For each slice, the policy selects actions, updates a replay buffer and the shared covariance matrix $A^{-1}$, and then retrains UtilityNet.
• This mimics an online learning environment without the prohibitive cost of live API calls.

3. Key Contributions

1. NeuralUCB for LLM Routing: The first application of NeuralUCB to the specific problem of cost-aware LLM routing, effectively handling non-linear query-to-model mappings under partial feedback.
2. Gated Exploration Mechanism: Introduction of a gating branch that dynamically decides whether to apply the UCB exploration bonus, preventing unnecessary exploration on queries where the optimal model is already clear.
3. Comprehensive Evaluation Framework: Implementation of a robust offline replay pipeline on RouterBench that simulates online learning dynamics, allowing for the tracking of cumulative rewards and cost-quality tradeoffs over time.

4. Experimental Results

The method was evaluated on RouterBench (36,497 samples, 86 domains, 11 candidate LLMs) against baselines including Random, Min-Cost, and RouteLLM-BERT.

• Performance vs. Baselines:
  • Reward: NeuralUCB consistently achieved the highest average reward (approx. 0.59–0.61), significantly outperforming RouteLLM-BERT (0.35), Min-Cost (0.52), and Random (~0.32).
  • Cumulative Reward: The gap between NeuralUCB and baselines widened over time, demonstrating effective learning and adaptation.
• Cost-Quality Tradeoff:
  • The NeuralUCB policy achieved competitive quality while using only ~33% of the inference cost of the "max-quality" reference (which always picks the best model regardless of cost).
• Ablation Studies:
  • Text Encoders: Smaller encoders like all-MiniLM-L6-v2 and all-mpnet-base-v2 performed as well as or better than larger, more complex encoders (e.g., multilingual-e5-large-instruct), suggesting diminishing returns for encoder size in this specific routing context.

5. Significance and Future Work

• Practical Impact: The study demonstrates that it is possible to dynamically route LLM queries to achieve a superior balance between cost and performance, potentially saving significant computational resources in production systems.
• Limitations: The authors note that for many samples, the "max-quality" model is already the reward-optimal choice, limiting the potential gains from routing.
• Future Directions: The paper suggests focusing on collecting more diverse training data and refining the utility model to improve action discrimination, particularly in scenarios where cost constraints are tight and the optimal model is not obvious.

In conclusion, this paper presents a robust, theoretically grounded approach to LLM routing that leverages neural contextual bandits to navigate the complex landscape of model capabilities and costs, offering a practical path toward more efficient LLM deployment.