GraphRouter: A Graph-based Router for LLM Selections

Q1. The initial LLM node feature is based on text embedding, which is likely not the best strategy to model numerical information like token pricing and context length when there is a need to capture the subtlety in numerical differences.

Response: Thanks for your valuable comments. We agree with your insights that to learn effective LLM representation, node initialization is critical. However, how to use embeddings to understand and model numerical differences is difficult, which remains an open question [1, 2, 3]. In fact, the focus of our paper is on how to learn effective LLM embeddings that enhance the ability to select LLMs. We achieve this by analyzing interaction data between task queries and LLMs, and leveraging the robust message-passing capabilities of GNNs. Thus, we adopt a simple approach where we represent the LLMs through the direct PLM’s embeddings of its descriptions. Our experiment has demonstrated that even with the simplest way of initializing LLM node’s embeddings, our framework still significantly outperforms the existing baseline (Ding et al., 2024;Chen et al., 2023;Dai et al., 2024). While we can use techniques, such as adding relative magnitude tokens (RMT) to better discretize numerical feature values (Yan et al., 2024), to enhance the LLM nodes’ embeddings, this is not the main focus of this paper, and We will leave this for future studies to improve and enhance.

[1] Limitations of language models in arithmetic and symbolic induction, arXiv 2022.

[2] Towards Cross-Table Masked Pretraining for Web Data Mining, WWW 2024.

[3] Making pre-trained language models great on tabular prediction, arXiv 2024.

Q2. From L245-246, the paper assumes a universally optimal LLM selection for each query, which may not be the case in practice with constraints like inference cost and latency. This should also be discussed in the limitation and future work sections.

Response: Thanks for the reviewer’s insightful feedback. In fact, we do not assume a universally optimal LLM selection for each query. As outlined in Section 3.4, the weights (e.g., $\alpha$ and $\beta$) and elements (e.g., Performance and Cost) included in the reward framework inherently result in varying optimal LLM selections. Due to the high uncertainty and difficulty in quantifying inference cost and latency in real-world scenarios, we follow the settings of existing work (Ding et al., 2024; Chen et al., 2023; Dai et al., 2024), which primarily consider performance and cost in the reward—these two quantifiable metrics (as specifically defined in section 3.4) are often trade-offs against each other and also represent the main concern of the users. We demonstrate through experiments under 3 different reward settings that GraphRouter achieves the best results compared to baselines in various scenarios.

Q3. The discussion of related works can be further improved for proper credit attribution. For example: The authors cite "Xie et al., 2022" for label propagation while the credit at least should also be properly attributed to the original work "Zhu & Ghahramani, 2002" [1]. R-GCN is likely the first heterogeneous GNN to the best of my knowledge. It should also be cited in discussing heterogeneous graph neural networks [2]. The authors employ an attention-based heterogeneous GNN and should also discuss the first few attention-based heterogeneous GNNs, such as [3]. [1] Zhu & Ghahramani. Learning From Labeled and Unlabeled Data With Label Propagation. 2002. [2] Schlichtkrull et al. Modeling Relational Data with Graph Convolutional Networks. 2017. [3] Wang et al. Heterogeneous Graph Attention Network. 2019.

Response: Thank you for your insightful comments. Following the reviewer's suggestions, we revised the corresponding discussion in the related work section: Traditional graph algorithms, such as label propagation (Xie et al., 2022; Zhu & Ghahramani, 2002), directly utilize edge relationships to propagate known labels to target nodes. Additionally, Heterogeneous Graph Neural Networks (HeterGNNs) (Peng et al., 2019; Hu et al., 2020; Schlichtkrull et al., 2017) and Heterogeneous Graph Attention Networks (HGATs) (Wang et al., 2019) have also been proposed to handle more complex graph data.

Q1. The evaluation setting is simplified. On one hand, the experimental tasks primarily focus on reasoning and question answering, overlooking specialized areas such as code generation. On the other hand, the available LLMs are quite limited, with nearly half originating from the same series (i.e., LLaMA). Can the empirical evaluation be conducted in a more complex scenario, such as by adding more LLMs or introducing a broader variety of tasks?

Response: Thanks for your constructive feedback. Following the reviewer’s advice, we added two new datasets and four new LLMs (covering all the APIs currently available on Together AI) to the experiment. We first expanded upon the dataset from Section 3.1 by adding two new datasets. 1) HumanEval [1] is a dataset that measures LLMs' coding capabilities, specifically consisting of 164 original programming problems that assess language comprehension, algorithms, and simple mathematics, with some problems comparable to basic software interview questions. 2) HotpotQA [2] is a question-answering dataset with 113k entries featuring natural, multi-hop questions with strong supervision for supporting facts to enable more explainable question-answering systems. Similarly, we summarized the metrics and data volume for these datasets in the first table below. Additionally, we incorporated four more LLMs from the Together AI API: Qwen-2 (72b), Code Llama (34b), Mixtral-8x22B, and Upstage. Likewise, the size and cost information of these LLMs are summarized in the second table below. Further, similar to section 3.2, we constructed a dataset based on the extended dataset and the interaction data from LLMs. This dataset was then split into training, validation, and test sets in a 70%:10%:20% ratio, based on distinct queries. We compared the performance of GraphRouter with other baselines on this extended dataset and reported the experimental results in the third table below. We observed that GraphRouter improved the Reward by at least 12.3% compared to the baselines, confirming GraphRouter's strong generalizable abilities to more datasets and other LLMs. All the discussions and experimental results mentioned above have been updated in Section B.1 of the Appendix and Table 20-22 of the revised pdf version.

Table 1: Overview of extended datasets

Table 2: Statistics of larger LLMs set and their costs on Together API.

Dear Reviewer pHun,

We recognize that the timing of this discussion period may not align perfectly with your schedule, yet we would greatly value the opportunity to continue our dialogue before the deadline approaches.

We hope that our responses and additional experiments have effectively addressed your concerns. We truly appreciate all the valuable advice we have received. Could you let us know if your concerns have been adequately addressed? If you find that your concerns have been resolved, we would appreciate it if you could reconsider the review score.

Thanks!

Q1. Authors only provide intuitive explanations for why graph structure should help with LLM selection, lacking analysis on why edge prediction correlates with routing performance. Also, the paper fails to explain why the heterogeneous graph structure (Figure 5) is optimal for capturing LLM-query relationships.

Response: Thanks for your valuable questions and insightful feedback. We answer your questions step by step.

[Why edge prediction correlates with routing performance]. In the paper, we introduced in [lines 240-242] why edge prediction correlates with routing performance. Specifically, we model the LLM selection problem as predicting and selecting the LLM node with the highest reward for each query node. To achieve this, we construct edge labels by connecting the query node to the LLM node with the highest reward for the training of GraphRouter. This enables the GraphRouter, during the test phase, to select the LLM node connection with the highest probability for each query node as the result of the LLM selection.

[Design of heterogeneous graph structure]. In fact, we do not claim that the heterogeneous graph structure is optimal for capturing LLM-query relationships. We believe that our proposed heterogeneous graph structure is sufficiently elegant and naturally models LLM-query relationships, and it has achieved the best results compared to baselines during the experimental process. Furthermore, in response to Reviewer pHun's Q4, we have explored the use of another lightweight graph structure, such as SGC [1], to address the LLM selection problem. We have included discussions and experimental comparisons of SGC and our graph structure in Section C.4 and Table 27 of the Appendix in our revised PDF.

[1] Simplifying graph convolutional networks, ICML 2019.

Q2. L. 327 presents results across different LLMs but doesn’t analyze how model architectures affect routing decisions. There is no discussion of how to efficiently update the graph structure when new LLMs are added; Example: The experiments in Table 5 show impressive few-shot performance but don't evaluate beyond 10 LLMs, leaving questions about larger-scale deployments.

Response: We thank the reviewer for raising the above concerns. We would like to answer your questions one by one.

[Graph structure update]. Indeed, we have introduced the details of updating the graph structure in [Figure 5]. For new queries, LLMs, and tasks, the fundamental update to the graph's structure essentially involves partially adding new edges and nodes. Then, as discussed in the introduction [lines 76-94], leveraging the powerful inductive capabilities of GNNs, GraphRouter can efficiently and effectively generalize when faced with updates to the graph.

[Larger-scale deployments]. As shown in our related work session, the problem of LLM selection has already been studied by a few existing works (Ding et al., 2024; Chen et al., 2023; Dai et al., 2024). While these works all experiments with fewer than five LLMs, we extend our evaluation with ten LLMs, which is already the largest cale evaluation in this task to date. Moreover, in our experiments, we have utilized all the LLM APIs available for inference on the Together AI platform (some APIs on the platform may have failed).

In addition, the large-scale deployments of GraphRouter are also made possible by its GNN backbone. In established and ordinary GNN tasks, such as those detailed in the Open Graph Benchmark (OGB) [1,2], the number of nodes and edges can escalate to hundreds of millions. On these datasets, GNNs are even capable of training in less than a minute and performing inference in seconds on a standard laptop. In real-world, large-scale deployment scenarios like recommendation systems, GNNs have demonstrated the capability to efficiently manage up to tens of billions of nodes and trillions of edges [3,4]. In contrast, the number of LLM nodes falls short of reaching the million mark, indicating that technically, utilizing GraphRouter on larger-scale scenarios is feasible and efficient without issue. Moreover, in real-world application scenarios, the massive interaction data generated by numerous users enables GraphRouter to learn more refined LLM embeddings, thereby allowing it to better generalize to other LLMs.

[1] Open graph benchmark: Datasets for machine learning on graphs, Neurips 2023.

[2] Graph neural networks: A review of methods and applications, AI open 2020.

[3] Graph convolutional neural networks for web-scale recommender systems, KDD 2020.

Q3. In L. 219, task-query edges are initialized uniformly to 1, which seems overly simplistic given the rich task-query relationships that could be captured. Have you considered incorporating more sophisticated edge features, particularly for task-query relationships? How might this affect the framework's performance? LLM-query edge features only use performance and cost concatenation, ignoring other potentially valuable signals like response length or generation time. The ablation studies in Section 4.3 don't explore alternative edge feature designs Example: In Table 4, some performance variations could potentially be explained by inadequate edge features, but this is not analyzed.

Response: Thanks for the reviewer’s constructive feedback. We answer the questions from two aspects.

[Embedding similarity-based task-query edges]. We have compared the performance of GraphRouter when using more sophisticated edge features versus our current implementation of initializing all edges with a value of 1. To model more complex task-query relationships into the edge, we assign edge values based on the embedding similarity between task descriptions’ embedding and queries embedding, and we name this modified version of our model Edge-similarity. We compared the Reward of GraphRouter and Edge-similarity across three scenarios. As shown in the following table, additional task-query edge information would harm the selection process, as Edge-similarity underperforms GraphRouter across all three scenarios. Although we agree that adapting more advanced relation extraction models could potentially improve the performance, under the current framework, we conclude that complex semantic relationships between queries and tasks would not facilitate the selection performance. Additionally, since this task requires substantial effort and falls outside the primary focus of our paper, we leave this topic for future research. We have concluded these discussions and experiments in section C.2 and Table 25 of the Appendix in our revised PDF.

Comparison of GraphRouter and Edge-similarity Across Different Scenarios Focusing on Reward. It evaluates reward variations between two models under three distinct scenarios. Each number is formatted to three decimal places for precision.

[Impact of different LLM-query edge features] Following the reviewer’s advice, we investigated the Reward of GraphRouter with additional edge features. Specifically, we established three different edge feature combinations, Plus length and Plus time respectively represent the addition of token length and LLM inference time to the GraphRouter base edge features, and Plus length & time represents the addition of both edge features to GraphRouter. As shown in the table below, we compared the Reward values of the origin GraphRouter with the GraphRouter adding the above edge feature combinations across three scenarios.

We observe that while in the Cost First scenarios, an improvement has been shown, it is the least frequently encountered scenario in real-world applications as the actual quality of the LLM’s response is often more prioritized. In the other two scenarios, the gains are minimal or even result in a performance decline. This may be because, in the Performance First and Balance scenarios, Performance significantly contributes to Reward, and both token length and LLM inference time are more closely related to the Cost metric. Therefore, adding these edge features in these scenarios creates certain redundancies, making it difficult to enhance performance. Conversely, in the Cost First scenario, since Cost has a greater impact on Reward, incorporating token length and LLM inference time into the edge features better aids prediction.

Therefore, given the lack of a significant performance boost and the added complexity of incorporating these features, we choose to keep the model simple by only considering the performance and cost features. All the discussions and experimental results mentioned above have been updated in Section C.1 of the Appendix and Table 24 of the revised PDF version.

Impact on Reward with different edge features. We compared the Reward metric under four different combinations of edge features across three scenarios.

Q6. Could you elaborate on how the framework would handle dynamic updates to LLM capabilities, such as when models are fine-tuned or updated? Would this require rebuilding the entire graph?

Response: Thanks for your insightful question. When LLM models are fine-tuned or updated, we can view these LLM models as new LLMs. In addition, the strong inductive power enables GraphRouter to effectively learn new LLM nodes' embeddings by modeling their few-shot interactions with task nodes and query nodes. In this process, we can simply add on the previous graph, so there is no need to rebuild the entire graph. The only update required is the new LLM nodes' connections with the few-shots queries. In [section 4.2 of our paper and Table 5], we demonstrated through experiments that our method generalizes well to new LLMs. Therefore, when models are fine-tuned or updated, GraphRouter can seamlessly adapt by incorporating these changes without the significant computational overhead.

We greatly appreciate you raising this question, as it aligns with a similar point mentioned by Reviewer pHun in Q2. This has also inspired us to explore a method for addressing dynamic updates through a zero-shot approach, requiring only minor modifications to the edges in GraphRouter. Specifically, we draw an analogy to recommendation systems, which rely on the social networks of new and old users to address the cold start problem of new users. In our task, we considered using the addition of inner edges within LLM nodes to enable zero-shot capabilities for new LLMs (dynamically updated) in GraphRouter. We connect LLM nodes within the same size or the same LLM series (such as those within the LLaMA series). Based on this, we followed the training data settings described in section 4.2, but without the need for few-shots data, to conduct zero-shot LLM selection experiments. We set the first six LLMs in Table 3 as visible during training, while the last four are used solely for zero-shot experiments. We connected nodes within the same size or LLM series and conducted zero-shot experiments using this version of GraphRouter in the Balance scenario, as shown in the table below. We observed that under the zero-shot setting, GraphRouter (zero-shot) not only had an extremely low time cost but also approached the reward of the strongest baseline, C2MAB-V, in this scenario. All these findings highlight the great potential of leveraging the inner links of LLM nodes to enhance the zero-shot capabilities of GraphRouter. All the discussions and experimental results mentioned above have been updated in Section D of the Appendix and Table 29 of the revised pdf version.

Dear Reviewer c8tQ,

We recognize that the timing of this discussion period may not align perfectly with your schedule, yet we would greatly value the opportunity to continue our dialogue before the deadline approaches.

We hope that our responses and additional experiments have effectively addressed your concerns. We truly appreciate all the valuable advice we have received. Could you let us know if your concerns have been adequately addressed? If you find that your concerns have been resolved, we would appreciate it if you could reconsider the review score.

Thanks!

Q1. The studied setting is not quite realistic. The proposed method constructs a graph with task, query, and LLM nodes. For each query, it may select a different LLM to answer the query, which is quite unrealistic in practice.

Response: Thank you for highlighting concerns regarding the studied setting. We answer your questions as follows:

[LLM selection is realistic in real-world applications]. As mentioned in the introduction [lines 35-45] , as the number of LLMs available rapidly increases, choosing the right LLM to use is also becoming a critical issue. This claim is supported by the increasing number of LLM API servers, such as Together AI [1] and GroqCloud [3], and these servers usually host multiple different LLM models (usually larger than 10). Moreover, as introduced in section 2.2 and involved in existing research (Ding et al., 2024, Chen et al., 2023, Dai et al., 2024), different LLMs excel at different queries and tasks. Therefore, the selection of LLMs naturally becomes a practical issue to consider, and designing an efficient LLM selection mechanism capable of addressing various questions and tasks is of practical significance. Furthermore, in the industry, query-level routing is very common. For example, for models, such as ChatGPT [2], which usually serve over 100 million users concurrently, [2], it is impossible to process all user queries on a single model at the same time; it inherently requires a router to distribute queries across multiple LLM models.

[Modeling LLM selection with graph is natural and useful]. As discussed in [section 2.2], the interactions between tasks, queries, and LLM nodes provide critical contextual information regarding the performance and cost of how LLMs solve different queries and tasks. This heterogeneous contextual information can be naturally modeled as a heterogeneous graph (Peng et al., 2019; Hu et al., 2020). Moreover, in our experiment, GraphRouter substantially surpasses existing routers, delivering a minimum performance improvement of 12.3%. This validates the effectiveness of modeling pass interaction data with heterogeneous graphs, and framing the LLMs selection problem into predicting the most probable edge in the graph.

We hope that the further introduction provided above can clarify the importance of our studied setting. We also kindly anticipate that you, together with the other reviewers who have recognized the significance and authenticity of this issue, will advocate for collective efforts toward its resolution and development.

[1] https://docs.together.ai/docs/introduction

[2] https://explodingtopics.com/blog/chatgpt-users

[3] https://console.groq.com/docs/overview

Q2. The model performance especially cost may vary a lot on different hardware settings. It is also unrealistic to make sure the real hardware used can align with the numbers in the training data. And we cannot curate the training data for different hardware settings.

Response: We appreciate that the reviewer pointed out the concerns on model performance. However, we believe that the reviewer might have misunderstood our definitions of model performance and cost. We have also emphasized our definitions in multiple sections of the paper [sections 3.2 and 3.4]. We also follow the settings of existing work (Ding et al., 2024, Chen et al., 2023, Dai et al., 2024), and define model performance as the average quality of the responses under task-specific metrics [lines 345-348], and cost as the average price of tokens used by the LLM to answer queries [lines 348-350]. Our definitions of model performance and cost do not involve training data for different hardware settings. Moreover, other reviewers have already recognized the reasonableness of our setting, we kindly hope that the reviewer will also acknowledge the significance and practicality of our problems.