Large Language Model Routing with Benchmark Datasets

Venue: colmweb_org_COLM_2024_Conference
Published: 2024-03-22 23:47:56

Abstract: The number of open-source Large Language Models (LLMs) grows daily, as does the number of available benchmark datasets used to evaluate LLMs. While some models dominate these benchmarks, no single model achieves the best accuracy in all tasks and use cases. In light of this observation, we address the challenge of selecting the best LLM from a collection of pre-trained models, given a new task.  While related work relies on evaluating each candidate model on a set of labeled examples, our new formulation does not assume any labeled data from the new task is available.  Instead, we repurpose a collection of benchmark datasets---which may focus on different tasks than the one at hand---to learn a ''router'' model for  LLM selection from inputs only; this problem reduces to a collection of binary classification tasks. Empirically, our strategy consistently improves performance over using any single model for all tasks.

Main paper: 
 introduction: 
 Figure 1 : We learn the strengths of candidate LLMs (T5, Falcon, Llama) on various tasks (emojis inside boxes: QA, reasoning, summarization, etc.) and domains (4 sections within each box: finance, legal, general knowledge, etc.) from benchmark datasets. We accomplish this by training a binary classifier per LLM (decision boundaries marked with colors in the upper part of the figure).For a new task (paper stack), we score each LLM with these binary classifiers and recommend an LLM (here Falcon) to the user.Large Language Models (LLMs) have demonstrated groundbreaking abilities to solve diverse tasks across a variety of NLP domains (Devlin et al., 2018; Brown et al., 2020) . Today, researchers in both academia and industry are releasing new LLMs daily.1 These models perform tasks ranging from text classification to question-answering, summarization, and dialogue.The popularity and influx of open-source LLMs and the diversity of their potential use cases made it crucial to develop comprehensive benchmarks, i.e., collections of datasets representing different tasks and domains to compare LLMs. For example, HELM (Liang et al., 2022) consists of 42 scenarios covering a variety of uses, MMLU (Hendrycks et al., 2020 ) is a multiple-choice question answering benchmark with 57 tasks organized by topics, Open LLM Leaderboard (Beeching et al., 2023) combines MMLU with other question-answering datasets, and LM Evaluation Harness (Gao et al., 2021) supports over 200 tasks. While there always will be an LLM that is the best on average across benchmarks, there is unlikely to ever be a model that is strictly the best on each of the hundreds of datasets comprising various benchmarks. Meanwhile, a practitioner typically wants to know what is the best model for their specific use case and is less concerned about average performance on a plethora of other datasets.In this paper, we study the problem of identifying the best LLM for a new task. To learn about the strengths and weaknesses of candidate LLMs we use benchmark datasets that give insights into the performance of LLMs across tasks and domains. For example, suppose the new task is answering math questions. In that case, it is more intuitive to consider models that do well on other STEM question-answering datasets and discount performance on, e.g., sociology or toxicity detection. We make this idea precise by casting the learning of model strengths as a binary supervised learning task, where the features are input embeddings of samples across tasks and the labels are whether the model "did well" on the corresponding inputs, e.g., generated correct class label, answered a question correctly, or followed input instructions sufficiently well. See Figure 1 for an illustration. Such information is collected during benchmark evaluations and can be reused for training model routers without having to run expensive LLM inference again. The resulting router is also efficient at test time, as it only requires calling the chosen LLM.Our contributions are summarized below:• We formalize the problem of learning the strengths and weaknesses of LLMs for downstream routing, i.e., selecting the best model, as a collection of binary classification problems. The goal of each classification problem is to predict whether a given LLM will be "correct" on an input. • We propose three scores for selecting LLMs for a new task using these correctness predictors. Our third score is designed to account for mistakes a correctness predictor can make on the (out-ofdistribution) data from a new task that is likely to be different from the datasets in benchmarks used for training the correctness predictors. We establish connections to meta-learning to obtain theoretical insights into the efficacy of these scores. • We verify the efficiency of our model routing scores empirically on 29 datasets from HELM (Liang et al., 2022) representing scenarios like question answering, text classification, knowledge, and reasoning, and MixInstruct (Jiang et al., 2023) , a collection of datasets for evaluating instruction following capabilities of LLMs. • We discuss and empirically investigate questions concerning the efficacy and utility of learning LLM routers from benchmarks: generalization of correctness predictors to new tasks, the importance of a larger pool of benchmarks, and the potential of routing smaller LLMs to reduce costs. 
related work: 
 Benchmarking Comparing models or algorithms across various tasks is a standard practice in ML and AI literature. Prior to Foundation Models (Bommasani et al., 2021) , it was typical to apply the same learning algorithm to train a model on each of the datasets and compare the performance against other learning algorithms. The UCI Machine Learning Repository (Kelly et al., 2023) is one prominent example of such a collection of datasets often used to compare learning algorithms. With the emergence of Foundation Models, i.e., models with billions of parameters trained on massive datasets using large compute clusters, the paradigm changed to evaluating the same model (or a fewshot tuned version of it) on a variety of tasks (Bojar et al., 2014; Goyal et al., 2019; Li et al., 2022) .In the context of Large Language Models, many benchmarks (Wang et al., 2018; 2019; Hendrycks et al., 2020; Gao et al., 2021; Srivastava et al., 2022; Liang et al., 2022; Beeching et al., 2023; Jiang et al., 2023) were proposed to help determine the most capable LLM. Benchmarks typically average the performance of models across tasks and provide a final ranking, discarding the rest of the information. In this work, we use the byproducts of benchmark evaluations, i.e., the per-sample performance of various LLMs across tasks, to learn about their individual strengths and identify the best LLM for a new task.Model selection Selecting the best model, or model selection, is a classical topic in statistics and ML (Bishop & Nasrabadi, 2006; Hastie et al., 2009; Raschka, 2018) . However, the typical problem setting is quite different: classical methods like cross-validation aim to estimate the population error of a model trained on samples from the population distribution. In other words, the goal is to find the best model for in-distribution test data, i.e., data sampled from the same distribution as the train data. The notion of "train" data is quite elusive for LLMs, as they are usually trained on massive datasets with trillions of tokens with a simple task of next token prediction (Radford et al., 2019; Brown et al., 2020) . However, the tasks we evaluate them on are often more structured, e.g., classification and question-answering, and are specific to domains that may or may not be sufficiently represented in the train data. In addition, techniques like k-fold cross-validation require training the model multiple times, which is infeasible for LLMs.Out-of-distribution model selection Recognizing the limitations of the model selection methods for in-distribution test data (Gulrajani & Lopez-Paz, 2021; Koh et al., 2021) , recent work has proposed a variety of methods to select models when deployed on data that may differ from the train data. These methods rely on ideas such as bootstrapping (Xu & Tibshirani, 2022) , reweighing (Chen et al., 2021b; Maity et al., 2023) , agreement of models or ensembles (Jiang et al., 2021; Chen et al., 2021a; Ng et al., 2023) , or aligning model accuracy in-distribution with a confidence threshold (Guillory et al., 2021; Garg et al., 2022; Yu et al., 2022) . Most of these methods are nontrivial to extend to generation use-cases of LLMs; some require training multiple models, and some need well-defined in-distribution data related to the new task. (Fu et al., 2023; Wang et al., 2023; Liu et al., 2023 ) may be used instead of such scoring models, but such an approach would require producing generations with every candidate LLM and evaluating each of these generations with a typically larger/commercial LLM. In this paper, instead, we use data from benchmarks to learn the strengths and weaknesses of LLMs across tasks and domains. The resulting model router requires generating outputs only with the chosen LLM at test time, making the proposed method a lot more cost-efficient. 
learning from benchmarks: 
 We start by introducing notation to describe the majority of NLP benchmarks. Let {x d 1 , . . . , x d n d } D d=1 be a collection of inputs across D tasks. Each input text x d i corresponds to a reference answer r d i , i.e., an ideal generation for the corresponding input. Finally, there is a metric F d (x, o, r) that can be task-dependent and measures how well a response o for an input x corresponds to the reference r.To test an LLM m , m ∈ {1, . . . , M }, on the benchmark, for each task d = 1, . . . , D, its responses are generated {o d im = LLM m (x d i )} n d i=1 and compared to the corresponding references to obtain performance metrics{f d im = F d (x d i , o d im , r d i )} n d i=1 .2 At this point, the majority of the benchmark studies will take a (weighted) average of the performance metrics and report a single score for every LLM to rank them in performance. Instead, we reuse these evaluation results to formulate a supervised learning problem to better understand the strengths and weaknesses of various LLMs based on their performance on data points and tasks.Supervised learning from benchmarks Our goal is to learn a simple routing function g m (x) for each LLM, m = 1, . . . , M , that can predict{f d ′ im } n d ′i=1 , i.e., the performance of the corresponding LLM on a new task d ′ . Then it is trivial to select the best LLM for this task. For efficiency at test time, we restrict the routers {g m } M m=1 to only depend on the input x. This is in contrast to the majority of prior works on LLM routing that first obtain generations with every candidate LLM and then use them to choose the best model (Liu & Liu, 2021; Ravaut et al., 2022; Jiang et al., 2023) .With thousands of open-source LLMs, it is simply infeasible to obtain generations with every LLM for every input at test time.To complete the problem formulation, we denote the "correctness" of model m on an input x by y(x, m) ∈ {0, 1}. Correctness is evaluated as follows: generate a response o d im with LLM m on input x d i , compare it to the corresponding reference r d i , and output 1 if the model's response is good enough, i.e., f d im > η d , and 0 otherwise, where η d is some threshold that can be task and/or metric specific. For tasks like classification or multiple-choice QA, y(x d i , m) = f d im , while for various evaluation metrics used in summarization and instruction following tasks (Zhang et al., 2020; Sellam et al., 2020; Yuan et al., 2021) , the notion of correctness can help to account for the heterogeneity of popular metrics and task difficulty levels. In Section 5.2, we also present results with raw metrics instead of correctness.To train a predictor of an LLM correctness, for each LLM, m = 1, . . . , M , we solve the following optimization problem:EQUATIONwhere we choose ℓ to be a binary cross-entropy loss and g m is any standard probabilistic classifier, i.e., g m (x) estimates P (y(x, m) = 1|x In this work, we use a simple model for the correctness predictor: we embed all inputs with a sentence transformer (Reimers & Gurevych, 2019) and use a k-nearest neighbors classifier (Cover & Hart, 1967) as {g m } M m=1 . kNN is a simple non-parametric classifier that allows us to fit a potentially complicated decision boundary of an LLM's correctness across multiple tasks without extensive hyperparameter tuning. We choose this approach for learning correctness predictors to emphasize the utility of learning from benchmarks even with a basic method and instead focus on the question specific to our problem that has not been studied in prior works on OOD generalization: Can we improve the quality of LLM routing with an imperfect correctness predictor? 
llm routing with (imperfect) correctness predictors: 
 The goal of LLM routing is to identify an LLM that will have the highest frequency of being correct on a new task d ′ , given the inputs {xEQUATIONHere, S(m, d ′ ) is the "oracle" score that we want to estimate. The most intuitive estimator is simply using the correctness predictorEQUATIONbut prior work has shown that accurately estimating P (y(x, m) = 1|x), i.e., training calibrated correctness predictors g m s, is challenging on OOD data (Ovadia et al., 2019) . Meanwhile, g m may still produce accurate predictions after thresholding the predicted probability even if the class probabilities are not estimated well, which is often the case with neural networks (Guo et al., 2017) . This motivates another score:EQUATIONwhere t ∈ (0, 1) is some threshold, e.g., t = 0.5, I is an indicator function, and ḡm (x) ∈ {0, 1} can be interpreted as the prediction of g m on x. This score, however, does not take into account the potential "imperfection" of g m , i.e., g m , as a binary classifier, is likely to have lower accuracy on OOD data from task d ′ . To address this issue, we model the out-of-distribution confidence of the predictions ḡm .A simple OOD confidence model We model LLM correctness as follows:EQUATIONi.e., p(d ′ , m) ∈ [0, 1] is the probability that ḡm is the correct prediction on a data point from task d ′ .The above model can be condensed as follows:EQUATIONIn this simplistic (and approximate) model, we assume that p(d ′ , m) does not depend on the input x after conditioning on the task d ′ . The assumption is analogous to the homoscedastic error term assumption in linear regression models and allows us to interpret p(d ′ , m) as the marginal/overall accuracy of ḡm on data from the task d ′ .Prior work has studied the problem of estimating OOD accuracy given the inputs from a new task, but existing methods are challenging to combine with our approach. For example, Garg et al. (2022) learn a threshold on model confidence, which is hard to apply when using kNN classifiers, and Ng et al. ( 2023) require data augmentations that can be challenging to identify given the diversity of tasks in benchmarks. Prior methods also do not take into account the partition of the train data into tasks inherent in our problem setting.We treat the problem of estimating p(d i=1 and predict p(d ′ , m) using the fitted regression model. In general, one can consider more sophisticated, higher-dimensional task descriptors u(d), but here, for simplicity, we keep it 1-dimensional and use a Gaussian kernel smoother (also known as the Nadaraya-Watson estimator) as the non-parametric regressor. We provide details in Appendix A.Finally, given the model of LLM correctness 6, S(m, d ′ ) is a random variable (corresponding to S(m, d ′ )) distributed as a (scaled) sum of two Bernoulli random variables. To arrive at our final score for LLM routing, we take its expected value:EQUATIONWhen selecting an LLM with S 3 , we consider an alternative to the arg max criterion based on our correctness model 6, which defaults to the best model on average across benchmark datasets when we are not sufficiently confident that a candidate model will be better:EQUATION)where m 3 = arg max m S 3 (m, d ′ ), i.e., the best LLM for the new task according to S 3 , and m * = arg max m D d=1 S(m, d), i.e., the best LLM across the benchmark datasets. In the experiments, we set η = 0.6.We summarize our LLM routing procedures in Appendix A.The OOD confidence model in equation 6 is a meta-model of routing across multiple tasks, and fitting it entails a form of meta-learning. Consider the meta-learning problemEQUATION) where ḡm and p(•, m) are meta-parameters and adaptation step ḡm → ḡm (•)p(•, m) adaptively shrinks the router output towards ambiguity. We exploit this connection to theoretically demonstrate the potential advantages of routing LLMs using S 3 over S 2 .In expectation/in the population, equation 9 fits a larger model class than equation 1, so the risk of the adaptively shrunken router is at most that of the non-adaptive router:D d=1 E ℓ(ḡ m (X d )p(d, m) + (1 -ḡm (X d ))(1 -p(d, m)), y(X d , m)) ≤ D d=1 E ℓ(ḡ m (X d ), y(X d , m)) .(10) This suggests (subject to standard assumptions on the loss function) that adaptive shrinkage routing leads to better approximations of the oracle router. Lemma 4.1 confirms this intuition. Lemma 4.1. Let ℓ(y 1 , y 2 ) = ρ(y 1 -y 2 ) for some subadditive ρ : R → R (e.g. ρ(x) = 1 2 x 2 for the square loss). We have We present the proof in Appendix D. Combining equation 10 and Lemma 4.1, we expect the adaptive router based on S 3 to outperform its non-adaptive counterpart based on S 2 . That said, it is unclear whether adaptive shrinkage will improve the performance of the adaptive router in finite samples: the expected performance of the adaptive router may be offset by the inflation in variance from fitting the larger (adaptive) model class. Fortunately, our empirical results show that task-specific adaption, i.e., using S 3 as a score for routing, generally improves performance. The two-step method for fitting ḡm and p in Section 4 approximately minimizes equation 9 with a single Gauss-Seidel pass through the decision variables.ℓ(S 2 , S) ≤ E ℓ(ḡ m (X d ), y(X d , m)) ), ℓ(S 3 , S) ≤ E ℓ(p(d, m)ḡ m (X d ) + (1 -p(d, m))(1 -ḡm (X d )), y(X d , m)) ). 
model routing on helm: 
 Data. We select 29 datasets from the HELM benchmark (Liang et al., 2022) representing scenarios such as question answering, text classification, language, knowledge, and reasoning, among others. We present additional information about these datasets in Table 4 .We evaluate 18 open-source models ranging in size from 3B to 70B, including base and chat variations of Llama 2 in different sizes. All models are summarized in Table 5 .Model routing. The best model on average (BMA) across the 29 considered HELM datasets is llama-2-70b (followed by llama-2-70b-chat). Our goal is to show that learning model routers from benchmark data can simultaneously outperform BMA and reduce inference costs by recommending smaller LLMs for tasks where they can perform well. We compare models selected with the three scores, S 1 , S 2 , and S 3 , presented in Section 4 to the performance of llama-2-70b, i.e., the BMA. All correctness predictors g m s are kNN classifiers with k = 5. We also consider Mixture of Experts (MoE) (Masoudnia & Ebrahimpour, 2014) as a baseline where an LLM is selected for each input individually based on the correctness frequency of candidate LLMs on the same k neighbors of this input. Results for neural networks as correctness predictors are in Appendix B.1.We also report the performance of the best model according to the "oracle" score S, which is the upper bound on what can be achieved with model routing, and S3 , which corresponds to S 3 with the true p(d ′ , m), i.e., the accuracy of (an imperfect) g m on d ′ . Finally, we compare the scoring LLMs with the average log-likelihood (LL) (or negative perplexity) of the response they generate on the inputs from the task of interest. This last baseline requires producing generations with every LLM at test time to make a selection, while all of our scores only require generating with the chosen LLM.Results. We conduct 29 sets of experiments, each time selecting 28 of the datasets as the benchmark data for training the LLM routers and using the remaining task as the new task d ′ for evaluating the quality of the LLM selection for this task. In Table 1 we report averages across experiments for the performance of the selected model (Acc.), ratio of this performance to the performance of the best model for the corresponding new task (Ratio to Best), Pearson and Spearman rank correlations between model accuracies and model scores, number of parameters of the selected model (# Params), rank of the selected model out of 18 considered (Rank). We also report the fraction of times the BMA is selected by a method (% BMA). Best results are highlighted with bold and second best with an underline (excluding Oracle).First, we notice that accounting for imperfections of the correctness predictors (their average accuracy is 0.59) has clear benefits: when we have access to the true accuracy of correctness predictors, the corresponding score, S 3 true p, noticeably outperforms all other scores. Our simple kernel smoothing estimator of this accuracy (MAE= 0.116) allows us to obtain a practical model routing score S 3 that outperforms BMA (llama-2-70b) while choosing smaller models for some of the tasks (as evident by the average number of parameters of the chosen models). S 2 sacrifices some accuracy but chooses even smaller performant models. Overall, learning from benchmarks allows us to obtain LLM routers that can improve overall performance while utilizing smaller models where appropriate. Finally, we note that log-likelihood (LL) also performs well. However, routing with it requires passing each test input through each candidate LLM, which has 347B parameters in total. Reducing the OOD gap. The average accuracy of correctness predictors across tasks and models for the experiments in Table 1 is 0.59. It is a fairly low accuracy for binary classification, which we attribute to the diversity of tasks in HELM leading to substantial distribution shifts when predicting the correctness of LLMs on held-out tasks. We investigate the quality of model routing when we reduce this OOD gap. A simple strategy to reduce this gap is to collect a small number of labeled in-distribution samples. This can be accomplished by asking a practitioner to provide reference answers (r d ′ i s) for a small number of inputs from their task to evaluate the correctness of candidate LLMs on these in-distribution inputs and use it to improve correctness predictors.We simulate this scenario by moving min(αn d ′ , 50) samples from the data from a new task d ′ to the data for training the correctness predictors. The upper limit of 50 samples is to maintain practical utility while accounting for varying dataset sizes (see Table 4 ). In this experiment, we additionally compare to Few-Shot, which selects an LLM for the new task based on the performance only on these labeled samples from the new task. We conduct 29 sets of experiments, repeating each one 10 times to obtain standard deviations (randomness is due to random selection of data points from a new task for reducing the OOD gap). We summarize the average accuracy of models selected with various routing scores for varying α in Figure 2 (α = 0 corresponds to Table 1 ). Results for Pearson correlation are in Figure 6(a) .We see that even a small number of in-distribution samples (α = 0.05) can reduce the OOD gap (corresponding average accuracy of correctness predictors is 0.65; see Figure 6(b) ) and noticeably improves the model routing performance of all three of our scores. When the number of in-distribution samples further increases, S 1 starts to outperform S 3 . We attribute this observation to kNN being well-calibrated in-distribution, i.e., the correctness predictors provide reliable estimates of their own confidence P (y(x, m) = 1|x), which are used by S 1 in equation 3. Finally, we note a fairly large variance in the results due to random selection of the in-distribution training samples from d ′ , suggesting that active learning (Settles, 2009) can help to improve LLM routing further.We further demonstrate our approach in a different setting and task type, on the MixInstruct benchmark dataset (Jiang et al., 2023) . The dataset is composed of instruction-following tasks, divided into train/validation/test sets of 100K/5K/5K samples, and includes evaluations of N = 11 opensource LLMs using common metrics, e.g. BERTScore (Zhang et al., 2020 ), BARTScore (Yuan et al., 2021) , and BLEURT (Sellam et al., 2020) . In Jiang et al. (2023) , this benchmark was used to compare different LLM ranking methods in per-instance model selection. We follow the same setting and apply our score S 1 (m, d ′ ) to the test set, per-instance, where we use the 100K-sample train set as the benchmark data for training our LLM router. See Appendices A and C for details on the score computation and the experiment parameters, respectively. Due to the per-instance setting and the test set constructed from in-distribution data, we focus on our simplest router model S 1 , equation 3.We compare our approach with the scoring methods examined by Jiang et al. (2023) , as well as scoring based on the average log-likelihood (LL) of the model responses to the inputs. Additionally, We report the results of BERTScore, BARTScore and BLEURT in Table 2 , along with the number of model calls per instance (MCPI), for N LLMs, performed during inference time. All compared methods require model generations for every point in the test set, by each of the examined LLMs, whereas our approach requires only one model generation and one call to some general embedding function. In addition, all methods, except for LL, require training auxiliary language models, whereas our approach is a simple kNN classifier on the embedded inputs. While our approach does not consistently outperform the compared methods, these results demonstrate the potential of using benchmark datasets for model routing with significantly better inference-time efficiency.Effect of benchmark dataset sparsity. To highlight the potential of our approach in this setting, we examine the effect of the reference benchmark data sparsity. We apply our method to different subsets of the test set, X test , where the subsets are defined by limiting the maximal average distance of each test set point to the closest points from the reference (train) set, denoted by NN train , i.e.X ′ C = x ′ ∈ X test 1 |NNtrain(x ′ )| x∈NNtrain(x ′ ) dist(x ′ , x) < C, where C is the maximal average distance and X ′ C is the resulting subset of the test set. Figure 3 presents the metric scores for the different subsets using our method, the oracle (best possible choices), and LL scoring. We also report the percentage of the test set that is used in each subset. This figure depicts that our predictor approaches the oracle metrics as the average distance to the reference points decreases. This suggests that adding more benchmark datasets to reduce the sparsity of the reference space may lead to better LLM selections with our approach. 
discussion and conclusion: 
 How useful are smaller LLMs? While a given LLM may work best on average, these models tend to be the biggest and therefore most expensive to run. Practitioners can achieve gains in cost, compute, and latency if we can successfully predict whether a smaller LLM can be adequate for a given task. Identifying good smaller models for tasks of interest will also redefine the cost/benefit tradeoff behind automating certain tasks, potentially incentivizing the automation of new tasks that were previously cost-prohibitive to automate with larger LLMs. To evaluate the potential of smaller LLMs we revisit our HELM experiment in Figure 2 . In Figure 4 , we perform LLM routing using only models with ≤ 13B parameters and compare it to the performance of Llama 2 70B. Oracle's performance demonstrates that it is conceptually possible to outperform a large model by routing smaller LLMs. Results with our scores S 1 and S 2 (see Figure 7 for breakdown by scores) demonstrate that it is also practically feasible to match the performance of the 70B model by combining learning from benchmarks with a small number (α = 0.04, i.e., 2-40 samples) of labeled samples from a new task that a practitioner can provide to save on the inference costs in their LLM application. Pearson Corr(S3 eq. ( 7), Accs.) Learning from more benchmarks. We anticipate learning LLM routers from benchmarks to be the most effective when new tasks are similar to the benchmark tasks, thus reducing the OOD gap without any labeling burden for a practitioner. To empirically investigate this hypothesis, in Figure 5 we visualize the relation between the quality of model routing with S 3 , measured with Pearson correlation between model scores and accuracies of candidate LLMs, and the distance u(d ′ ) from a new task d ′ to the available benchmark data for training the routers. In this experiment, we aggregate results across different α values from Figure 2 . For smaller distance values the correlation is approaching 1, while for large distances it sometimes deteriorates. Results for other scores demonstrate a similar trend and are presented in Appendix B.3 along with additional details. This experiment and the benchmark dataset sparsity analysis presented in Figure 3 for MixInstruct illustrate that learning with more benchmarks can improve the efficacy and reliability of the LLM routers as new tasks are more likely to be closer to a large collection of datasets.Future work. Our work demonstrates the potential of learning from benchmarks for LLM routing and investigates 3 model scores in the context of OOD generalization when routing LLMs for new tasks. We summarize potential next steps for improving the quality and efficacy of LLM routers.The major challenge of LLM routing is OOD generalization of correctness predictors. Thus, using more benchmarks and modern methods for improving OOD generalization to learn correctness predictors is a promising next step. A practitioner can also provide labels for a few samples from their task, possibly guided by active learning techniques, to adapt or fine-tune correctness predictors. Even when reducing the OOD gap is too challenging, our score accounting for the (potentially low) accuracy of correctness predictors demonstrated strong results when this accuracy, p(d ′ , m), is known for a new task, thus encouraging the development of methods for estimating it better.We also anticipate that routing "expert" LLMs fine-tuned for a specific domain can improve the results. Regions of the sample space where such models are "correct" should mostly align with the domains of their expertise (recall Figure 1 ), making it easier to learn the corresponding correctness predictors and simplifying LLM routing when a new task is from a specific domain.Our experiments in Figure 4 demonstrate the utility of LLM routing with smaller models, which can reduce costs and facilitate using LLMs in a broader set of domains. Thus, we want to explore modifications to our scores that will encourage the selection of smaller LLMs when their anticipated performance is comparable to the larger, more reliable models. Prior work on frugal API selection (Chen et al., 2020; 2023) provides a good starting point to explore this direction.Finally, it is interesting to explore other approaches to correctness prediction. For example, Mixture of Experts (Masoudnia & Ebrahimpour, 2014) whereK(u(d ′ ), u z ) = exp -(u(d ′ )-uz) 2 2σ 2. We set σ = 0.09 in all experiments, which is the value we found to perform well through some preliminary experimentation.Finally, we note that the proposed confidence model, including the definitions of the dataset distance and kernel smoother, can be combined with any classifier g m , and is not restricted to the kNN classifier used for the correctness predictor in our experiments.Additional notes regarding S 3 Recall that when selecting a model with S 3 (m, d ′ ) we use an additional step described in equation 8 that facilitates the selection of the best model on average when we are not sufficiently confident in the model with the highest S 3 (m, d ′ ) score. Probability expression, P ( S(m 3 , d ′ ) > S(m * , d ′ )), required for this step is not available in closed form, as S is distributed as a (scaled) sum of two Bernoulli random variables, but it is straightforward to estimate via Monte Carlo sampling from the corresponding Bernoulli distributions.When reporting correlations for S 3 (e.g., Pearson and Spearman correlations in Table 1 ), we use S 3 (m, d ′ ) as is, i.e., as defined in equation 7. 
 
 Appendix: We present LLM routing with more sophisticated correctness predictors. In this section, we use a small Multi-Layer Perceptron (MLP) as g m to classify the embedded inputs. Each MLP comprises three fully connected layers (two hidden layers, 1500 units each) with ReLU activation and sigmoid output. These MLPs were trained with the Adam optimizer (Kingma & Ba, 2014) with learning rate 0.01 to minimize the binary cross entropy loss described in equation 1.Training occurred over 100 epochs with early stopping should the validation accuracy not increase for 10 consecutive epochs. We also explore the combination of S 3 with ATC (instead of the kernel smoother), a prior OOD accuracy estimator proposed by Garg et al. (2022) .Results We conducted the model routing experiment analogous to the experiment in Table 1 of the main paper. Results are reported in Table 3 . We note several differences compared to results with the kNN classifier in Table 1 . First is the clear improvement in the S 1 score; the MLP classifier performs better on OOD data. As before S 3 true p outperforms the other scores indicating the potential value of using the proposed confidence model 6. However, in this experiment, S 3 with the kernel smoother no longer improves upon S 1 . We also report results with ATC as the estimator of the accuracy of correctness predictors, which performs noticeably worse than our kernel smoothing estimator. The corresponding MAE is 0.118 for the kernel smoother and 0.177 for the ATC, demonstrating the advantage of our estimator in this application. Finally, across all scores the selected model sizes have considerably increased; this suggests that the MLP classifier underestimates the performance of smaller models compared to kNN. (c) S3 equation 8Figure 7 : LLM routing with ≤ 13B parameter models compared to Llama 2 70B.We present additional results for this experiment in Figure 6 . (a) shows Pearson correlation improvement as we increase α, similar to the trends in accuracy improvement in Figure 2 ; (b) demonstrates that the accuracy of correctness predictors g m s improves as we increase the number of samples from d ′ used for training them, thus reducing the OOD gap; (c) shows the mean absolute error (MAE) of our kernel smoothing estimator of the accuracy of correctness predictors p(d ′ , m) -the estimator does not improve as much with increased α, thus S 3 eventually becomes worse than S 1 in terms of correlation and accuracy of the selected models.The dataset distance u(d ′ ) is computed as in equation 12. As evident from equation 13, dataset distance will usually decrease for larger values of α as inputs from d ′ are moved into D (assuming that inputs from d ′ are on average closer to each other than they are to inputs from other tasks).In this experiment, this serves as a mechanism to study the performance of LLM routing on closer datasets, providing insights into the benefits of learning LLM routers on more benchmarks where it is more likely that dataset distance for a new task is small.In Figure 8 we present relations between dataset distance u(d ′ ) and Pearson correlation between various model scores and accuracies of candidate LLMs. For results with S 3 see Figure 5 . Correctness predictor and metrics In the experiments on the MixInstruct dataset (Jiang et al., 2023) , we construct S 1 (m, d ′ ) following the scoring approach described in Appendix A, where the MixInstruct train set was defined as the benchmark dataset. Then, instead of computing a per-dataset score for the entire test set, we compute the score for each test point, i.e., S 1 (m, x d ′ i ) = g m (x d ′ i ), and select a model m per-point based on this score. The reported metrics in Table 2 and Figure 3 are averaged over the output evaluations of these per-point model selections. In our experiments, to compute g m (x d ′ i ) we use the BERTScore metric on the closest train set points (as y(x, m) in 11). This was motivated by the conceptual relation between the implementation of our approach and the BERTScore, which relies on embedding space distances, and was validated empirically.kNN parameter We set k = 10 for the kNN classifier, slightly higher than in the HELM experiments. This choice was motivated by the in-distribution properties of the test set in MixInstruct, which is constructed from different parts of the same datasets that comprise the train set. We note that the metrics did not significantly vary for different choices of k ∈ [5, 100].D PROOF OF LEMMA 4.1 Lemma D.1 (Lemma 4.1). Let ℓ(y 1 , y 2 ) = ρ(y 1 -y 2 ) for some subadditive ρ : R → R (e.g. ρ(x) = 1 2 x 2 for the square loss). We have Hugging Face currently hosts22,482 models for text generation.We omit dependency on the prompt when generating with an LLM and, w.l.o.g., consider the same LLM with a different prompting strategy as a different LLM. 


Decision: This paper focuses on addressing the following problem: given a number of LLMs, pick the best one per input to optimize for quality. The proposed approach relies on training a binary classifier per LLM on existing benchmarks. At inference time, it does not rely on labels and only operates on inputs and the classifiers -- a collection of binary classification jobs are run to predict which LLM will be the most accurate. The idea of using sentence transformer for embedding the data and using KNNs as correctness predictors for new tasks is interesting. The approach is computationally efficient to implement, fast at inference Extensive results (e.g. on HELM) shows the effectiveness of the approach considering that it only relies on the input, although for truly out of distribution tasks the generalization results are understandably less promising.

Reviews: review: paper_summary: paper_summary: This paper aims to selecting the best LLM from a collection of pretrained models for an unseen task. It is useful and motivates well. It also quantifies a rooting function to predict performance. The authors propose repurposing a collection of benchmark datasets to learn a "router" model for LLM selection from inputs only, reducing the problem to a collection of binary classification tasks.
strength_weakness: strength_weakness: I am clear about the idea and motivation. However, it is really confusing with the detailed methods, for Sections 3 and 4, especially with so many new notions.

Language and terminology should be refined for clarity and precision, especially in technical descriptions and explanations.
 It would be beneficial to include visual aids such as figures or tables to illustrate key concepts or experimental results for better comprehension.
paper_summary: paper_summary: This paper propose a method to select the best LLM from a collection of pre-trained models, given a new task. The proposed method does not assume any labeled data from the new task. It learns a “Router” model for LLM selection from inputs only. Experimental results show the effectiveness of the proposed method.
strength_weakness: strength_weakness: 1.	The improvements are limited according to Table 2.

2.	No human evaluation, I think there might at least be one to generally prove that the proposed method could predict the quality of a LLM’s output as expected.
questions: questions: N/A
paper_summary: paper_summary: This paper focuses on learning to route different LLMs for downstream tasks without using task-related labeled data. The authors propose to utilize the performance of these LLMs on benchmark datasets and train a router model to do binary classification on the LLMs' performance given the task inputs. The prediction scores are used to conduct the LLM selection. The approach leverages embedded inputs and task descriptions from benchmark datasets and employs k-nearest neighbors classifiers to predict LLM correctness, aiming to generalize to unseen tasks. Experimental results on HELM and MixInstruct datasets demonstrate the potential of selecting suitable LLMs to reduce the cost and improve overall performance. However, the authors mentioned that out-of-distribution generalization is the main challenge of LLM routers.
strength_weakness: strength_weakness: - Lack of explanation of the usage of KNN models. How does the proposed method deal with the out-of-distribution challenge.
- Lack of evaluation on more realistic tasks or instructions. Some instructions require the LLMs to use multiple or a combination of capabilities to solve, which is hard to reflect on the benchmarks. This composition problem needs to be discussed for the proposed method.
paper_summary: paper_summary: This paper propose a new approach to LLM selection without relying on labeled data from the new task. Instead, they develop a "router" model trained solely on inputs from various benchmark datasets, which reduces the selection problem to a collection of binary classification tasks. Empirical results show that this method consistently outperforms using a single model for all tasks.
strength_weakness: strength_weakness: (1) If I understand it correctly, you still assign a unique LLM to solve some specific task, but not different LLMs to solve a task? Basically, it would be intuitive, even if a capable model can still underperform a weak model on some instances of a specific task. Thus, it seems more reasonable if you route instances but not datasets to LLMs.

(2) Efficiency can be potentially an issue: since you use lots of LLMs, how do you effectively control the effiency cost of maintaining multiple LLMs like MoE?

(3) There lacks a deep discussion how this approach is compared with MoE, in both advantages and disvantages.

Comments: title: Response to Authors of Submission 1078
comment: I have read the rebuttals and changed my rating score accordingly.
Dear Reviewer, as the end of the discussion period is approaching, we'd appreciate it if you could comment on our rebuttal. Thank you.
Thank you. We will make sure to incorporate the comments into the revised version.
Thank you for updating the review. Please find our comments below.

> The improvements are limited according to Table 2

We agree that *performance* improvements over Llama 2 70B in Table 2 are small. However, the improvements in terms of *parameter count* are fairly substantial. For example, S3 on average selects an LLM with 0.71 of the Llama 2 70B parameters, which can reduce the inference cost, while providing some performance improvements.

We also note that the results in Table 2 are for the most challenging scenario when new tasks have no labeled data. In Figure 2 we show that the performance gains from selecting LLMs with our scores can be further increased when a few labeled samples from the new task are available.

> No human evaluation, I think there might at least be one to generally prove that the proposed method could predict the quality of a LLM’s output as expected.

HELM tasks are amenable to automatic evaluation and do not require human evaluation. For the MixInstruct (open-ended instruction following) experiments in Section 5.2, we used the metrics proposed by the dataset's authors [1]. Although they did not perform human evaluations—which we agree could strengthen the empirical findings—they verified that the considered automatic evaluation metrics correlate well with scores obtained using ChatGPT.

[1] Jiang et al. LLM-Blender: Ensembling Large Language Models with Pairwise Ranking and Generative Fusion.
I marginally raised my rating score, while I hope these comments would be incorporated into the revised version.
I am very much sorry with this mistake, the reviews are edited accordingly with my apology .