PaperlayerView exact revision on arXiv

Source-linked AI summary

Rethinking Predictive Modeling for LLM Routing: When Simple kNN Beats Complex Learned Routers

Yang Li

arXiv:2505.12601v2cs.LG

TL;DR

Whether sophisticated learned routers are necessary remains an open question, despite routing’s importance for cost-effective deployment and user experience. We evaluate simple kNN routing across text and multimodal benchmarks and find it matches or outperforms complex learned approaches, while running 13–14× faster.

  • Problem

    Whether increasingly sophisticated learned routing is necessary remains unclear, despite routing’s importance for cost-effective deployment and user experience.

    Source: §1 · p. 1

  • Method

    We evaluate simple kNN routing against complex learned approaches across standardized instruction-following, question-answering, and reasoning benchmarks and text and multimodal settings.

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

  • Results

    Across text and multi-modal benchmarks, kNN matches or outperforms complex learned approaches and is 13-14× faster than complex alternatives.

    Source: §8 · p. 8

  • Takeaways & Limitations

    Simple routing can reduce computational and engineering overhead for deploying sophisticated multi-model systems.

    Source: §1 · p. 2

    kNN may be less suitable for sparse longtail neighborhoods, rapidly changing performance landscapes, or compositional reasoning across domains, and its memory requirements grow linearly with support set size.

    Source: §8 · p. 8

Abstract

from arXiv · show

As large language models (LLMs) grow in scale and specialization, routing--selecting the best model for a given input--has become essential for efficient and effective deployment. While recent methods rely on complex learned routing strategies, their dependence on disparate training data and evaluation setups makes comparison and generalization difficult. In this work, we revisit LLM routing through the lens of simplicity. We show that a well-tuned k-Nearest Neighbors (kNN) approach not only matches but often outperforms state-of-the-art learned routers across diverse tasks. To support systematic evaluation, we introduce a suite of standardized routing benchmarks spanning instruction-following, question-answering, and reasoning tasks, as well as the first multi-modal routing dataset involving visual inputs. Our findings reveal that the locality properties of model performance in embedding space enable simple non-parametric methods to achieve strong routing decisions with lower sample complexity than parametric approaches. This challenges the prevailing trend toward sophisticated architectures and highlights the importance of thoroughly evaluating simple baselines before investing in complex solutions. To support reproducibility and further exploration, we will release all benchmarks and code upon publication.

1. Introduction

The paper questions whether complex learned routers are necessary, given inconsistent comparisons and limited evaluation of simple baselines. It argues that tuned kNN performs strongly across routing settings and introduces standardized text and vision-language benchmarks.

1. Introduction

  • Current routing methods use varied training data, evaluation protocols, and assumptions, making fair comparison and generalizability difficult.The field has increasingly adopted complex architectures without rigorous comparison to simple baselines.Source: §1 · p. 1
  • When carefully implemented and tuned, kNN matches or often outperforms complex learned routers across diverse routing scenarios, with computational-efficiency and robustness advantages.The paper also reports greater stability under distribution shift and lower sample complexity for non-parametric methods, which it links to locality in embedding spaces.Source: §1 · pp. 1–2
  • The benchmark suite standardizes evaluation across instruction-following, question-answering, and reasoning tasks.The paper also introduces the first benchmark for routing between vision-language models.Source: §1 · pp. 1–2
  • Embedding-space locality, where similar queries benefit from similar models, supports effective non-parametric routing with lower sample complexity than parametric approaches.The evaluation also reports greater robustness under distribution shift for simple methods.Source: §1 · p. 2
  • The authors argue that simpler routing can reduce computational and engineering overhead, and recommend evaluating simple baselines before investing in complex solutions.They present these as practical implications and guidance for practitioners.Source: §1 · p. 2

2. Background and Related Works

The paper situates routing among strategies for balancing model performance and computational cost, then highlights the need for systematic comparison of complex routers against simpler alternatives.

2. Background and Related Works

  • Ensembles aggregate outputs from multiple models but incur latency and computational overhead that scales linearly with the number of models.The related-work discussion contrasts this cost with cascading and routing approaches.Source: §2 · p. 2
  • Cascading invokes models sequentially, potentially lowering average cost, but challenging queries can require multiple calls and auxiliary quality estimation.The latter can add latency and system complexity.Source: §2 · p. 2
  • Routing selects one LLM per query, avoiding multiple calls; existing approaches predict model scores or pairwise preferences, or learn contextual-bandit policies.The paper focuses on routing for its efficiency potential and practical relevance.Source: §2 · p. 2
  • Routing research has increasingly favored complex architectures without systematic comparison to simpler alternatives.The paper calls for evaluating simple and complex methods systematically.Source: §2 · p. 2

3. Predictive LLM Routing

The paper frames routing as selecting a model that balances response quality against cost. It distinguishes direct model-selection policies from utility prediction and describes how support sets enable neighborhood-based routing.

3. Predictive LLM Routing

  • LLM routing selects a model from a candidate pool to process a query while balancing performance and cost.The performance score can represent quality, while cost can reflect latency or compute; λ sets the user-defined trade-off.Source: §3 · pp. 2–3
  • Model-selection policies directly map queries to models, while utility-prediction methods estimate each model’s score and cost before selecting the highest-utility option.The paper describes these two classes as complementary routing approaches.Source: §3 · p. 3
  • Support sets record outcomes for semantically similar prompts; kNN retrieves neighbors and aggregates their recorded performance and cost to estimate utility.Parametric routers can also process a query and its neighbors.Source: §3 · p. 3
  • Neighborhood-based routing can adapt at test time to distribution shifts without complete retraining when similar queries tend to benefit from similar models.The paper reports that this locality assumption holds well in its experiments.Source: §3 · p. 4

4. Routing Benchmarks

The paper introduces standardized text and vision-language routing benchmarks and separate evaluation protocols for utility-prediction and model-selection methods. It notes that fixed-preference accuracy can obscure cost-performance trade-offs.

4. Routing Benchmarks

  • The benchmark suite is designed to address inconsistent evaluation protocols that have hindered meaningful comparisons between routing methods.The text benchmark expands on RouterBench’s 11 models across 6 tasks with a broader range of models and tasks.Source: §4 · p. 4, §4.1 · p. 4

4.1. Text-Based Routing Benchmarks

  • The text benchmark draws on AlpacaEval, Open LLM Leaderboard v2, and HELM-Lite, using reported evaluation metrics for performance and API pricing based on token counts for cost.It selects three model families from each leaderboard, with multiple variants per family.Source: §4.1 · p. 4

4.2. Vision-Language Routing Benchmarks

  • The paper introduces the first vision-language routing benchmark, built on vHELM and curated across five visual question-answering and reasoning datasets.The benchmark supports evaluating routing where model choice depends on visual content as well as text.Source: §4.2 · p. 4

4.3. Evaluation Protocol

  • Utility-prediction methods are assessed by the area under a convex-hull cost-performance curve across λ values, while selection methods are tested at three fixed preferences.Fixed-preference accuracy can obscure true cost-performance trade-offs; the paper therefore focuses its main results on utility prediction and reports selection results in an appendix.Source: §4.3 · p. 4

5. Routing Approaches

The section compares routing methods of increasing complexity under consistent input representations and both utility-prediction and model-selection formulations. It covers neighborhood-based, linear, neural, graph, and attention approaches.

5. Routing Approaches

  • All methods support utility prediction and model selection using consistent input representations.Text routing uses BERT embeddings, while multimodal routing uses VLM2Vec embeddings.Source: §5 · p. 5
  • kNN retrieves the k most similar training examples and aggregates their outcomes, using averages for utility prediction or majority voting for model selection.It requires no training and adapts to new queries through local neighborhood information.Source: §5 · p. 5
  • Linear routers use regression over query embeddings or learnable model embeddings to predict routing outcomes.The linear matrix-factorization router models interactions between query and model embeddings.Source: §5 · p. 5
  • Neural routers use multilayer perceptrons to capture nonlinear query-to-outcome mappings, including through prompt-model interactions.The MLP uses multiple 100-dimensional hidden layers; the MLP matrix-factorization router extends the factorized linear approach.Source: §5 · p. 5
  • The Graph Router models queries and models as nodes linked by utility-valued edges, while attentive routers use support examples and cross-attention for contextual decisions.The Double Attentive Router adds attention across candidate models as well as queries.Source: §5 · p. 5

6. Quantitative Results

Across text benchmarks, kNN and other simple routers are competitive with or outperform more complex approaches, while kNN also shows lower routing time and stronger performance under distribution shift. The simple-method pattern extends to vision-language routing, with attention-based architectures showing strength on some multimodal tasks.

Text-Based Routing Results

  • Routing-method rankings remain consistent across benchmarks and model pools, and increasing k from 10 to 100 generally improves kNN performance.Linear models achieve an average AUC of 53.14.Source: §6 · p. 5
  • Across text benchmarks, kNN (k=100) averages AUC 52.68, competitive with or superior to MLP (51.71) and Graph Neural Networks (51.82).The text benchmarks use AUC of the Pareto front to capture performance-cost trade-offs.Source: §6 · p. 5, Table 2 · p. 6

Computational Efficiency Analysis

  • 65.69 seconds is the cumulative routing time for kNN (k=100), while graph and attention-based methods require over 866 seconds, approximately 13–14× slower.The timing excludes one-time index-building and model-training costs.Source: Table 3 · p. 6, §6 · p. 6

Robustness Under Distribution Shift

  • Under cross-dataset distribution shift, simple methods maintain more stable performance, with kNN adapting locally through relevant training examples.The evaluation compares in-distribution settings, where train and test datasets match, with out-of-distribution settings using different training datasets.Source: Table 4 · p. 6, §6 · p. 6

Multi-Modal Routing Results

  • In vision-language routing, kNN (k=100) achieves an average AUC of 72.12, outperforming most neural approaches.The results extend the simple-method pattern to multimodal inputs.Source: §6 · p. 6, Table 5 · p. 7
  • Attention-based architectures show particular strength on some multimodal tasks, such as MME.The passage attributes this likely to their ability to model visual-textual interactions more explicitly.Source: §6 · p. 7

Statistical Validation and Embedding Analysis

  • Across five runs, kNN (k=100) achieves 77.31 ± 0.27 AUC, Linear 77.52 ± 0.21 AUC, and MLP 76.94 ± 0.33 AUC.The small standard deviations indicate that the reported relative performance is consistent across random seeds or data splits.Source: §6 · p. 7
  • Switching from BERT to SFR embeddings provides modest improvements, while the relative ranking of routing methods remains consistent.Detailed results are reported in Appendix I.Source: §6 · p. 7

7. Theoretical Analysis

The theoretical analysis links kNN routing to locality in query-performance space and compares its sample complexity with parametric routers. Empirical evidence supports the locality property and the low-dimensional conditions associated with kNN’s sample-complexity advantage.

Theoretical Analysis

  • δ-locality formalizes the condition that model utility changes locally across query embeddings, supporting routing from nearby queries’ performance patterns.The paper describes semantically similar queries as expected to have similar utility scores from the same model.Source: §7 · p. 7
  • Under δ-locality, kNN requires fewer training samples than a parametric router to achieve the same expected-regret bound.The theorem gives kNN a sample-size requirement tied to intrinsic dimension d, while a parametric router with L Lipschitz-continuous layers requires Ω(L/ϵ(δ)^2).Source: §7 · p. 7
  • The analysis predicts a kNN sample-complexity advantage when embedding dimension is low and ϵ(δ) decreases rapidly with δ.The paper states that these conditions can yield the same performance guarantees with significantly fewer samples than parametric routers.Source: §7 · p. 7

Empirical Validation

  • As embedding distance increases, model-performance ranking agreement decreases; the reported correlations are r=-0.815 for Arc-Challenge and r=-0.875 for GSM.Agreement is measured as the correlation between performance scores across models for each query pair, with pairs binned by embedding distance.Source: §7 · pp. 7–8, Figure 1 · p. 8
  • Measured intrinsic dimensions are approximately 2–28 for RouterBench and 13–18 for vision-language embeddings, below their ambient dimensions.The paper connects low intrinsic dimension combined with strong locality to the conditions under which kNN gains sample-complexity advantages.Source: §7 · p. 8

8. Discussion and Conclusion

Across text and multimodal benchmarks, simple kNN routers match or outperform complex learned approaches, with efficiency and robustness advantages. The paper links this performance to embedding-space locality and gives diagnostics for deciding when kNN is suitable.

8. Discussion and Conclusion

  • 52.68 AUC on text and 72.12 on multimodal benchmarks, kNN routing was 13–14× faster than complex alternatives.Under distribution shift, kNN degraded by 2.63 points versus 6.67 points for complex methods.Source: §8 · p. 8
  • Strong locality and low intrinsic dimensionality in embedding spaces help explain kNN’s routing effectiveness with lower sample complexity.The benchmarks show intrinsic dimension d ≈ 2–28; the reported locality is associated with semantically similar queries benefiting from similar models.Source: §8 · p. 8
  • Practitioners can assess kNN suitability using embedding-distance/performance-agreement correlation and intrinsic dimension, then monitor neighbor distances and agreement for query-level uncertainty.The reported benchmark locality correlations are r = −0.815 to −0.875, with intrinsic dimension d ≈ 2–28 versus ambient dimensions of 768–3584.Source: §8 · p. 8
  • kNN may be less suitable for sparse long-tail neighborhoods, rapidly changing performance landscapes, or compositional reasoning across domains, and its memory use grows linearly with support-set size.The authors identify these as settings where complex approaches may be justified.Source: §8 · p. 8
  • The authors argue that competitive performance, efficiency, and robustness make simple routing attractive for deployment, while calling for work on embeddings, training signals, and identifying when complexity is needed.They frame reduced computational and engineering overhead as a potential way to broaden access to effective multi-model systems.Source: §8 · p. 8

Impact Statement

The paper presents routing benchmarks, evaluation protocols, and model implementations for text and vision-language tasks. It also describes potential deployment benefits and the connection between routing and recommendation systems.

Impact Statement

  • The paper argues that reducing computational and engineering overhead could broaden access to effective multi-model systems, while warning that simple methods may be unsuitable with weak locality or high dimensionality.It says diagnostic tools can help practitioners choose an appropriate method.Source: Impact Statement · p. 9

A. Additional Related Works

  • The paper connects routing to recommendation systems: queries act like users and models like items, suggesting analogies to neighborhood, matrix-factorization, and content-based methods.It points to cold-start handling and hybrid routing as possible opportunities for adapting recommendation techniques.Source: A. Additional Related Works · p. 11

B.1. Text-Based Routing Benchmarks

  • The benchmark suite combines AlpacaEval, Open LLM Leaderboard v2, HELM-Lite, and RouterBench, with text tasks spanning instruction following, reasoning, and other task categories.RouterBench includes six tasks, while the selected model families and evaluation metrics vary by benchmark.Source: §B.1 · pp. 11–12

B.2. Vision-Language Routing Benchmarks

  • The paper introduces its first benchmark for routing between vision-language models, using vHELM outcomes and five datasets with varied visual-task challenges.The datasets include visual perception, image description, visual mathematics, broad VLM capabilities, and academic-domain understanding.Source: §B.2 · p. 13

B.3.1. AUC SCORE CALCULATION METHODOLOGY

  • For utility prediction, the evaluation selects models across cost preferences, forms a non-decreasing convex hull of actual cost-performance pairs, and scores its normalized area.This AUC approach evaluates trade-offs across cost-performance preferences.Source: §B.3.1 · p. 14

B.4. Data Splits and Reproducibility

  • The benchmarks use 70/10/20 prompt splits, prevent leakage between support sets and test queries, and fix and document random seeds.The authors describe the benchmark construction as spanning diverse tasks, models, and modalities.Source: §B.4 · p. 14

C. Routing Approaches

  • Text queries use normalized BERT embeddings and multimodal inputs use Qwen7B-based VLM2Vec representations; the evaluated routers include kNN, linear, factorization, MLP, graph, and attention-based approaches.The kNN router uses cosine similarity and predicts performance and cost by weighted averages over neighbors.Source: §C.1 · p. 14, §C.2 · p. 15

C.3. Training Protocols

  • Utility-prediction models train with a weighted sum of score and cost mean-squared errors, while selection models use cross-entropy with labels defined by the optimal model under a trade-off parameter.The coefficient α balances performance and cost prediction.Source: §C.3 · p. 16

D. Selection-based Routing Evaluation

Selection-based evaluation compares routers at fixed cost-performance preferences, but preference settings may not control cost monotonically and utility scores can hide trade-offs. The section also reports detailed benchmark breakdowns and scalability findings for kNN.

D. Selection-based Routing Evaluation

  • Averaged utility scores and fixed-preference accuracy can obscure actual cost-performance trade-offs, including whether similar scores come from costlier, higher-performing models or more cost-efficient ones.The authors recommend utility-prediction evaluation and visualizing actual cost-performance points for a fuller comparison.Source: D. Selection-based Routing Evaluation · pp. 16–17
  • Preference settings often fail to produce a monotonic trend, making precise cost-budget control difficult for selection-based routing.The section presents utility-prediction evaluation as a more flexible way to explore the Pareto front.Source: D. Selection-based Routing Evaluation · pp. 17–18
  • In many cases, kNN achieves comparable or better performance than neural architectures at similar cost efficiency in text and VLM routing cost–quality comparisons.The figures visualize actual cost-performance relationships rather than only combined utility scores.Source: D. Selection-based Routing Evaluation · p. 17, Figure D.1 · p. 18, Figure D.2 · p. 18

E. Detailed Results

  • The appendix provides RouterBench results averaged across six datasets and detailed results by dataset and preference setting for text and multimodal routing.Tables E.1 and E.2 cover utility prediction and selection, while E.3 and E.4 break down three preference settings.Source: E. Detailed Results · p. 18, Table E.1 · p. 18, Table E.2 · p. 19

F. Scalability Analysis

  • The scalability analysis reports that kNN routing scales efficiently to production scenarios with millions of queries.The section examines memory, index construction, and retrieval latency.Source: F. Scalability Analysis · p. 19
  • kNN support-set memory grows at about 3 KB per query, reaching 30 GB for 10M queries; index construction is about 0.28 ms per query and under an hour at 10M queries.The authors also report that adding examples to an index takes milliseconds, unlike complete retraining for parametric methods.Source: §F.1 · p. 19, §F.2 · p. 20

G. Latency Analysis

Across latency and distribution-shift evaluations, kNN is faster than complex routing methods and shows a smaller average performance drop on out-of-distribution queries. Cross-dataset results also vary by model and dataset pair.

G. Latency Analysis

  • 10.95 seconds per dataset was the average processing time for kNN (k=100), the fastest method overall.Graph and attention-based methods took over 144 seconds per dataset on average, approximately 13–14× slower than kNN.Source: G. Latency Analysis · p. 21, Table G.1 · p. 22

H. Out-of-Distribution Queries

  • OOD generalization varied by dataset pair: models trained on hellaswag performed well on mmlu but poorly on mbpp, while gsm-trained models showed the reverse pattern.The authors attribute these patterns to differences in domain shift between datasets.Source: H. Out-of-Distribution Queries · p. 22
  • 2.63 points was the smallest OOD performance drop, for kNN (k=100), versus 6.67 points for Linear (MF).All routing approaches degraded on OOD queries on average.Source: H. Out-of-Distribution Queries · p. 22, Table H.1 · p. 23
  • No routing model universally dominated across cross-dataset evaluations, underscoring the importance of data diversity and robustness-aware design.Latency measurements excluded one-time index-building and model-training costs; CPU and GPU methods ran on different hardware.Source: G. Latency Analysis · p. 20, H. Out-of-Distribution Queries · p. 22

I. Embedding Analysis

The embedding analysis finds that routing-method rankings remain similar across embedding models, while the proof compares kNN and parametric routers’ sample complexity under locality and intrinsic-dimension conditions.

I. Embedding Analysis

  • Switching from BERT to SFR embeddings modestly improved most methods, but the relative ranking of routing methods remained similar.The authors say neighborhood-size importance across embeddings supports locality as a fundamental property of routing rather than an artifact of a particular representation.Source: §I · p. 24, Table I.1 · p. 24

J. Proof

  • The theorem states that a parametric router with L Lipschitz-continuous layers needs Ω(L/ϵ(δ)^2) training samples for the stated regret bound.The bound targets expected regret O(ϵ(δ)) with probability 1−α.Source: J. Proof · p. 24

Part 1: kNN Router Sample Complexity

  • If every query has at least k neighbors within distance δ, locality bounds each model’s utility-estimation error by ϵ(δ), yielding kNN regret O(ϵ(δ)).The proof establishes neighborhood coverage with probability at least 1−α.Source: Part 1: kNN Router Sample Complexity · p. 26
  • For small ϵ(δ), the parametric router’s sample-complexity lower bound grows quadratically with 1/ϵ(δ).The stated dominant term for small ϵ and moderate L is Ω(L/ϵ^2).Source: Part 1: kNN Router Sample Complexity · p. 27
  • With low intrinsic dimension and strong locality, kNN requires significantly fewer training samples than a parametric router for the same regret bound.The comparison says kNN complexity depends on 1/δ^d and logarithmically on 1/α, while the parametric bound grows quadratically with 1/ϵ(δ).Source: Part 1: kNN Router Sample Complexity · p. 27

Related papers

Loading 2505.12601v2…