Paper: Large Language Model Routing with Benchmark Datasets (ICLR-2024)

Reviewer comment (Experiments):
The paper only compares with relatively straightforward instance based approach while in related routing LLM sections, it does mention more approaches but not compare with them.

Author response:
As discussed in the corresponding related work section, all prior methods consider instance-based LLM selection and (1) require a sufficient amount of **labeled data from the new task** to train a scoring or ranking model and (2) require generating outputs with **every candidate LLM**. In contrast, we consider a setting where no (or very little) labeled data/references on a new task are available, which we believe is a more practical scenario. The only baseline applicable in such a problem setting is LL (or perplexity), which we compare to in our experiments. We note that LL still suffers from (2), while our method avoids this problem while also outperforming this baseline.

One exception is the experiment in section 5.2, where the train and test data are from the same distribution and thus it is possible to compare to additional baselines mentioned in the related work section, e.g., PairRanker (note that these baselines continue to be inefficient dues to (2)).

In Figure 2, we also added an additional baseline, Few-Shot, which uses the small number of labeled samples from the new task to select an LLM.

Reviewer comment (Evaluation):
There is no clear conclusion that can be drawn from the paper that in practice, what score to use (and with O.O.D score, whether one should use a score). For example, the NN experiments show that S1 performs the best while in Table 1, S3 performs the best. Note that due to the fact that we work on dataset, there is only 28 datapoints. This is not to blame the authors but with this data size and different results, it seems impossible to draw conclusions.

Author response:
Based on the experiments in Figure 2, we recommend using kNN as the correctness predictor together with S3 while labeling \~5 samples from the new task for the best tradeoff between the router performance and labeling efforts. The reason to use kNN over MLP is that it is straightforward to have it take advantage of the few labeled samples from the new task, while MLP would require retraining/fine-tuning with up-weighted new-task labeled samples. When it is possible to obtain more labeled samples from the new task (\~25-50), we recommend S1 instead of S3 as the kNN classifier becomes better calibrated, thus reducing the need to model correctness with S3. This recommendation is also supported by the results in Figure 2 for higher $\alpha$ values and is briefly discussed in the last paragraph of Section 5.1.

That said, rather than a concrete practical recommendation, the main contribution of our work is the new problem formulation of LLM routing using *benchmark datasets* that has the potential to improve the collective capabilities of open-source LLMs. The use of benchmark datasets is important as it allows researchers with limited computational resources to contribute to LLM research by reusing publicly available LLM (per-sample) evaluation data (e.g., HELM, Open LLM Leaderboard, evaluation data released as part of our code) instead of running costly evaluations themselves. Thus, our main focus is on bringing forth various aspects of the problem along with methodological approaches to address them.

Reviewer comment (Writing):
For S3, the authors say that "we assign a task descriptor u(d)", what is this u(d) please? This is a bit confusing since Appendix A further uses u(d) as dataset distance.

Author response:
By "task descriptor" we simply mean some scalar or vector representation of a task. The specific choice of task descriptor u(d) in our paper is described in Appendix A and is based on dataset distance, although other choices (including those not based on distance) are conceptually possible. We have edited Appendix A to clarify this.

Reviewer comment (Evaluation):
By using S3, how many tasks end up using BMA? Can you comment on the difference between Table 1 and Table 3?

Author response:
In both Table 1 and Table 3, the column titled "\% BMA" specifies the fraction of times that the BMA was selected by each of the methods. For example, based on Table 1, BMA is chosen $29\times 0.48\approx 14$ times by S3. The difference between Table 1 and Table 3 is the type of correctness predictor, $g_m(x)$, used. In Table 1, we use non-parametric kNN classifiers, and in Table 3, we train Multi-Layer Perceptrons (MLPs) as the correctness predictors.

Reviewer comment (Theory):
I'm somewhat confused about whether it's crucial to use "imperfect" one if we have a "perfect" correctness predictor. In other words, why do we opt for an imperfect correctness predictor, such as a non-parametric classifier, instead of a parametric one?

Author response:
If we have a "perfect" correctness predictor, the problem is fairly trivial and all proposed scores will perform as the oracle. Unfortunately, the perfect correctness predictor is generally unattainable, analogous to perfect accuracy being unattainable in most supervised learning tasks, especially with distribution shift. The "imperfection" is not related to the choice of correctness predictor (nonparametric vs parametric) but to the challenging nature of the binary classification problem corresponding to correctness prediction. In Appendix B.1 we present results using a neural network as the correctness predictor, i.e., a parametric classifier.

Reviewer comment (Novelty):
From my perspective, this work bears similarities to the Mixture of Experts (MoE) model, where experts in MoE are replaced with LLMs. So, what is the distinguishes between this work and MoE, where LLMs serve as experts? Would the non-parametric method remain efficient if we use it for the traditional MoE?

Author response:
The traditional application of MoE is to decide which expert(s) to use based on the input. In this paper, we mainly consider the problem of selecting a single model for a task, i.e., a collection of inputs. It is perhaps possible to use some variation of MoE to improve correctness predictors ($g_m$s), thus potentially improving the effectiveness of model routing with our scores, however, it is unclear how MoE alone can be used in our problem setting.

One exception is the experiment in section 5.2, where the model is selected per sample and some of the baselines, e.g., PairRanker (which requires inference with all experts unlike our method), can be viewed as variations of MoE.

We have added a brief discussion regarding MoE in the conclusion.

Reviewer comment (Experiments):
This paper doesn't seem to clarify the difference between this method and certain fine-tuning techniques, nor does it address whether the proposed method outperforms the current fine-tuning methods. If we fine-tune the selected LLM, would it certainly perform better than an LLM that hasn't been selected? Or should we use the selected LLM directly after the router has chosen it?

Author response:
We consider the problem of selecting an LLM for a new task where no (or very few) labels/references are available, thus fine-tuning is not a viable option. In our experiments, the LLM selected by the router is directly applied to the samples from the new task.

When the new task comes with a sufficient amount of labels/references, choosing a model most suitable for fine-tuning is an interesting problem, but is beyond the scope of this work. In the future work part of the conclusion, we briefly discuss the potential impact of routing models that were already fine-tuned for various domains, i.e., "expert" LLMs, where we anticipate our method to continue being effective.

Reviewer comment (Experiments):
Comparison with Existing Methods: The concept of routing is a prevalent strategy in conventional Mixture-of-Experts (MoE) solutions. More comprehensive discussion and experimental comparisons are encouraged.

Author response:
The traditional application of MoE is to decide which expert(s) to use based on the input. In this paper, we mainly consider the problem of selecting a single model for a task, i.e., a collection of inputs. It is perhaps possible to use some variation of MoE to improve correctness predictors ($g_m$s), thus potentially improving the effectiveness of model routing with our scores, however, it is unclear how MoE alone can be used as a baseline.

One exception is the experiment in section 5.2, where the model is selected per sample and some of the baselines, e.g., PairRanker (which requires inference with all experts unlike our method), can be viewed as variations of MoE.

Following your suggestion, we have added a comment regarding MoE in the conclusion.

Reviewer comment (Experiments):
Results of "S3 true p" need further practical analysis. These results are only achieved when the model has access to the true accuracy of correctness predictors.

Author response:
We present this score to demonstrate the upper bound on the proposed model routing score that takes into account the potential inaccuracies of the correctness predictors. As discussed at the top of page 7, we find it encouraging that even with correctness predictors that on average have accuracy as low as 0.59 on a binary classification task, it is conceptually possible to route LLMs efficiently.

Reviewer comment (Writing):
$g_m(x)$ is defined to evaluate the correctness of model $m$ on an input $x$ and gold label $y$. The lack of $y$ in $g_m(x)$ causes confusion.

Author response:
$g_m(x)$ is a classifier which estimates the probability that an LLM $m$ will generate the correct answer for an input $x$. The probability is determined entirely from $x$ so including $y$ in $g_m(x)$ would be incorrect.

Reviewer comment (Theory):
Eq.1 is a little bit confusing. It estimates the loss of $g_m(x)$ and $y(x,m)$ given $g_m$. $g_m$ seems to be both independent variable and dependent variable.

Author response:
For our approach, we are training the probabilistic estimator $g_m(x)$ to predict the gold label $y(x, m)$. Eq. 1 states that we are selecting the $g_m(x)$ that minimizes the binary cross-entropy loss between the prediction and the gold label. In other words, we are simply training a standard binary classifier in eq 1. We have clarified this in the updated draft.

Reviewer comment (Theory):
The problem of OOD in Eq.3 lacks necessary discussion. Eq.3 does not contain notation of $P(y|x)$, it estimates $g_m$. Although the target of $g_m$ is to estimate the correctness of model $m$ that can be potentially affected by OOD, there still lacks necessary clarification about what $P(y|x)$ represents in Eq.3.

Author response:
Eq. 3 uses $g_m$, which is an estimator of $P(y(x,m)=1|x)$ as stated after eq. 1. We are simply stating that a (probabilistic) binary classifier may fail to generalize out-of-distribution, particularly in terms of its calibration quality. We have clarified this in the updated draft.

Reviewer comment (Theory):
The relation of solution of Eq.4 and OOD is not clear.

Author response:
Eq. 4 presents one of the model routing scores based on predictions with $g_m$s. As a binary classifier, $g_m$ is likely to have lower accuracy when evaluated on OOD data, thus the score based on these predictions may be less effective when selecting an LLM for a new task. We have clarified this in the updated draft.

Reviewer comment (Novelty):
This approach takes a lot of efforts on choosing LLMs based on a prediction function. For example, Eq.3 chooses a language model directly based on the prediction of $g_m$. Eq.4 introduces the threshold and can get better generalization results. These tricks are valuable, but some of them are popular in traditional classification methods.

Author response:
Could you please clarify the concern? If the question is regarding the novelty of the proposed scores, then we agree that S2 (eq. 4) is a standard trick. However, the subsequent discussion on modeling LLM correctness and the construction of S3 is novel to the best of our knowledge. We also want to emphasize that the main contribution of our work is the new problem formulation of LLM routing using *benchmark datasets* that has the potential to improve the collective capabilities of open-source LLMs. The use of benchmark datasets is important as it allows researchers with limited computational resources to contribute to LLM research by re-using publicly available LLM (per-sample) evaluation data (e.g., HELM, Open LLM Leaderboard, evaluation data released as part of our code) instead of running costly evaluations themselves. Thus, our main focus is on bringing forth various aspects of the problem along with methodological approaches to address them.