RouterDC: Query-Based Router by Dual Contrastive Learning for Assembling Large Language Models
Recent works show that assembling multiple off-the-shelf large language models (LLMs) can harness their complementary abilities. To achieve this, routing is a promising method, which learns a router to select the most suitable LLM for each query. However, existing routing models are ineffective when multiple LLMs perform well for a query. To address this problem, in this paper, we propose a method called query-based Router by Dual Contrastive learning (RouterDC). The RouterDC model consists of an encoder and LLM embeddings, and we propose two contrastive learning losses to train the RouterDC model. Experimental results show that RouterDC is effective in assembling LLMs and largely outperforms individual top-performing LLMs as well as existing routing methods on both in-distribution (+2.76\%) and out-of-distribution (+1.90\%) tasks. Source code is available at https://github.com/shuhao02/RouterDC.
Paper
Similar papers
Peer review
Summary
The paper introduces RouterDC, a novel LLM query router that employs contrastive learning losses to train an encoder and LLM embeddings for routing queries efficiently. The motivation behind this approach is the variability in LLM performance across tasks and domains, and the computational efficiency of routing over ensemble methods. The proposed methodology involves using a small LLM encoder, `mDeBERTaV3-base`, and learnable LLM embeddings for candidate LLMs. RouterDC employs a supervised learning approach with ground truth annotations, using contrastive losses to optimize the routing mechanism. Experimental results demonstrate that RouterDC outperforms existing routing methods on both in-distribution and out-of-distribution tasks.
Strengths
- The paper addresses a significant challenge in LLM utilization, focusing on efficient routing to optimize performance across different tasks and domains. - The contrastive learning approach for query routing is innovative, leveraging both *sample-LLM contrastive loss* and *sample-sample contrastive loss* to enhance training stability and performance. - Experimental results show that RouterDC significantly outperforms existing routing methods, indicating the effectiveness of the proposed approach.
Weaknesses
- The parameter efficiency of RouterDC compared to existing routing methods is not addressed, which is important for understanding the scalability of the approach. For example, the training cost is likely significant when the number of LLMs scales up (each training query needs to be evaluated on each LLM), and retraining is required when any new LLMs are incorporated. This could impact the practicality of the system. - The paper does not incorporate LLM costs into the loss function. When multiple LLMs perform well (which is the motivation of the work), it would be natural to choose the cheapest one to enhance cost efficiency. - Table 3 shows that removing certain LLMs, such as `Chinese-Mistral-7B` or `dolphin-2.6-mistral-7b`, does not affect overall performance, and removing `Mistral-7B` even improves it. This indicates that including incapable LLMs can lead to unnecessary computational overhead and performance degradation. There is a lack of analysis for this issue, which can be crucial for an effective routing system in the real world.
Questions
1. Table 1 lacks a comparison between different routing methods’ training time or training computational cost. How does RouterDC compare in terms of parameter efficiency with existing methods? This information is crucial for understanding the scalability of the approach. 2. The performance improvements after incorporating sample-sample contrastive loss seem marginal. Can this be further analyzed or explained? 3. On OOD tasks, RouterDC performs worse than the best-performing LLM on certain individual tasks. Is there any analysis of the reasons? 4. Evaluation results on specialized routing benchmarks like [1] could provide a more comprehensive assessment of RouterDC's performance. 5. A normalized average score, instead of the simple average among the scores of different benchmarks, can provide a better comparison in Tables 1 and 2, considering the varying complexity and score scales across benchmarks. 6. There is no need to include RandomSelect as a weak baseline. 7. In the "Routing to Different Numbers of LLMs" evaluation, why were LLMs added in the chosen order? Based on Table 1, adding them in performance-descending order might yield smaller accuracy enhancements. 8. Based on Table 3, the incorporation of certain incapable LLMs can lead to unnecessary computational overhead and performance degradation. Would an LLM-dropping mechanism improve overall performance? 9. The visualization of query embeddings from different benchmarks is confusing. Did you cluster queries within one benchmark or across several benchmarks in the sample-sample contrastive learning loss function? --- [1] Hu, Qitian Jason, Jacob Bieker, Xiuyu Li, Nan Jiang, Benjamin Keigwin, Gaurav Ranganath, Kurt Keutzer, and Shriyash Kaustubh Upadhyay. "ROUTERBENCH: A Benchmark for Multi-LLM Routing System." arXiv preprint arXiv:2403.12031 (2024).
Rating
7
Confidence
4
Soundness
3
Presentation
4
Contribution
3
Limitations
The broader implications of this routing approach, including its scalability and cost efficiency in practical applications, should be more adequately addressed.
Thank you for the detailed responses. I am happy to vote for acceptance. Please ensure that these discussions and results are included in your revision. However, after re-reading the paper, I noticed that it currently lacks a comparison or at least a discussion of the existing cascade-based approaches [1-4] for assembling LLMs. Unlike ensemble-based methods, cascade approaches do not necessarily invoke all models and can also achieve good cost-performance trade-offs as routing approaches (indeed, they may invoke multiple LLM calls for one query, but this does not necessarily mean that their cost-efficiency is worse). Including such a discussion or comparison would provide a more comprehensive understanding of your work within the context of existing research. [1] Model Cascading: Towards Jointly Improving Efficiency and Accuracy of NLP Systems, EMNLP 2022 [2] Language Model Cascades: Token-level Uncertainty and Beyond, ICLR 2024 [3] Large Language Model Cascades with Mixture of Thoughts Representations for Cost-efficient Reasoning, ICLR 2024 [4] Online Cascade Learning for Efficient Inference over Streams, ICML 2024
Reply to Reviewer 4ok1
Thanks again for your positive rating. We will certainly add the above discussions and experiments to the revision. --- > Q12. However, after re-reading the paper, I noticed that it currently lacks a comparison or at least a discussion of the existing cascade-based approaches [1-4] for assembling LLMs. **A12.** We have discussed a cascade-based method (i.e., FrugalGPT) in the related work of the paper (Lines 75-77, Lines 182-184). Thanks for bringing the four other cascade-based methods to our attention. We agree that cascade-based methods are one direction to achieve cost-effectiveness when choosing LLMs, but they are different from RouterDC in terms of **settings, inference cost, and tasks**. We discuss the differences between our RouterDC and the mentioned cascade-based methods below. (i) RouterDC considers **a different setting** compared with the cascade-based methods [1-4]. The cascade-based methods usually **assume that the capacity of LLMs depends on the model size**. Their intuitive idea is to query LLMs from weak (small) to strong (large) until a satisfactory answer is obtained instead of calling the strong LLMs for all queries. Our RouterDC does not require this assumption and can select a suitable LLM from multiple small or large candidate LLMs. Hence, routing-based methods are more general. Furthermore, even if LLMs are of the same size, they may have different specialized capabilities. (ii) Cascade-based methods [1-4] may call LLMs **multiple times** for a query (in the worst case, all candidate LLMs need to be called), but our RouterDC only needs to call the selected LLM **once** in inference/testing. (iii) RouterDC is general and can be used for **generation tasks**, but Model Cascading [1] and Online Cascade Learning [4] are limited to **classification task** (e.g., SST2, MRPC, IMDB). Generation tasks are usually more useful and challenging than classification tasks in NLP. We will include the above discussion and related works in the revision. --- #### References [1] Model Cascading: Towards Jointly Improving Efficiency and Accuracy of NLP Systems, EMNLP 2022 [2] Language Model Cascades: Token-level Uncertainty and Beyond, ICLR 2024 [3] Large Language Model Cascades with Mixture of Thoughts Representations for Cost-efficient Reasoning, ICLR 2024 [4] Online Cascade Learning for Efficient Inference over Streams, ICML 2024
Summary
The authors propose a routing between different LLMs, based on classification of embeddings from a fine-tuned transformer model (mDeBERTaV3-base), on a per-query basis. To sharpen the classification, they chose positive and negative samples among the training tasks, and to stabilise training they add a loss term that encourages cohesive clusters. The results suggest that the system is capable of routing appropriately.
Strengths
* Experiments use a realistic pool of local LLMs * Practical to implement - though makes most sense if the upstream LLMs are being paid for per-token, rather than by availability Original Rating: 4 Original Confidence: 3
Weaknesses
* Since the LLMs in the pool appear to be pretty decisively good at some of the tasks (Figure 9), the meat of the router task is to see whether it's possible to classify these different tasks via a small LM and embeddings. On its face, it doesn't seem too remarkable that this works. * Sample-Sample Contrastive Loss : Smells like a post-experiment fix, rather than a principled choice + Particularly since a pretrained (frozen?) mDeBERTaV3-base is used to determine which samples 'belong' to which clusters Minor point(s) * L30: "Figure 3 shows the scores of seven LLMs for an example query, where the top three LLMs have significantly higher scores." + Seems unclear whether top-3 is so different from top-4 : Maybe a better example could be used
Questions
* For the "OOD" results, how disjoint (really) are (i) CMMLU and C-EVAL; (ii) HumanEval and MBPP? * Table 3: Robustness of RouterDC to LLM losses during inference - isn't this proving that (for instance) Meta-Llama3-8B is being consistently chosen for CMMLU? * L46: "To improve the training stability, we cluster the training queries" - intra-group vs inter-group (within batch?) + L225: The number of clusters N is set to 5 : Is it a coincidence that this matches the number of training tasks? - Appendix B : Is it troubling that task-identity is roughly the same as pretrained cluster labels? Perhaps the task prompts are a give-away... - L257: "Moreover, increasing N leads to higher average accuracy when N is small (≤ 4), but the accuracy saturates quickly." (ditto) * How could this approach work in a chat context? Is it just for single queries?
Rating
6
Confidence
4
Soundness
3
Presentation
3
Contribution
2
Limitations
* The router being learned here is a task-wise classifier. There is likely commercial value in having this kind of routing, but implementing/testing this doesn't seem particularly novel (apologies for the vagueness).
Summary
This paper studies the problem of assembling off-the-shelf LLMs to harness their complementary strengths. The authors propose a novel query-based router by Dual Contrastive learning (RouterDC), i.e., a sample-LLM contrastive loss and a sample-sample contrastive loss. The former contrastive loss aims at training the router such that it can assign suitable LLMs for queries, while the latter is for training stability. Experiments on various challenging tasks demonstrate that RouterDC performs better than existing routing methods and individual top-performing LLMs in both in-distribution and out-distribution settings.
Strengths
1. The proposed dual contrastive losses for training a query-based router are novel. The sample-LLM contrastive loss seems more sound for learning the router than the KL loss used in previous work ZOOTER. 2. Extensive experiments on both in-distribution and out-distribution tasks show that the proposed RouterDC outperforms existing routing methods on average. The performance shown in Fig.2 confirms that RouterDC can harness the complementary strengths of off-the-shelf LLMs. 3. The paper is well-written and easy to follow. It provides many comprehensive ablation experiments, e.g., Fig. 11 shows the advantage of the sample-sample contrastive loss in improving the training stability of RouterDC.
Weaknesses
1. In Line 157, the authors claim that “The reason is that some similar queries can have dissimilar embeddings and may be routed to different LLMs.” Evidence should be provided to support this claim. 2. For the OOD setting (Table 2), RouterDC fails to beat the best individual LLMs on all tasks (e.g., 38.81 vs 39.72 on PreAlgebra). 3. RouterDC may require a large amount of labeled data to train the router.
Questions
1. How about the performance of RouterDC without the sample-sample contrastive loss? 2. In Fig. 9, it seems that no samples are routed to the Chinese-Mistral-7B model, why? 3. Can you visualize the embeddings of training samples extracted by the encoder $\mathcal{E}(x;w)$ of RouterDC?
Rating
7
Confidence
4
Soundness
4
Presentation
3
Contribution
3
Limitations
The authors are encouraged to discuss limitations in the conclusion section.
Summary
This paper aims to improve the ability of LLMs by assembling them. The proposed method use a contrastive learning strategy. It uses a sample-LLM contrastive loss which pull the query embedding closer to the top-performed LLM embedding. It also employs a sample-sample contrastive loss, which learns the distribution of input with the help of clustering.
Strengths
The contrastive learning idea is intuitive. The routing model is only 86m, which is quite small compared to the LLMs. And the method achieve better results compared with simple voting or learning the routing with a reward model.
Weaknesses
The contribution of this work lies in two part, first, it uses sample-llm contrastive loss, instead of directly learning the routing as a classification; second, it use sample-sample contrastive loss to make the training more stable. However, the relation between the two is missing. I am wondering how the two parts contribute to the final performance. Is it possible to use ZOOTER with the sample-sample loss? What do you expect as the result? I think the response partially addressed my concern. So I raised my score to 5. I do understand there are hyper-parameters that one could control to affect the results, but are the proposed two contrastive losses the most important ones to be consider in this scenario? I would love to see more discussions in the paper.
Questions
See the weakness part.
Rating
5
Confidence
5
Soundness
2
Presentation
3
Contribution
2
Limitations
N/A
* A1+A8 - It was already clear that you are doing query-based rather than task-based classification. One of my concerns was that for the data you trained/tested over, the two things were so closely matched that the different datasets would leak the task just due to the wording, etc. Your Table R2 helps. * A4 : I'm surprised that there are *any* strict overlaps between the datasets : Interesting! But that doesn't change the point that the the claim that the alternate datasets are 'OOD' seems like an overreach. I'm happy to update a little: New Rating: 5 New Confidence : 4
Reply to Reviewer GGL7
Thanks for your further comments and **raising the score**. For the remaining concerns, we address them as follows. --- > Q11. One of my concerns was that for the data you trained/tested over, the two things were so closely matched that the different datasets would leak the task just due to the wording, etc. **A11.** We understand the reviewer’s concern that certain words in the query may leak the task identity, making it easy for RouterDC to perform like a task classifier. To resolve this concern, we conducted an additional experiment in **a single task setting**, i.e., we train the router on the training set of HumanEval and evaluate it on the testing set. The single task setting is an edge case where **all queries may contain the same task information**. Hence, the router needs to learn how to route queries appropriately based on the query itself instead of some possible task information contained in the query. Table below reports the testing accuracy. As can be seen, **RouterDC largely outperforms the best candidate LLM (i.e., dolphin-2.9-llama3-8b) and existing routing methods**, demonstrating that the router can select appropriate LLMs for queries based on query characteristics. We will add the experiment and discussion to the revision, which definitely improve our work. \begin{array}{lc} \hline \text{Method} & \text{HumanEval} \newline \hline \text{Mistral-7B} & 28.98 \newline \text{MetaMath-Mistral-7B} & 29.80 \newline \text{zephyr-7b-beta} & 22.04 \newline \text{Chinese-Mistral-7B} & 21.43 \newline \text{dolphin-2.6-mistral-7b} & 45.10 \newline \text{Meta-Llama-3-8B} & 26.73 \newline \text{dolphin-2.9-llama3-8b} & 49.39 \newline \hline \text{ZOOTER} & 39.38 \newline \text{CosineClassifier} & 52.45 \newline \text{RouterDC} & \mathbf{56.32} \newline \hline \end{array} --- > Q12. I'm surprised that there are any strict overlaps between the datasets : Interesting! **A12.** Thanks for your comments! We guess there was a typo in the comment, it should be "there are **NOT** any strict overlaps between the datasets : Interesting!" > Q13. But that doesn't change the point that the the claim that the alternate datasets are 'OOD' seems like an overreach. **A13.** Thanks for your insightful comments! We appreciate the reviewer raising the concern about the definition of OOD. In the paper, C-EVAL and MBPP are treated as OOD tasks as they are different task distributions or question-answer instruction. We briefly summarize their differences below. (i) CMMLU and C-EVAL have **different task distributions**. CMMLU contains more culture- and region-related tasks, while C-EVAL have more STEM tasks. Moreover, CMMLU and C-EVAL use **different prompts** to ask multiple-choice question. C-EVAL uses continuous underscores to indicate the answer's location, whereas CMMLU employs no special notation for referencing the answer, except for brackets when the answer is within a sentence. (ii) HumanEval and MBPP assess the code generation proficiency of LLM from **two distinct perspectives**. A HumanEval query **gives the header** of a python function and some comments, requiring the LLM to **implement the rest** of the function. On the other hand, a MBPP query **gives an intent** and asks the LLM to **generate the function from scratch**. To further resolve this concerns, we evaluate the learned router on **one more OOD task: JavaScript** [R1], which aims to generate JavaScript code to solve problems. Different from HumanEval, which generates Python code to solve problems, JavaScript can be viewed as **a distant OOD task**. Table below reports the testing accuracy. As can be seen, RouterDC outperforms existing routing methods by a large margin, demonstrating that our RouterDC is more effective in routing queries of the distant OOD task. We will include the additional experiments and discussions in the revision. \begin{array}{lc} \hline & \text{JavaScript} \newline \hline \text{Mistral-7B} & 29.88 \newline \text{MetaMath-Mistral-7B} & 31.83 \newline \text{zephyr-7b-beta} & 11.71 \newline \text{Chinese-Mistral-7B} & 17.68 \newline \text{dolphin-2.6-mistral-7b} & 45.00 \newline \text{Meta-Llama-3-8B} & 37.07 \newline \text{dolphin-2.9-llama3-8b} & 53.84 \newline \hline \text{ZOOTER} & 41.64 \newline \text{CosineClassifier} & 37.32 \newline \text{RouterDC} & \mathbf{48.66} \newline \hline \end{array} --- ### references [R1] CodeGeeX: A Pre-Trained Model for Code Generation with Multilingual Benchmarking on HumanEval-X. KDD 2023.
A12. """Thanks for your comments! We guess there was a typo in the comment, it should be "there are NOT any strict overlaps between the datasets : Interesting!" """ Actually, there was no typo. Your claim is that one dataset is OOD compared to the other : Surely then, it should be surprising that they have textually identical pairs in them? A11(+A13) : This is good additional work, that goes far further in proving your case. Thank you. And for that part, I'll inch up one (last) time: New Rating: 6 New Confidence : 4
Reply to Reviewer GGL7 (2)
We are glad that our additional experiments have resolved your concerns, and we will add the experiments to the revision. --- > Q14. Actually, there was no typo. Your claim is that one dataset is OOD compared to the other : Surely then, it should be surprising that they have textually identical pairs in them? **A14.** We apologize for our misunderstanding of your comments. As discussed previously, most of the CMMLU and C-EVAL samples (about 99%) are different. Thus, the tiny overlap may not significantly affect the evaluation.
A Gentle Reminder for Reviewer ytgM
Dear Reviewer ytgM, We sincerely thank you again for your effort to improve our work. We have provided a detailed response to resolve your concerns. We would like to kindly remind the reviewer that the close date of the reviewer-author discussion is approaching. Please let us know if you have any additional questions or comments. Best, The Authors
Decision
Accept (poster)