Source-linked AI summary
BEST-Route: Adaptive LLM Routing with Test-Time Optimal Compute
Dujian Ding, Ankur Mallick, Shaokun Zhang, Chi Wang, Daniel Madrigal, Mirian Del Carmen Hipolito Garcia, Menglin Xia, Laks V. S. Lakshmanan, Qingyun Wu, Victor Rühle
TL;DR
LLM routing must reduce deployment cost without sacrificing response quality, but prior approaches often overuse expensive models when small-model responses are insufficient. BEST-Route combines difficulty-aware model selection with cost-aware best-of-n sampling to improve smaller-model performance. Across real-world evaluations, it reports up to 60% cost reduction with less than 1% performance drop, while depending on accurate proxy-model ranking.
Problem
Prior routing methods often underuse smaller models because a single inexpensive response may not match the quality of an expensive model, limiting cost savings.
Method
BEST-Route uses a cost-efficient difficulty-aware router and test-time best-of-n sampling to select models and computational resources for each query.
Results
Up to 60% cost reduction is achieved with less than 1% performance drop across evaluations on real-world datasets.
Takeaways & Limitations
Adaptive routing and cost-aware best-of-n sampling provide a flexible way to trade response quality against inference cost.
Takeaways & Limitations
Best-of-n selection depends on proxy reward model accuracy, and misalignment with ground-truth evaluations can produce suboptimal response choices.
Abstract
from arXiv · showhide
Large language models (LLMs) are powerful tools but are often expensive to deploy at scale. LLM query routing mitigates this by dynamically assigning queries to models of varying cost and quality to obtain a desired trade-off. Prior query routing approaches generate only one response from the selected model and a single response from a small (inexpensive) model was often not good enough to beat a response from a large (expensive) model due to which they end up overusing the large model and missing out on potential cost savings. However, it is well known that for small models, generating multiple responses and selecting the best can enhance quality while remaining cheaper than a single large-model response. We leverage this idea to propose BEST-Route, a novel routing framework that chooses a model and the number of responses to sample from it based on query difficulty and the quality thresholds. Experiments on real-world datasets demonstrate that our method reduces costs by up to 60% with less than 1% performance drop.
1 Introduction
BEST-Route targets the cost-quality tension in LLM deployment by routing queries according to difficulty and using adaptive test-time compute. Its best-of-n strategy improves the usefulness of smaller models, achieving substantial savings with minimal quality loss.
- Motivation: Prior routing methods often send all but the easiest queries to large models, limiting potential cost savings.They do not adequately use extra compute to improve smaller-model performance.
- Approach: BEST-Route combines a difficulty-aware router with test-time best-of-n sampling to balance response quality and inference cost.The router selects a model and allocates computational resources based on the query.
- Approach: Easy queries receive minimal sampling from smaller models, while harder queries can use larger models or additional computation.The framework predicts whether sampled small-model responses can match a powerful reference model.
- Implications: The framework is intended to provide a flexible cost-performance balance for LLM service providers and users.The authors connect adaptive routing and test-time optimization with more efficient LLM deployment.
2 Related Work
Related work addresses LLM efficiency through routing among models, test-time sampling, and decoding acceleration. BEST-Route builds on routing and best-of-n sampling, while differing from speculative decoding in its focus on distributing queries across models.
- Efficient ML Inference: Efficient inference research seeks to reduce the substantial computational and financial costs of deploying increasingly large LLMs.Prior approaches include static efficiency optimizations and smaller models.
- LLM Routing: LLM routing dynamically assigns queries among models to balance response quality and inference cost.Recent work extends routing from binary choices to larger collections of candidate models.
- Test Time Optimal Compute: Test-time optimal compute improves difficult-query outcomes by sampling multiple responses and selecting among them.Increasing the number of samples can raise the probability of finding correct solutions, especially for coding tasks.
- Other Multi-LLM Inference Techniques: Speculative decoding accelerates generation within one expensive model, whereas LLM routing distributes query traffic across multiple models.The two techniques optimize different parts of the inference process.
3 Problem Formulation
The problem formulation treats query difficulty, model choice, response sampling, evaluation quality, and inference cost as jointly relevant to routing. It motivates using smaller models for easy queries and larger models or more computation for difficult ones.
- 3.1 Motivation: Queries vary in difficulty, with straightforward transformations often suitable for smaller models.Difficulty depends on complexity, ambiguity, and task requirements.
- 3.1 Motivation: Complex or ambiguous queries may require larger models or additional test-time sampling to preserve response quality.The paper contrasts these queries with easy requests that can be served more cheaply.
- 3.1 Motivation: Routing only between a small and a large model can underuse the broader cost-quality spectrum available in modern model platforms.The paper argues that routing across available models is important for better trade-offs.
- 3.1 Motivation: Prior multi-model routers may still default to the largest model, while best-of-n sampling can enhance small-model response quality.Best-of-n generates multiple responses and selects the best one.
- 3.2 Problem Setting: The system chooses between one response from a powerful reference model and best-of-n responses from smaller, cheaper models.The router must select the model and sampling strategy for each query.
- 3.3 Evaluation Metric: Response quality is difficult to evaluate automatically because traditional metrics can align weakly with human judgment across tasks.The paper therefore adopts an LLM-based evaluator.
- 3.3 Evaluation Metric: armoRM is used to assess response quality because it is reported to perform strongly on evaluation-model benchmarks while remaining feasible for large-scale evaluation.It is a fine-tuned Llama3-8B model.
- 3.3 Evaluation Metric: Inference cost can be measured by FLOPs, latency, or monetary expense, with token-based expenses especially relevant to service users.The paper distinguishes hardware-independent FLOPs from practical deployment costs.
4 Routing Framework
BEST-Route combines proxy-guided best-of-n sampling with adaptive routing to choose a model and sample count for each query. It uses match probabilities, cost estimates, and quality thresholds to select cost-effective inference options.
- Memory Efficient Best-of-n Sampling: The proxy reward model is intended to preserve ground-truth response rankings so best-of-n selection can identify the highest-quality response.It is trained with pairwise ranking loss on response pairs ordered by ground-truth reward.
- Memory Efficient Best-of-n Sampling: BEST-Route uses a proxy reward model to score n generated responses and selects the response with the highest proxy score.The procedure generates samples, computes Rproxy for each, and returns s∗ = arg maxs∈S Rproxy(s).
- Test-time Optimal LLM Routing: BEST-Route compares a reference-model response with smaller-model best-of-n responses using a predicted match probability.The router estimates whether a smaller model's best-of-n response is at least as good as the reference response under ground-truth reward.
- Test-time Optimal LLM Routing: The inference pipeline predicts match probabilities, computes costs, samples from the selected model, and applies proxy-based best-of-n selection.Costs include estimated output-token costs and input-token costs, while output length is estimated from average training-data lengths.
- Test-time Optimal LLM Routing: A cost-efficient multi-head router shares a BERT-style query representation across K × N lightweight heads for model-and-sample-count choices.Each head predicts the probability that one model with a specific n matches the reference model.
- Test-time Optimal LLM Routing: At inference, users set match-probability thresholds, and BEST-Route selects among qualifying model-and-sampling combinations using estimated cost.Higher thresholds favor the reference model, while cost estimation chooses the least-cost option that meets the threshold.
5 Evaluation
BEST-Route is evaluated across diverse datasets, LLM portfolios, routing baselines, quality and cost metrics, in-distribution and out-of-distribution settings, and latency overhead. Results show improved cost-quality trade-offs, including substantial savings with small quality degradation and gains from adaptive sampling and specialized models.
- Evaluation setup: The evaluation uses 10K instruction examples across question answering, coding, and safety, evaluates eight LLMs, and includes OOD testing on MT-Bench.The dataset is split into 8K/1K/1K training, validation, and test examples, with 20 responses generated per example.
- Evaluation setup: BEST-Route is compared with N-class, N-label, and clustering-based routing using armoRM response quality and inference cost.The baselines use fixed or model-selection routing strategies, while BEST-Route supports tunable cost-accuracy trade-offs.
- Routing performance: 60% cost reduction incurs only a 0.8% quality drop for BEST-Route, up to 4.28% better than all baselines.N-label routing instead incurs over a 5% performance drop at the same cost reduction, while fixed-routing baselines largely use the reference model.
- Adaptive sampling: BEST-Route achieves 20% cost reduction with a 0.21% quality drop and 40% reduction with a 0.47% drop when maximum sampling is n = 5.Best-of-n sampling improves routing performance, whereas fixed model-and-n pairings can yield lower-quality responses.
- Specialized-model routing: Adding Codestral-22b redirects coding queries from GPT-3.5-turbo and enables up to 20% cost reduction while exceeding GPT-4o performance.Under strict constraints without Codestral-22b, BEST-Route primarily selects Mistral-8x7b; higher budgets shift traffic toward larger models.
- Overhead and generalization: At n = 20, routing overhead is 18.7× faster than Llama-3.1-8b, while increasing n from 1 to 20 raises generation latency by 30%–59.3% across tested models.The measured overhead includes router prediction and best-of-n reward-model selection; API latency is not measured because of network and queuing delays.
6 Limitations
BEST-Route has limitations involving reliance on proxy reward model accuracy and scalability to extremely large model pools. Misalignment can affect response selection, while scaling to hundreds of LLMs remains unexplored.
- Proxy reward model misalignment may result in suboptimal response selection despite strong alignment with ground-truth evaluations.
- BEST-Route’s effectiveness with extremely large model pools, such as hundreds of LLMs, remains unexplored.
- Scaling the router to very large model pools may require additional optimizations.
7 Conclusion
BEST-Route is an adaptive LLM-routing framework that combines cost-efficient routing with test-time optimal compute through best-of-n sampling. Evaluations on real-world datasets show up to 60% cost reduction with less than 1% performance drop.
- BEST-Route combines adaptive model selection with best-of-n sampling to optimize inference costs while maintaining high response quality.
- 60% cost reduction with less than 1% performance drop was achieved on real-world datasets.
- The multi-head router supports fine-grained trade-offs between accuracy and efficiency.
A.1 Dataset
The evaluation dataset covers diverse tasks and supports training a generic routing framework. Baselines include router architectures, clustering-based routing, model cascades, and response-cost estimation procedures.
- Dataset: The dataset spans question answering, coding, and safety evaluation tasks.
- Dataset: The dataset contains 8K training, 1K validation, and 1K testing examples.
- Baselines: N-class Routing uses a BERT-based router to predict the best LLM for each query.
- Baselines: N-label Routing predicts all capable LLMs and selects the cheapest when multiple candidates exist.
- Baselines: Clustering-based Routing applies K-Means to TF-IDF query features and selects the most effective LLM for each cluster, using K = 50.
- Baselines: Model Cascades sequentially sample K responses per LLM and stop when the most consistent response exceeds a predefined threshold.
- Baselines: Model Cascades use exact match, BLEU, and ROUGE as agreement functions with K = 5.
- Cost estimation: BEST-Route estimates response costs from output-token counts and unit output-token prices, using training-split average output length per LLM.
B.2 Performance Results Compared to Model Cascades
Compared with model cascades, BEST-Route delivers higher-quality responses at lower cost. It achieves 60% cost reduction with a 0.8% quality drop, up to 6.46% better than cascading-based routers.
- 60% cost reduction with a 0.8% quality drop is achieved by BEST-Route compared with the reference model.
- BEST-Route delivers higher quality responses while achieving higher cost savings than all model-cascade baselines.
- BEST-Route is up to 6.46% better than cascading-based routers.
- Figure 8 reports routing performance results compared with model-cascade baselines.
C armoRM Score - A Case Study
armoRM aggregates 19 response-quality perspectives into a unified score. In the fever example, the higher-scoring response adds safety guidance about persistent high fever and seeking medical advice.
- armoRM aggregates 19 perspectives, including Helpfulness, Correctness, Coherence, and Verbosity, into a unified score ranging from -1 to 1.The metric is built from ground-truth ratings, normalization, de-biasing, and weighted summation.
- This example illustrates how a small armoRM score increase can reflect improvements across multiple quality dimensions.
- For the question about whether fever is normal during illness, both responses explain fever as a natural defense mechanism.