PaperlayerView exact revision on arXiv

Source-linked AI summary

Large Language Model Routing with Benchmark Datasets

Tal Shnitzer, Anthony Ou, Mírian Silva, Kate Soule, Yuekai Sun, Justin Solomon, Neil Thompson, Mikhail Yurochkin

arXiv:2309.15789v1cs.CLcs.LG

TL;DR

The paper addresses selecting the best LLM for a new task when no model dominates every use case and aggregate benchmark rankings lose task-specific information. It repurposes benchmark datasets to train binary correctness predictors and route inputs to candidate models, improving performance over always using one model while exposing out-of-distribution limitations.

  • Problem

    No single LLM is typically best across all tasks, while aggregate benchmark rankings do not identify the best model for a practitioner’s specific new task.

  • Method

    The paper repurposes benchmark results to train binary correctness predictors for candidate LLMs and uses their scores to route new tasks.

  • Results

    Routing improves overall performance over always using one model and can match a 70B model using smaller LLMs with 2–40 labeled samples from a new task.

  • Takeaways & Limitations

    Learning from benchmarks can select smaller, less expensive LLMs while retaining strong performance across tasks.

  • Takeaways & Limitations

    The major challenge is out-of-distribution generalization of correctness predictors across new tasks, although a few labeled task samples can reduce this gap.

Abstract

from arXiv · show

There is a rapidly growing number of open-source Large Language Models (LLMs) and benchmark datasets to compare them. While some models dominate these benchmarks, no single model typically achieves the best accuracy in all tasks and use cases. In this work, we address the challenge of selecting the best LLM out of a collection of models for new tasks. We propose a new formulation for the problem, in which benchmark datasets are repurposed to learn a "router" model for this LLM selection, and we show that this problem can be reduced to a collection of binary classification tasks. We demonstrate the utility and limitations of learning model routers from various benchmark datasets, where we consistently improve performance upon using any single model for all tasks.

1 Introduction

The paper targets selecting an LLM for a specific new task rather than relying on aggregate benchmark rankings. It repurposes benchmark results to learn binary correctness predictors that score candidate models for routing.

  • Motivation: Benchmark suites compare LLMs across diverse tasks and domains, but average rankings discard task-specific strengths.Practitioners typically need the best model for a particular use case, not the best average model across many datasets.
  • Benchmark landscape: Benchmark examples include HELM with 42 scenarios, MMLU with 57 tasks, and LM Evaluation Harness with over 200 tasks.
  • Approach: The proposed router trains one binary classifier per candidate LLM and uses their scores to recommend a model for a new task.
  • Problem: The paper studies how to identify the best LLM for a new task using benchmark datasets that reveal model performance across tasks and domains.

2 Related work

Prior work commonly ranks LLMs using aggregate benchmark performance or selects models after generating outputs from every candidate. This paper instead addresses task-specific and potentially out-of-distribution selection without requiring every candidate generation at decision time.

  • Benchmarking: LLM benchmarks typically average performance across tasks into a final ranking, discarding information about task-level variation.
  • Model selection: Classical model selection targets in-distribution test data, whereas LLM routing must select models for new tasks whose distributions may differ from benchmark data.
  • Routing LLMs: Prior LLM routing methods generally generate outputs from every candidate LLM before choosing the best model, which can be computationally prohibitive for large model pools.
  • Routing LLMs: The paper’s routing approach differs by predicting candidate performance from inputs rather than first obtaining generations from all models.

3 Learning from Benchmarks

The paper converts benchmark evaluations into supervised learning data for predicting each LLM’s correctness on inputs and routing new tasks to promising models. It emphasizes efficient inference and the challenge of generalizing predictors across task distributions.

  • Learning from Benchmarks: Benchmark inputs, references, and task-specific metrics are reused to learn from individual data-point performance instead of only producing aggregate model rankings.
  • Routing: At test time, routers use only the new task input to score candidate LLMs, avoiding generation by every model.
  • Correctness labels: For each candidate LLM, correctness is defined by whether its response exceeds a task- or metric-specific quality threshold.
  • Correctness predictors: The method trains one probabilistic binary classifier per LLM to estimate the probability that the model will be correct on a new input.
  • Limitation: A central limitation is out-of-distribution generalization because predictors must estimate performance on unseen tasks whose distributions may differ from training benchmarks.
  • Predictor implementation: The study uses sentence-transformer embeddings with k-nearest neighbors to fit potentially complex correctness boundaries without extensive hyperparameter tuning.

4 LLM routing with (imperfect) correctness predictors

The paper routes a new task to an LLM by predicting each candidate’s correctness and combining predicted correctness with an estimate of out-of-distribution reliability. Its adaptive score accounts for imperfect predictors and is theoretically expected to better approximate oracle routing than a non-adaptive score.

  • Routing objective: LLM routing identifies the candidate expected to be correct most often on a new task from that task’s inputs.The routing objective is to select the LLM with the highest correctness frequency on the unseen task.
  • OOD confidence model: Because calibration is difficult out of distribution, the method models predictor accuracy p(d′, m) as the probability that a correctness prediction is correct on task d′.The model treats predictor outputs as correct with probability p(d′, m) and incorrect otherwise.
  • OOD confidence model: The method estimates p(d′, m) from task descriptors using a non-parametric regression model, specifically a one-dimensional Gaussian kernel smoother.Training uses task-level predictor accuracies and distances from each task to the others; the same descriptor is computed for a new task.
  • Adaptive routing: The final score S3 combines the base score S2 with estimated predictor accuracy and can default to the best model on average when confidence is insufficient.The alternative selection rule uses a probability threshold η and otherwise chooses the benchmark-average best model.
  • Connection to meta-learning: Adaptive shrinkage fits a larger model class, so under standard loss assumptions its population risk is no greater than that of non-adaptive routing.The authors expect S3 to outperform S2, while noting that finite-sample variance could offset this advantage; empirical results generally improve with S3.

5 Experiments

Experiments on HELM and MixInstruct show that benchmark-trained routers can improve model selection while reducing inference costs, though performance depends on distribution shift and benchmark coverage.

  • HELM model routing: Across 29 HELM leave-one-task-out experiments, benchmark-trained routers are compared with the best model on average, oracle routing, and log-likelihood scoring.The evaluation reports selected-model accuracy, ratio to the best model, rank correlations, and model-call or parameter costs.
  • HELM model routing: Accounting for correctness-predictor imperfections enables S3 to outperform the best-model-on-average baseline while selecting smaller models where appropriate.The correctness predictors have average accuracy 0.59, while the practical S3 score estimates their accuracy with MAE=0.116.
  • HELM model routing: Log-likelihood also performs well, but requires passing each test input through every candidate LLM, whose total size is 347B parameters.This creates a substantially more expensive inference-time routing procedure than the proposed approach.
  • Reducing the OOD gap: Adding a small number of in-distribution samples reduces the out-of-distribution gap and improves all three routing scores.With α = 0.05, correctness-predictor accuracy reaches 0.65; further in-distribution data causes S1 to outperform S3, although variance remains large.
  • MixInstruct model routing: On MixInstruct, the approach uses one model generation and one embedding call per instance, whereas compared methods generate with every examined LLM.The method does not consistently outperform all comparators, but demonstrates improved inference-time efficiency.
  • MixInstruct model routing: On MixInstruct subsets, routing approaches oracle metrics as the average distance to reference points decreases, suggesting denser benchmark coverage may improve selection.The subsets are formed by limiting each test point’s average distance to its closest reference-set points.

6 Discussion and Conclusion

Smaller-model routing can reduce inference costs while matching a much larger model, but out-of-distribution generalization remains the central challenge. More benchmark coverage and limited task labels are presented as paths toward more reliable routing.

  • How useful are smaller LLMs?: Routing smaller models can reduce costs, compute, and latency when a smaller model is adequate for the task.The paper connects this capability to broader use of LLMs across domains.
  • How useful are smaller LLMs?: Smaller-model routing can match Llama 2 70B using benchmark learning plus α = 0.04, or 2-40 labeled samples from a new task.The result targets lower inference costs while preserving performance.
  • Learning from more benchmarks: For smaller-distance new tasks, S3 correlation with candidate-model accuracy approaches 1, whereas correlation sometimes deteriorates at larger distances.The analysis supports learning routers from benchmarks that make new tasks closer to the available data.
  • Future work: The major challenge is out-of-distribution generalization of correctness predictors.The paper identifies more benchmarks, OOD methods, and a few task labels as potential ways to improve routing.
  • Future work: Routing expert LLMs fine-tuned for specific domains may simplify prediction when a new task belongs to that domain.The paper expects regions where such models are correct to align with their areas of expertise.

A Correctness predictors and confidence estimation

The method embeds benchmark inputs and learns per-model correctness predictors with nearest neighbors. It then estimates how reliably those predictors transfer to a new task using dataset distance and kernel smoothing for confidence-aware routing.

  • Correctness predictors: Benchmark inputs are embedded with a sentence transformer, and each candidate model receives a binary kNN correctness predictor.The predictors estimate whether a model is correct on a new embedded input from nearby benchmark examples.
  • Correctness predictors: The routing method computes per-model scores S1 and S2 from the correctness predictors for a new task.These scores are defined using the predictor outputs described in the preceding equations.
  • Confidence estimation: The confidence model combines a one-sided Chamfer-style dataset distance with a Gaussian kernel smoother estimating predictor accuracy p(d′, m).The estimated accuracy is used in the S3 score for confidence-aware model selection.
  • Confidence estimation: Kernel-smoother training pairs are generated by holding out benchmark datasets and measuring predictor accuracy after training on the remaining datasets.In HELM, the resulting training set contains Z = 4200 distance-accuracy pairs.
  • Confidence estimation: The confidence model can be combined with any classifier, although the experiments use kNN correctness predictors.The experiments set the kernel bandwidth to σ = 0.09.

B.1 Reducing the OOD gap

Adding a small amount of in-distribution data reduces the OOD gap and improves routing, while further additions can make S1 outperform S3 as predictor calibration improves.

  • Reducing the OOD gap: Increasing α improves Pearson correlation similarly to the accuracy trends, while more new-task samples improve correctness-predictor accuracy and reduce the OOD gap.The additional results are summarized in Figure 6.
  • Reducing the OOD gap: Kernel-smoothing MAE does not improve as much with increasing α, so S3 eventually becomes worse than S1 in correlation and selected-model accuracy.This identifies a limit of the confidence estimate as more in-distribution data are added.

B.2 Dataset distance and Pearson correlation

The dataset distance decreases as new-task inputs are mixed into the benchmark data, enabling analysis of routing on tasks closer to the training distribution. S3 routing quality improves as this distance becomes smaller.

  • Dataset distance: Increasing α usually decreases u(d′) because new-task inputs are moved into the benchmark set.This assumes new-task inputs are, on average, closer to one another than to inputs from other tasks.
  • Dataset distance and Pearson correlation: Figure 8 relates dataset distance u(d′) to Pearson correlations between model scores and candidate-LLM accuracies.The figure presents these relationships across routing scores, with S3 results shown separately in Figure 5.

C Additional details for model routing on MixInstruct

For MixInstruct, the routing procedure scores each test point separately and selects a model per point. The experiments use BERTScore-based nearest training examples and k=10 in the kNN classifier.

  • Correctness predictor and metrics: Each MixInstruct test point receives a per-point score for every model, and the router selects a model based on that score.Reported metrics are averaged over the output evaluations from these per-point selections.
  • Correctness predictor and metrics: BERTScore on the closest training points is used to compute the model-quality estimate g_m(xᵢ).The choice was motivated by the conceptual relation between the approach and BERTScore's use of embedding-space distances and was empirically validated.
  • kNN parameter: k=10 is used for the kNN classifier because MixInstruct test examples are in-distribution with respect to datasets represented in training.Metrics did not significantly vary for k∈[5, 100].

D Proof of Lemma 4.1

The appendix presents the proof of Lemma 4.1 by establishing an upper bound for one loss term and indicating that the other follows similarly. The supplied passages also identify the lemma's subadditive-loss setting.

  • Lemma statement: Lemma D.1 considers losses of the form ℓ(y₁, y₂)=ρ(y₁−y₂), where ρ is subadditive.The passage gives ρ(x)=1 as an example.
  • Proof: The proof begins by establishing an upper bound for ℓ(S₂, eS).The displayed derivation itself is not included in the supplied passage.
  • Proof: The upper bound for ℓ(S₃, eS) is obtained through a similar argument.The proof passage notes that subadditive functions are convex in one step of the derivation.
  • Supporting tables: The supplied appendix also labels tables describing HELM datasets and candidate LLMs.No table contents are provided in the cited passages.
Loading 2309.15789v1…