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

Because LLM inference can be prohibitively costly, this paper addresses routing when the candidate pool changes dynamically. UniRoute uses prediction-error features and prompt clustering to route among unseen LLMs; experiments report gains over K-NN on EmbedLLM and consistent outperformance of ZeroRouter across four datasets.

  • Problem

    The paper addresses routing when the candidate LLM pool changes dynamically, in a setting where LLM inference cost can be prohibitive.

    Source: §1 · p. 1, §3 · p. 3

  • Method

    UniRoute represents LLMs with prediction-error features from representative prompts and learns a router with unsupervised or supervised prompt clustering to generalize to unseen LLMs.

    Source: §1 · p. 2, §4.1 · p. 5

  • Results

    On EmbedLLM (>30 unseen LLMs), UniRoute shows compelling gains over K-NN on all metrics and consistently outperforms ZeroRouter across four datasets.

    Source: §7.2 · p. 10

  • Takeaways & Limitations

    UniRoute provides principled strategies for routing among multiple unseen test-time LLMs.

    Source: §8 · p. 11

    UniRoute assumes each new LLM can be efficiently evaluated on a modest set of representative prompts to compute its representation.

    Source: §4.3 · p. 6

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

UniRoute addresses dynamic LLM-pool routing by representing models with prediction-error features on representative prompts, enabling routing to unseen models without retraining. It introduces clustering-based instantiations and evaluates routing among more than 30 unseen LLMs.

Introduction

  • Model routing seeks the lowest-cost LLM that can reasonably address each prompt, using high-cost models sparingly on harder inputs.Source: §1 · p. 1
  • Existing routing work largely assumes a fixed LLM pool, while frequent model changes make retraining impractical because it incurs overhead and requires sufficient labels for each new model.Source: §1 · pp. 1–2
  • Its two instantiations use unsupervised or supervised prompt clustering, with cluster-based routing selecting the model with the smallest cost-adjusted average error for the prompt’s cluster.The prompt embedder can be unsupervised or trained with labels from LLMs different from those available at test time.Source: Figure 1 · p. 2, §1 · p. 2
  • UniRoute represents each LLM by its prediction-error vector on representative prompts and learns a router that generalizes to unseen LLMs without retraining.The paper formalizes routing with a dynamic pool of LLMs.Source: §1 · p. 2
  • The paper accompanies its clustering-based methods with an excess-risk bound.Source: §1 · p. 2
  • Experiments on EmbedLLM, SPROUT o3-mini, RouterBench, and Chatbot Arena illustrate routing among >30 unseen LLMs.Source: §1 · p. 2

2 Background: Model Routing with a Static LLM Pool

This background defines LLM prediction as minimizing task loss and static routing as choosing among cost-ordered models under a fixed budget, then surveys learned routing strategies.

2 Background: Model Routing with a Static LLM Pool

  • An LLM’s conditional output distribution is converted into a target prediction by a task-specific prediction function, with the goal of minimizing loss.The function may use a standard decoding algorithm or additional processing, such as stripping non-numeric tokens.Source: §2 · pp. 2–3
  • Static routing selects among cost-ordered LLMs to meet a fixed budget while minimizing prediction risk.The setup assumes at least two models with inference costs ordered from lowest to highest.Source: §2 · p. 3
  • Learned routers can predict each model’s expected loss and trade off predicted quality against inference cost; this rule is described as a plug-in estimator of an optimal rule [Hu et al., 2024b, Narasimhan et al., 2022, Somerstep et al., 2025].The trade-off is controlled by a nonnegative hyperparameter.Source: §2 · p. 3
  • The section lists text embeddings as a basis for expected-loss prediction, alongside matrix factorization over frozen embeddings and K-NN estimation [Ong et al., 2025, Zhao et al., 2024a, Hu et al., 2024b].The matrix-factorization approach may be fitted using pairwise comparisons or pointwise correctness labels.Source: §2 · p. 3

3 Model Routing with a Dynamic LLM Pool

The section formalizes routing when the available LLM pool changes, then derives an optimal cost-adjusted routing rule and a plug-in approach for new candidate models.

3 Model Routing with a Dynamic LLM Pool

  • The section frames model routing as a problem in which the set of available LLMs may vary dynamically.The test candidate set may differ from the training set, including having no models in common.Source: §3 · p. 3, §3.1 · p. 4

3.1 Problem Setup

  • The dynamic setup targets routers that generalize from a training pool of LLMs to test pools that may differ completely, avoiding repeated retraining and deployment as models change.A fixed-pool router estimates parameters only for its observed LLMs; new releases and deprecated models make repeated annotation and iterative retraining costly.Source: §3.1 · pp. 3–4

3.2 Optimal Routing with a Dynamic Pool

  • For each prompt and candidate set, the Bayes-optimal router selects the model minimizing expected loss plus a cost adjustment λ_H · c(h), trading off quality against average cost.The multiplier λ_H can be swept to trace a cost-quality curve; the paper focuses on 0-1 loss while noting the rule adapts to other loss functions.Source: §3.2 · p. 4

3.3 Plug-in Routing with a Dynamic Pool

  • UniRoute estimates the optimal rule with a plug-in router: train its base router, embed each test LLM, then select a candidate for each new prompt.The approach estimates γ*(x, h) with γ(x, h) and uses test-model embeddings Ψ(h).Source: §3.3 · pp. 4–5

4 UniRoute: Universal Routing via an LLM Feature Representation

UniRoute enables dynamic routing by representing prompts and LLMs as feature vectors, allowing routing to extend to previously unseen LLMs when their representations can be computed.

4.1 The UniRoute Approach

  • UniRoute parameterizes routing with prompt and LLM feature maps, so an easily computed LLM representation supports routing to models unseen during training.Prompt features can use general-purpose text embeddings, while the LLM features are designed separately to avoid dependence on a fixed model pool.Source: §4.1 · p. 5

4.2 Representing an LLM via the Prediction Error Vector

  • LLMs are represented by prediction-error vectors on a small labeled validation set, with compressed features possible and K-NN on that set as a special case [Hu et al., 2024b].The approach assumes new LLMs can be evaluated efficiently on the validation prompts; similarity is defined by comparable performance on representative prompts [Thrush et al., 2024; Zhuang et al., 2024].Source: §4.2 · pp. 5–6

4.3 Discussion: UniRoute versus Standard Routing

  • With a modest validation set (e.g., O(10^3)), UniRoute avoids iterative router retraining, which can impose cumulative overhead and overfit as the LLM pool refreshes.Validation prompts may come from training data, domain-curated examples, or a standard benchmark suite.Source: §4.3 · p. 6
  • K-NN supports new LLMs without retraining but may generalize poorly on modest validation data and does not use the training set; its advantage is exploiting nonlinear structure [Hu et al., 2024b; Zhuang et al., 2024].UniRoute’s learned or compressed representation can incorporate training-set information, unlike K-NN’s prediction-error-based routing.Source: §4.3 · p. 6

5 UniRoute with Cluster-Based LLM Feature Representations

UniRoute represents LLMs using per-cluster errors and routes prompts through fixed or learned cluster assignments; its analysis bounds the risk gap to an optimal rule.

5.1 Representing an LLM via Per-Cluster Prediction Errors

  • For K > 1 clusters, UniRoute represents each LLM by its average validation error per cluster and each prompt by cluster membership, estimating performance from similar prompts.To avoid overfitting a modest validation set, it clusters training prompts and assigns validation prompts to those clusters. A new LLM’s per-cluster error vector is a one-off estimate that requires no gradient updates and can be reused for routing.Source: §5.1 · pp. 6–7

5.2 From Fixed to Learned Cluster Assignment Maps

  • Starting from the same clustering, UniRoute learns a prompt-to-cluster distribution from training labels to improve routing, fitting the map by log loss against correctness labels for training LLMs.The per-cluster errors used by the map are estimated from the validation set and do not depend on the learned assignment parameters.Source: §5.2 · p. 7

5.3 Excess Risk Bound

  • Under a regularity condition on P, Proposition 2 bounds the 0-1 risk gap between cluster-based routing and the optimal rule by the discrepancy between per-cluster and per-prompt errors.The discrepancy is between an LLM’s error on a prompt and the average error of the prompt’s constituent cluster.Source: §5.3 · p. 7

6 Related Work

The section maps model routing to prior router-learning strategies and distinguishes it from cascading, formal deferral frameworks, model fusion, and other ways to adapt computation across models or within a model.

Model routing

Model cascading

  • Cascading invokes models sequentially in cost order and uses output statistics such as confidence to decide whether to proceed to a costlier model.Cascading has a long history in computer vision [Viola and Jones, 2001, Wang et al., 2018, Swayamdipta et al., 2018, Rawat et al., 2021, Wang et al., 2022, Kag et al., 2023, Jitkrittum et al., 2023] and has also been applied successfully to LLMs [Varshney and Baral, 2022, Chen et al., 2023b, Gupta et al., 2024, Yue et al., 2024, Chen et al., 2024a, Nie et al., 2024, Chuang et al., 2025].Source: §6 · p. 8

Selective classification and learning to defer

  • Routing and cascading have formal connections to learning to reject, selective classification, and learning to defer to an expert [Chow, 1970, Bartlett and Wegkamp, 2008, Cortes et al., 2016] [Geifman and El-Yaniv, 2019, Narasimhan et al., 2024a,b] [Madras et al., 2018, Sangalli et al., 2023].Later work studies routing and cascading through these frameworks, following foundational studies by Trapeznikov and Saligrama [2013], Bolukbasi et al. [2017], and Mozannar and Sontag [2020] [Narasimhan et al., 2022, Mao et al., 2024a,b,c,d].Source: §6 · p. 8

Model fusion

Mixture of experts (MoE)

  • Classical mixture-of-experts models learn parameters for independent models and a routing rule [Jacobs et al., 1991, Jordan and Jacobs, 1993], while newer variants use sub-models within a larger model [Fedus et al., 2022, Zhou et al., 2022].Classical experts are typically the same size, so cost does not factor into router design; these models have also been proposed for specialization [Jang et al., 2023, Douillard et al., 2024].Source: §6 · p. 8

Early-exiting

  • Early-exiting adapts computation within one neural model by stopping at an intermediate layer, using either confidence or learned termination decisions [Teerapittayanon et al., 2016, Scardapane et al., 2020] [Xin et al., 2020, Schuster et al., 2022].Source: §6 · p. 9

Speculative decoding

7 Experiments

The experiments evaluate UniRoute on benchmarks with unseen test-time LLMs and find competitive quality-cost trade-offs, including robustness to validation-set size and cluster choice.

7.1 Experimental Setup

  • Experiments cover five benchmark sources and evaluate held-out LLMs using repeated example splits and quality-cost deferral curves.Across 400 trials, examples are split 60/10/30 for training, validation, and testing, except RouterBench and SPROUT o3-mini use 1% validation; EmbedLLM cost is parameter count, while RouterBench and SPROUT use API cost.Source: §7 · p. 9, §7.1 · pp. 9–10

7.2 Experimental Results

  • On EmbedLLM (>30 unseen LLMs), UniRoute gains over K-NN on all metrics, with statistically significant gains, and consistently outperforms ZeroRouter across all four datasets.The reported metrics are area under the deferral curve, area up to 50% cost, and quality-neutral cost.Source: §7.2 · p. 10
  • With small validation samples, UniRoute is often significantly better than K-NN, which cannot use training data to estimate test-model performance.UniRoute can exploit training data through unsupervised K-means or the supervised learned cluster map.Source: §7.2 · p. 10
  • UniRoute is effective across several cluster counts, and its hyperparameter selector picks an effective K.The paper reports this robustness in Appendix F.Source: §7.2 · p. 11
  • In static LLM pools, UniRoute typically performs comparably to most baselines.The paper reports this result in Appendix F.4.Source: §7.2 · p. 11
  • Prediction-error-vector embeddings show largely intuitive groupings of similar LLMs, including coding specialists.The paper presents this qualitative analysis in Appendix F.5.Source: §7.2 · p. 11

8 Conclusion and Future Work · Appendix

UniRoute presents principled strategies for routing among unseen test-time LLMs using prediction-error-vector representations. Future work could improve robustness to prompt-distribution shifts by dynamically varying the representative prompt set, reducing the need for frequent router retraining.

8 Conclusion and Future Work

  • The paper presents principled strategies for routing among unseen test-time LLMs using prediction-error-vector representations.Source: §8 · p. 11
  • Future work could improve routing robustness to prompt-distribution shifts by dynamically varying Sval, further reducing the need for frequent router retraining.Source: §8 · p. 11

A Limitations · B Societal Impact · C Proofs of Results in Main Body

The paper identifies limitations in static-pool performance and router design, discusses potential subgroup-bias risks and future constraints, and provides proof material on optimization reformulations.

A Limitations

  • In fully static LLM pools, UniRoute is not guaranteed to match methods such as the linear router (3), motivating hybrid static-dynamic routers.The limitation arises from dependence on the validation sample, which may be a subset of training data.Source: A Limitations · p. 23
  • A systematic study of other possible UniRoute instances beyond the proposed natural cluster-based routers remains future work.Source: A Limitations · p. 23

B Societal Impact

  • The work’s primary contribution is a mathematical formalism for routing among dynamic LLMs, instantiated with prompt clustering and evaluated on public benchmarks, typically using LLM accuracy.Source: B Societal Impact · p. 23
  • The authors foresee no immediate negative societal impact from the framework, while noting that routing could favor models with systematic bias on particular subgroups.Source: B Societal Impact · p. 23
  • The authors consider the framework amenable to constraints against undesirable routing outcomes and identify accounting for them as future research.Source: B Societal Impact · p. 23

C Proofs of Results in Main Body

C.1 Intermediate Results · C.2 Proof of Proposition 1

The intermediate results characterize how the penalty parameter shapes routing and establish that, under continuity assumptions, a penalty minimizer can solve the budget-constrained problem. The proof of Proposition 1 applies this result to dynamic routing and notes that the assumptions cover fixed-length sentence embeddings.

C.1 Intermediate Results

  • For fixed H and x, minimizing overall loss requires routing to the model with the lowest cost-adjusted loss L_H,x.Source: §C.1 · p. 24
  • As λ moves between two values, a router change occurs when the minimum among candidate models’ affine functions switches at a crossing.Each model’s cost-adjusted loss is affine in λ; if the minimizing model changes, two such functions cross between the penalty values.Source: §C.1 · p. 25
  • With continuous x and a finite model universe, the proof bounds router-change mass by finitely many crossing events whose probabilities vanish as the λ interval shrinks.Source: §C.1 · p. 26
  • Under continuous P(y|x) and P(x), every budget B in (0, C(r*_0)] admits a nonnegative penalty whose minimizer also solves the constrained problem.Source: §C.1 · p. 26

C.2 Proof of Proposition 1

  • For continuous random variables x and η(x)=P(y | x), any LLM meta-distribution H and eligible budget admit λ_H≥0 whose penalized minimizer is the constrained optimal dynamic router.The proposition applies to budgets B in (0, C(r*_0)].Source: §C.2 · p. 26
  • The continuity assumption covers the paper’s primary setting, where x is represented by a fixed-length sentence embedding in a Euclidean space.Source: §C.2 · p. 26

C.3 Proof of Proposition 2 … E Experimental Setup

The appendix justifies cluster-based routing through a proxy-risk argument, relates zero routing and K = 1 routing to cost-risk convex hulls, and points to additional experimental details.

C.3 Proof of Proposition 2

  • Under the regularity condition on P in Proposition 1, the difference in 0-1 risk between ˜r∗ and Bayes-optimal routing rule r∗ can be bounded.Source: §C.3 · p. 27
  • The proof uses a proxy-risk objective whose constrained minimizer has the cluster-based form, then bounds the difference between proxy and original risk.It uses the proxy minimizer’s optimality and the fact that the Bayes-optimal rule also satisfies the budget constraint.Source: §C.3 · p. 27

D Zero Routing

  • Zero routing randomizes between two LLMs on the non-decreasing convex hull of cost-risk pairs to satisfy budget B, and is described as a strong baseline [Hu et al., 2024b].The pair and mixing probabilities are selected to maximize expected quality on the validation sample while meeting the budget.Source: D Zero Routing · p. 29

D.1 Special Case: Optimal Cluster Routing Rule for K = 1

  • For K = 1, the optimal cluster-routing rule selects an LLM on the non-decreasing convex hull of cost-risk pairs.The passage says this rule is closely aligned with the Pareto-random router.Source: §D.1 · p. 29

E Experimental Setup

E.1 Splitting Data and LLMs … F.1 Hyper-parameter Choices & Statistical Significance

The appendices specify how evaluation separates prompts and models, measures cost–quality tradeoffs, and tunes routing implementations.

E.1 Splitting Data and LLMs

  • Experiments split examples into train, validation, and test sets and LLMs into training and testing pools, with testing models representing unseen deployment-time models.Training models are used to train routers, which are evaluated on test examples while routing among testing models.Source: §E.1 · p. 30
  • Clairvoyant oracle baselines can access testing-model correctness labels during training, while Math+Code has no training LLMs and therefore uses unlabeled training samples.The validation set represents new LLMs for UniRoute (K-Means), and test examples are reserved for evaluating routing methods.Source: §E.1 · p. 30

E.2 Evaluation: Deferral Curve

  • A deferral curve plots routed-model expected loss across cost budgets; quality metrics such as accuracy can instead show quality–cost tradeoffs.For each budget, the method computes a router and evaluates its resulting expected loss.Source: §E.2 · p. 30

E.3 Implementation Details of UniRoute (LearnedMap)

  • UniRoute (LearnedMap) uses a two-hidden-layer MLP to map prompts to a probability vector over clusters, with a frozen text embedding.The experiments use Gecko embeddings; training uses Adam for five epochs on EmbedLLM, MixInstruct, and RouterBench, with batch size 64 and learning rate 0.005.Source: §E.3 · pp. 30–31

E.4 Router Cost

  • Router overhead is excluded from deferral curves; including it would shift all curves right slightly without changing methods’ relative ordering.Inference uses a nearest-centroid lookup for K-Means or a small MLP for LearnedMap after prompt embedding.Source: §E.4 · pp. 32–33

F.1 Hyper-parameter Choices & Statistical Significance

  • K-NN, UniRoute (K-Means), and UniRoute (LearnedMap) select K by maximizing validation-set normalized area under the deferral curve.K-NN searches K from 5 to one third of the validation sample size; the UniRoute methods search from 3 to the validation sample size divided by 50, with LearnedMap trained using training-model correctness labels.Source: §F.1 · p. 33
  • The attribute-based variant embeds prompts using seven difficulty attributes, while Math+Code sets K to √Nval for K-NN and Nval/50 for K-Means.Attribute heads are trained against Gemini 1.5 Pro 002 annotations; significance uses 400 trials and a sign test at α = 0.01 for Figure 2 (Top), while bottom-plot confidence intervals use 100 trials for most datasets and 1000 for Math+Code. Clairvoyant MLP and matrix-factorization baselines are also specified.Source: §F.1 · pp. 33–34

F Additional Experimental Results

  • Appendix F adds experimental results omitted from the main text, including validation-set comparisons of area under the deferral curve across candidate K values.Testing models and test examples are excluded from hyperparameter selection.Source: §F.1 · p. 33

F.2 Deferral Curves and Additional Comparisons … F.5 Visualisation of LLM Embeddings

These appendices extend evaluation across dynamic and static LLM pools, test prompt-hardness representations under distribution shift, and visualize similarities among learned LLM embeddings.

F.2 Deferral Curves and Additional Comparisons

  • Figure 5 reports mean dynamic-pool deferral curves over 400 trials, from which the Area (50%), Area, and QNC statistics are derived.The trials randomly split examples from each dataset into training, validation, and testing sets.Source: §F.2 · p. 34, Figure 5 · p. 35
  • For Headlines, 10,000 prompts across 12 LLMs were split into six training and six test LLMs, with Gemma 2B output-probability quantiles used as prompt embeddings.Each trial divided examples into 4,000 training, 400 validation, and 5,600 testing examples.Source: §F.2 · p. 34
  • On Headlines, UniRoute (LearnedMap) has higher accuracy than K-NN throughout the entire cost range.The comparison is based on mean deferral curves over 400 trials in the dynamic pool setting.Source: §F.2 · p. 34, Figure 5 · p. 35

F.3 Train on Chatbot Arena and Test on EmbedLLM

  • Seven binary prompt-hardness attributes, annotated by Gemini 1.5 Pro, are distilled into a query embedder to improve robustness to test-time prompt-distribution shift.The attributes follow Li et al. [2024b].Source: §F.3 · p. 35
  • The cross-dataset test trains on Chatbot Arena pairwise preferences and evaluates on EmbedLLM correctness labels, using 26 shared LLMs and Bradley-Terry-Luce cluster quality estimates.After filtering, Chatbot Arena has 8,447 records; the test uses the full EmbedLLM dataset.Source: §F.3 · p. 35
  • K-means (Attributes) outperforms K-means (Gecko) under prompt shift and is the only method to reach the pool’s most accurate model, yielding finite QNC.The Pareto-random router’s declining trend is attributed to validation-set Pareto-optimal LLMs not being optimal on the test set.Source: §F.3 · p. 36, Figure 6 · p. 36

F.4 Static LLM Pool Setting

  • In the static pool, UniRoute is typically comparable to most baselines, while MLP often has a slight edge because it can learn its own LLM representation.That representation flexibility is possible only when all LLMs are seen during training.Source: §F.4 · p. 36, Table 2 · p. 37

F.5 Visualisation of LLM Embeddings

  • Embedding-similarity heatmaps show that Claude-family models are similar on RouterBench, while code-focused models are generally similar on EmbedLLM.The visualizations show Gaussian kernel similarity between pairs of cluster-based LLM embeddings.Source: §F.5 · p. 37, Figure 8 · p. 37, Figure 7 · p. 38

Related papers

Loading 2502.08773v2…