PaperlayerView exact revision on arXiv

Source-linked AI summary

Universal Model Routing for Efficient LLM Inference

Wittawat Jitkrittum, Harikrishna Narasimhan, Ankit Singh Rawat, Jeevesh Juneja, Congchao Wang, Zifeng Wang, Alec Go, Chen-Yu Lee, Pradeep Shenoy, Rina Panigrahy, Aditya Krishna Menon, Sanjiv Kumar

arXiv:2502.08773v2cs.CLcs.LG

TL;DR

The paper studies how to route prompts when new, previously unobserved LLMs enter a changing candidate pool, without repeatedly retraining the router. UniRoute represents each LLM using prediction-error features on representative prompts and learns routing over these representations. Across public benchmarks, it effectively routes among more than 30 unseen LLMs, while relying on a modest validation set that must support efficient evaluation of new models.

  • Problem

    Existing model-routing methods largely assume a fixed LLM pool, whereas changing pools make retraining costly and require labels for each new LLM.

  • Method

    UniRoute represents prompts and LLMs as feature vectors, deriving each LLM representation from prediction errors on representative prompts so routing can generalize to unseen models.

  • Results

    UniRoute effectively routes among more than 30 unseen LLMs across public benchmarks and provides competitive quality-cost trade-offs, outperforming K-NN on EmbedLLM and ZeroRouter across four datasets.

  • Takeaways & Limitations

    Prediction-error representations provide a principled basis for routing across dynamic LLM pools without retraining the router for every newly available model.

  • Takeaways & Limitations

    UniRoute assumes each new LLM can be efficiently evaluated on a modest validation set, and its performance depends on the choice of representative prompts.

Abstract

from arXiv · show

Model routing is a simple technique for reducing the inference cost of large language models (LLMs), wherein one maintains a pool of candidate LLMs, and learns to route each prompt to the smallest feasible LLM. Existing works focus on learning a router for a fixed pool of LLMs. In this paper, we consider the problem of dynamic routing, where new, previously unobserved LLMs are available at test time. We propose UniRoute, a new approach to this problem that relies on representing each LLM as a feature vector, derived based on predictions on a set of representative prompts. Based on this, we detail two effective instantiations of UniRoute, relying on cluster-based routing and a learned cluster map respectively. We show that these are estimates of a theoretically optimal routing rule, and quantify their errors via an excess risk bound. Experiments on a range of public benchmarks show the effectiveness of UniRoute in routing amongst more than 30 unseen LLMs.

1 Introduction

Model routing reduces LLM inference cost by selecting lower-cost models for easier prompts, but existing approaches generally assume a fixed model pool. UniRoute addresses changing pools by representing LLMs through prediction-error features and routing to previously unseen models without retraining.

  • Motivation: Model routing selects the lowest-cost LLM that can reasonably address each prompt, reserving expensive models for harder inputs.This provides an inference-efficiency strategy for pools containing models with different sizes and capabilities.
  • Problem: Frequent changes to the candidate pool make router retraining impractical because it requires computational overhead and sufficient labels for each new LLM.Older models may be deprecated as newer models become available.
  • Approach: UniRoute represents each LLM with a feature vector derived from prediction errors on representative prompts, enabling routing to previously unseen LLMs without retraining.The approach learns a router over these LLM features.
  • Contributions: The paper develops UniRoute instantiations based on unsupervised or supervised prompt clustering, with an accompanying excess risk bound.The cluster-based design is illustrated by partitioning validation prompts into representative clusters and using per-cluster errors.
  • Evaluation: > 30 unseen LLMs are evaluated across public benchmarks, illustrating UniRoute’s ability to route in a dynamic model pool.The experiments cover EmbedLLM, SPROUT o3-mini, RouterBench, and Chatbot Arena.

2 Background: Model Routing with a Static LLM Pool

Static model routing predicts each candidate LLM’s prompt-specific loss and combines it with inference cost to choose a suitable model. Its learned parameters are tied to the fixed candidate pool, motivating representations that can generalize across models.

  • Prediction setup: LLMs map prompts to predictions, and routing seeks a predictor that minimizes task loss while selecting among candidate models.The framework allows task-specific prediction functions and losses, including exact-match 0-1 loss.
  • Static routing: Static routing selects an LLM for each prompt from a fixed pool ordered by inference cost.The router aims to gain efficiency by using larger models sparingly on difficult inputs.
  • Routing strategies: A learned router predicts each LLM’s expected loss and selects a model using those predictions together with cost considerations.This general strategy includes approaches based on text embeddings and learned model-specific predictors.
  • Learned predictors: Text embeddings can parameterize prompt representations, while linear predictors may be fitted from training samples.The surveyed alternatives also include matrix factorization and K-NN estimators.

3 Model Routing with a Dynamic LLM Pool

Dynamic routing allows the candidate LLM set at test time to differ from the training set, including completely unseen models. The optimal rule selects the model with the lowest cost-adjusted expected loss, motivating a router that accepts both prompts and candidate sets.

  • Dynamic setting: Dynamic routing formalizes the case where the test-time LLM pool may differ from the training pool, including disjoint pools.This extends the static routing setup to changing model collections.
  • Motivation: Retraining a router for each new pool requires new model predictions, iterative training, and a fresh deployment, creating non-trivial overhead.This motivates an alternate routing setup for constantly refreshed pools.
  • Router formulation: A dynamic router takes both a prompt and a candidate LLM set as input and chooses the best available option from that set.The formulation targets generalization across sets of LLMs drawn from a broader collection.
  • Optimal rule: The Bayes-optimal rule routes to the LLM with the lowest expected loss after adding a cost penalty weighted by λ_H.The hyperparameter trades off expected quality and average inference cost.
  • Procedure: UniRoute operationalizes dynamic routing by training a base router, embedding test-time LLMs, and routing new prompts using those representations.The procedure computes each test LLM’s representation and then applies the routing rule to new inputs.

4 UniRoute: Universal Routing via an LLM Feature Representation

UniRoute represents prompts and LLMs as feature vectors so a router trained on one model pool can handle unseen LLMs. Its LLM features are derived from performance on representative validation prompts, while clustering provides concrete routing instantiations.

  • Core representation: UniRoute parameterizes routing as an interaction between a prompt feature map Φ(x) and an LLM feature map Ψ(h).The representation-based formulation is intended to support universal routing with dynamic LLM pools.
  • Unseen models: An easily computable Ψ(h) lets the router handle LLMs that were unobserved during training.Parameters associated with the feature maps can be fitted on the training set before routing over new models.
  • Prompt representation: Prompt features may use general-purpose text embeddings projected into a lower-dimensional space.The paper lists several general-purpose embedding models as possible starting points.
  • LLM representation: LLM features based on trained parameters are tied to the observed pool, while flattening model parameters is high-dimensional and unsuitable for many proprietary models.These limitations motivate performance-based representations.
  • Error-vector features: UniRoute represents an LLM using its prediction-error vector on a small validation set, optionally compressed into K dimensions.The representation treats LLMs with comparable performance on representative prompts as similar.
  • Relation to prior work: K-NN routing is a special case of UniRoute, whereas learning a compressed representation can use K ≪ Nval dimensions.The paper notes that prior feature representations did not generalize to unseen LLMs in the dynamic-routing setting.
  • Assumption: UniRoute assumes that each new LLM can be efficiently evaluated on a modest validation set, whose prompt choice affects the resulting representation.The validation prompts may be sampled, curated, or drawn from a standard benchmark suite.
  • Comparison: K-NN supports unseen LLMs without retraining but may generalize poorly with modest validation data and does not use the larger training set.This motivates the cluster-based UniRoute instantiation.

5 UniRoute with Cluster-Based LLM Feature Representations

UniRoute represents LLMs through cluster-level prediction errors and routes prompts using estimated performance on their assigned clusters. Its unsupervised and supervised cluster-map variants support efficient routing with previously unseen LLMs, while their quality gap from optimal routing is bounded by per-cluster versus per-prompt error discrepancies.

  • Unsupervised cluster-based routing: UniRoute represents each LLM by average errors over K predefined prompt clusters and each prompt by its cluster membership.The resulting estimate uses performance on similar prompts to approximate an LLM’s performance on the current input.
  • Unsupervised cluster-based routing: K-means clusters training prompts without labels, assigns validation prompts to clusters, and computes each new LLM’s per-cluster validation errors.The validation errors form the feature vector used for subsequent routing.
  • Unsupervised cluster-based routing: The cluster-based router incurs a one-off cost to estimate a new LLM’s error vector, after which routing is independent of later pool changes and requires no gradient updates.Further routing operates entirely on the estimated vector.
  • Supervised learned cluster map: The learned cluster map uses training labels to map prompts to a distribution over clusters, while retaining per-cluster errors estimated on validation data.Its parameters are selected by minimizing log loss against training-model correctness labels.
  • Excess risk bound: Under a mixture model and regularity condition, the excess 0-1 risk of cluster-based routing relative to the optimal rule is bounded by per-cluster versus per-prompt error discrepancies.The bound identifies the approximation gap as the difference between an LLM’s prompt-specific error and the average error of its constituent cluster.

6 Related Work

Related work includes routing, cascading, selective prediction, model fusion, mixture-of-experts, speculative decoding, and early-exiting approaches. These methods differ in whether they trade quality for cost, invoke multiple models, or operate over fixed model components.

  • Model routing: Model routing learns which candidate LLM should handle each prompt to reduce inference cost while preserving quality.The broader routing literature includes neural meta-models, nearest-neighbor methods, and matrix factorization.
  • Model cascading: Model cascading invokes models sequentially by cost and uses output statistics such as confidence to decide whether to continue.This differs from routing methods that select among candidate models for a prompt.
  • Selective classification and learning to defer: Selective classification and learning to defer provide formal foundations related to routing and cascading.The paper connects these approaches to learning to reject and deferring to an expert.
  • Model fusion: Model fusion primarily combines multiple models to improve quality, potentially at the expense of efficiency, rather than selecting one model to control inference cost.Fusion may invoke multiple models or produce a single fused model.
  • Mixture of experts: Classical mixture-of-experts methods learn routing among typically same-sized submodels, so cost considerations generally do not determine their routing rule.This contrasts with routing across LLMs of differing costs.
  • Other efficiency methods: Speculative decoding uses a smaller model to draft tokens and a larger model to verify them, while early-exiting stops computation within one neural model.Both are efficiency techniques related to, but distinct from, model routing.

7 Experiments

Experiments evaluate UniRoute for dynamic LLM pools across multiple public benchmarks, using deferral curves to measure quality-cost trade-offs. UniRoute generalizes to unseen models, performs well with small validation samples, remains robust across cluster counts, and is typically comparable to baselines in static pools.

  • Experimental setup: UniRoute is evaluated on EmbedLLM, RouterBench, Math+Code, SPROUT o3-mini, and Chatbot Arena under settings with unseen test-time LLMs.Training and testing LLM pools are separated for most datasets; Math+Code uses all four LLMs for testing and has unlabeled training data.
  • Evaluation: The experiments use binary accuracy and deferral curves that measure the trade-off between average quality and inference cost.Costs are parameter counts for EmbedLLM and API prices for RouterBench and SPROUT o3-mini.
  • Dynamic-pool results: UniRoute provides competitive quality-cost trade-offs, significantly improves over K-NN on EmbedLLM’s more than 30 unseen LLMs, and consistently outperforms ZeroRouter across all four datasets.The reported table metrics are area under the deferral curve, area up to 50% cost, and quality-neutral cost.
  • Validation-sample robustness: UniRoute is often significantly better than K-NN across validation sample sizes, because its variants can exploit training data whereas K-NN relies on retrieved validation neighbors.Figure 2 reports areas under the deferral curve with confidence intervals for varying validation-sample sizes.
  • Robustness and representation analysis: UniRoute remains effective across several cluster counts and typically performs comparably to most baselines in static LLM pools.Qualitative analysis also finds intuitive groupings, such as coding-specialist LLMs, in the prediction-error embedding space.

8 Conclusion and Future Work

The paper presents prediction-error-vector representations and principled routing strategies for multiple unseen test-time LLMs. It identifies robustness to prompt distribution shifts as a future direction, including dynamically varying representative prompts.

  • Conclusion: UniRoute provides principled strategies for routing among multiple unseen test-time LLMs using prediction error vector representations.The representation is the paper’s central mechanism for dynamic-pool routing.
  • Future work: Future work could improve robustness to prompt distribution shifts by allowing the representative prompt set to vary dynamically.The paper links this direction to reducing the need for frequent router retraining.

A Limitations

The paper’s limitations include no performance-recovery guarantee in fully static pools, an incompletely explored UniRoute design space, and potential bias-related societal risks.

  • B Societal Impact: Routers may overly favor models whose outputs are systematically biased on certain data sub-groups.The authors view incorporating constraints against such outcomes as a future research direction.

C.1 Intermediate Results

The intermediate results characterize optimal routing through cost-adjusted affine objectives and establish continuity-based properties of the optimal rule, including budget-constrained Lagrangian equivalence.

  • C.1 Intermediate Results: The optimal router selects the model minimizing a cost-adjusted loss for each fixed candidate pool and input.This follows from decomposing the overall objective into per-model affine functions of the Lagrange multiplier.
  • C.1 Intermediate Results: Changing the multiplier can change the optimal model only when two affine model objectives cross.The crossing point marks a transition in the lowest-valued affine function.
  • C.1 Intermediate Results: As the multiplier interval shrinks, the probability of encountering a crossing point tends to zero under the stated continuity assumptions.This yields convergence of the associated cost quantity as Δλ → 0.
  • C.1 Intermediate Results: For any admissible budget, some nonnegative Lagrange multiplier makes the unconstrained objective’s minimizer also solve the constrained routing problem.The result relies on continuity of the cost of the optimal rule as a function of the multiplier.
  • C.1 Intermediate Results: The continuity assumption applies when inputs are represented as fixed-length sentence embeddings.This is identified as the primary setting of the work.

C.3 Proof of Proposition 2

The proof of Proposition 2 replaces the original risk with a proxy risk, shows that cluster-based routing minimizes the proxy objective, and bounds the resulting excess 0-1 risk.

  • C.3 Proof of Proposition 2: The proxy objective is defined using an expectation over the joint distribution of inputs and proxy labels rather than true outputs.The proof then compares the proxy-optimal rule with the Bayes-optimal rule under a shared budget constraint.
  • C.3 Proof of Proposition 2: The proxy constrained optimization problem has a minimizer with the same form as the cluster-based routing rule.This establishes the cluster-based rule as optimal for the proxy objective under Proposition 2’s assumptions.
  • C.3 Proof of Proposition 2: The difference between the original 0-1 risk and the proxy risk can be bounded for any routing rule and fixed model pool.This bound connects optimization of the proxy objective to performance under the original risk.
  • C.3 Proof of Proposition 2: The K = 1 cluster-based rule returns the same LLM for every query and is closely aligned with the Pareto-random router.For general λ, the routing rule returns an LLM on the non-decreasing convex hull of cost-risk pairs.

E.1 Splitting Data and LLMs

The experiments split both examples and LLMs into training, validation, and testing portions to evaluate routing to previously unseen models. UniRoute represents test-time LLMs using validation data, while model-selection and routing overhead are handled separately.

  • Data and LLM splits: Each dataset separates examples into train, validation, and test splits, and LLMs into disjoint training and testing sets.Testing models are unavailable during ordinary router training and arrive at deployment time.
  • Data and LLM splits: Training examples provide correctness labels only for training models, whereas the validation split represents each test-time LLM as a feature vector.The validation set is specifically used to estimate representations for unseen LLMs.
  • Data and LLM splits: For Math+Code, no training LLMs are available, so the training sample is unlabeled.This is a dataset-specific boundary of the experimental setup.
  • UniRoute implementation: UniRoute (LearnedMap) maps frozen prompt embeddings to cluster probabilities using an MLP with two hidden layers.The implementation uses a frozen embedding model and a probability vector over K clusters.
  • UniRoute implementation: UniRoute (K-Means) routes by embedding a query and selecting its nearest centroid, while LearnedMap passes the embedding through a small MLP.K-means is run once on the training set; routing-model costs are excluded from plotted deferral curves.
  • Robustness and validation: K-Means Attributes outperforms K-Means Gecko in the distribution-shift experiment and is the only method reaching the pool’s most accurate model, attaining finite QNC.The result suggests prompt-hardness attributes improve robustness to prompt distribution shifts.

F.4 Static LLM Pool Setting

In the static-pool setting, UniRoute is generally competitive with existing methods, although an MLP often has a slight advantage because it can learn fixed-LLM representations directly. The embedding visualizations show meaningful similarity structure among LLMs.

  • Static-pool performance: In the static setting, UniRoute is typically comparable to most baselines, while the MLP baseline often has a slight edge.Here all LLMs are observed during training, unlike the dynamic setting that motivates the paper.
  • Static-pool performance: The MLP’s static-setting advantage reflects its flexibility to learn a separate representation for each fixed LLM.This flexibility is unavailable in the dynamic setting because new LLMs are not observed during training.
  • Evaluation: Table 2 evaluates static routing using area under the deferral curve, with best baseline and best UniRoute results separately highlighted.The static condition is defined by Htr = Hte.
  • Static-pool performance: UniRoute (K-Means) and UniRoute (LearnedMap) obtain area-under-deferral-curve values of .682 and .683, respectively, on EmbedLLM.The reported table compares UniRoute variants across multiple datasets and metrics; these values are the listed EmbedLLM results.
  • Embedding visualization: LLM-embedding heatmaps show Gaussian-kernel similarities between model pairs, including high similarity within the claude family on RouterBench.Code-focused models also tend to be highly similar on EmbedLLM.
Loading 2502.08773v2…