GraphRouter: A Graph-based Router for LLM Selections
Authors: Tao Feng
Yanzhen Shen
Jiaxuan You
Venue: ICLR 2025 Poster (2025)
Keywords: LLM selection
Graph-based router
Contextual interactions
New LLM settings
Primary area: learning on graphs and other geometries & topologies
TL;DR: We propose GraphRouter, a novel graph-based router that improves LLM selection for diverse tasks by leveraging an inductive graph framework to optimize performance-cost trade-offs and generalize across various models and scenarios.

Abstract: The rapidly growing number and variety of Large Language Models (LLMs)
present significant challenges in efficiently selecting the appropriate LLM for
a given query, especially considering the trade-offs between performance and
computational cost. Current LLM selection methods often struggle to generalize
across new LLMs and different tasks because of their limited ability to leverage
contextual interactions among tasks, queries, and LLMs, as well as their depen-
dence on a transductive learning framework. To address these shortcomings, we
introduce a novel inductive graph framework, named as GraphRouter, which
fully utilizes the contextual information among tasks, queries, and LLMs to en-
hance the LLM selection process. GraphRouter constructs a heterogeneous
graph comprising task, query, and LLM nodes, with interactions represented as
edges, which efficiently captures the contextual information between the query’s
requirements and the LLM’s capabilities. Through an innovative edge prediction
mechanism, GraphRouter is able to predict attributes (the effect and cost of
LLM response) of potential edges, allowing for optimized recommendations that
adapt to both existing and newly introduced LLMs without requiring retraining.
Comprehensive experiments across three distinct effect-cost weight scenarios have
shown that GraphRouter substantially surpasses existing routers, delivering a
minimum performance improvement of 12.3%. In addition, it achieves enhanced
generalization across new LLMs settings and supports diverse tasks with at least a
9.5% boost in effect and a significant reduction in computational demands. This
work endeavors to apply a graph-based approach for the contextual and adaptive
selection of LLMs, offering insights for real-world applications.

Published as a conference paper at ICLR 2025
GRAP HROUTER: A GRAPH-BASED ROUTER FOR LLM SELECTIONS
Tao Feng, Yanzhen Shen, Jiaxuan You Department of Computer Science University of Illinois Urhana Champaign Urbana, II., USA ftaofeng2, yanzhen4, Hlaxuan]eilling inoiiiiiiiii edu
ABSTRACT
The rapidly growing number and variety of Large Language Models (LLMs) present significant challenges in efficiently selecting the appropriate LLM for a given query, especially considering the trade-offs between performance and com- putational cost. Current LLM selection methods often struggle to generalize across new LLMs and different tasks because of their limited ability to leverage contextual interactions among tasks, queries, and LLMs, as well their dependence a transductive learning framework. To address these shortcomings, we introduce a novel inductive graph framework, named as GraphRouter, which fully utilizes the contextual information among tasks, queries, and LLMs to enhance the LLM selection process. GraphRouter constructs heterogeneous graph comprising task, query, and LLM nodes, with interactions represented as edges, which ef- ficiently captures the contextual information between the query's 's requirements and the LLM's capabilities. Through an innovative edge prediction mechanism, GraphRouter is able to predict attributes (the effect and cost of L.MM response) of potential edges, allowing for optimized recommendations that adapt to both existing and newly introduced LLMs without requiring retraining. Comprehensive experiments across three distinct effect-cost weight scenarios have shown that GraphRout e substantially surpasses existing routers, delivering a minimum performance improvement of 12.3%. In addition, it achieves enhanced generaliza- tion across new L.LMs settings and supports diverse tasks with at least a 9.5% boost in effect and significant reduction in computational demands. This work endeav- ors to apply a graph-based approach for the contextual and adaptive selection LLMs, offering insights for real-world applications. Our codes for GraphRoutez is released at https: : / //gtthb..... com/ ulab- wdic/Grraphoute
INTRODUCTION
The field of Large Language Models (LLMs) advancing quickly, offering an increasingly diverse range of models that vary in size, functionality, and computational demands (Bang, 2023; Liu et al., 2023). Although larger models tend to deliver better performance, their high computational costs make them inefficient for many less complex tasks (Snell et al., 2024; Chen Varoquaux, 2024). Additionally, LLMs demonstrate varied performance across different types of queries and tasks (Ahmed et al., 2024; Zhang et al., 2024), especially with the development of domain-specific LLMs (Singhal et al., 2022; Luo et al., 2022). These challenges make it difficult to recommend the optimal LLM services users that strike the balance between performance and cost for their specific needs. Therefore, our paper aims to raise attention to this pressing research question: Given the vast and continuously evolving landscape of L.MMs, how to recommend appropriate LLMS for various user queries with different implied tasks?
Existing researchers have proposed to develop a router to assign a specific LLM to each user query. Hybrid LLM (Ding al., 2024) trains binary score router function to determine whether to select a small LLM or a large LLM for specific query. Although it balances cost and performance, is limited just two LLMs, which falls short of addressing the real-world demand for wide range of LLMs. Some other studies (Dai et al., 2024; Chen et al., 2023) have introduced more advanced router models to address the challenge of selecting among limited set of L.MSs (typically 3 to 5). More

---

Published as a conference paper at ICLR 2025
Table Comparison of GraphRouter with existing methods from three perspectives: contex- tual info, generalization to new LLMs, and multi-task support. Compared to other approaches, GraphRoutez introduces an inductive graph framework that fully leverages contextual information, enabling it to generalize to new LLMs and adapt to a variety of tasks.
Method
Contextual Info
Generalization to New LLMs
Multi-task Support
Hybrid LLM (Ding et al., 2024)
Indices
X
FrugalGPT (Chen et al., 2023)
LLM name
X
C2MAB-V (Dai et al., 2024)
One-hot embedding
X
GraphRouter
Graph-based contexts
specifically, FrugalGPT (Chen et al., 2023) proposes a router model based on BERT (Devlin, 2018) to determine whether switch larger LLM or not, and C2MAB-V (Dai al., 2024) constructs bandit-based router to balance between the exploration and exploitation when choosing LLM for the user. However, as shown in Table they are still limited in the following aspects: 1) Relying solely on basic BERT-based embeddings to distinguish queries, and on names or indices to distinguish LLMs, they fail to fully leverage the contextual information from the interaction between the task, query, and LLM. This makes it challenging to achieve a router with strong generalization capabilities. 2) These methods rely on a transductive learning framework (Arnold et al., 2007; Joachims, 2003), which makes them ill-suited for real-world applications where new LLMs are frequently introduced. When new LLMS are presented, these approaches require retraining with few-shor interaction data before they can be used process that is impractical for recommending LLMs to large number of users in real-time; 3) They train a dedicated router for each specific task, greatly increases the computational overhead and complexity in real-world applications when multiple tasks are present.
To address these challenges, we introduce GraphRouter, a graph-based router for LLM selection. GraphRouter utilizes inductive graph framework to effectively leverage contextual information, allowing it to generalize to new LLMs and adapt to diverse tasks. Specifically, to fully utilize the contextual information for different queries and tasks, GraphRouter constructs a heterogeneous graph that contains three types of nodes; task node, query node and LLM node. The interaction information between them is represented as edges in a graph. For instance, the reward (including performance and cost) of an LLM responding to a query is modeled as an edge between the query node and the L.M node. Then, we are able to transform the task of predicting the cost and performance of an LLM--qurr pair to an edge prediction task. After forecasting the properties of the edges, we recommend the most suitable LLM to the user based on their preferences for performance and cost. In addition, in real-world scenarios, new LLMs are frequently developed, should also have the ability to accommodate these evolving models. In order to make GraphRout generalizable to new LLMs, make efforts in two key aspects. For the input, we utilize a generative LLM such as GPT-4o to generate a descriptive text for each LLM, outlining key details such as its strengths, token pricing, and context length. Based this, we derive an initial embedding for each I.M using a moderate-size pre-trained language model (e.g, BERT (Devlin, 2018)). This approach offers an advantage over directly using one-hot encoding, as it enables to generate inductive and more informative initial embeddings for new LLMs. For the GraphRout er model, we further develop a heterogeneous GNN that aggregates information from neighboring nodes of different types; given few-shot data, we verified that a trained GraphRouter can generalize to new LLM nodes without retraining.
effective framework
In summary, our main contributions are as follows:
To the best of our knowledge, we are the first work to build router for L.M selections from the graph perspective, which gives new insight to graph-enhanced LLM research.
We propose an inductive graph framework that fully leverages contextual information among tasks, queries, and LLMs, enabling it to generalize to new LLMs and adapt to variety of tasks without retraining.
In three experiment settings with different performance and cost tradeoffs, GraphRouter outperformed the baseline models by at least 12.3%. Furthermore, in scenarios where new LLMs are introduced in the testing data, our method not only saves significant training time but also improves performance by at least 9.5% compared to the baselines.
2

---

Published as a conference paper at ICLR 2025
Task
Query
LIM
Response
Performance Cost
Task Query
Alpaca
Give three for staying healthy?
LLaMA-3 (7b)
Nutrition plays..
0.19
103 x
Router
Choose
Alpaca
What are the three primary colors?
LLaMA-3 (70b)
Red, bluc, and yellow.
0.21
467.1
Response
GSM8K
How many musicians total...?
LLaMA-3 (7b)
There are 55...
0.00
250.1
Effect/ Cost
GSMBK
Hlw many seats ure there...?
?
User
LLM,
Figure 1: Overview of GraphRouter's 's LLM selection process. As depicted in the left section, the LLM selection process begins with the user inputting a query that belongs to a certain task. When the router receives the query and the task, it will analyze the input and choose the most appropriate LLM, for generation. Then, LLM,, is used to generate the response. In the end, this response, along with the measured effectiveness and cost, is returned the user. The right side of the figure illustrates example interaction records, which contain contextual information like task, user query, selected LLM, response, performance, and cost, in a table. These contextualized data are then utilized train the router.
2 GRAPHROUTER: GRAPH-BASED ROUTER FOR LLM SELECTION
2.1 PRELIMINARIES
We introduce the LLM selection problem in this section. As shown in the left part of Figure the process involves multiple steps, with the router being the most critical component. The router first receives the user query containing task information. Its goal is choose an appropriate LLM based on the incoming information in the user query to ensure optimal performance and minimal cost (LLM API cost). After calculation, the router chooses suitable LLM,, to answer the user query. Finally, the response is returned to the user with its performance and cost. Such interaction process generates rich contextual data, which contains information on tasks, queries, selected LLM, response, performance, and cost. We organize the data in table, as shown on the right of Figure 1.
2.2 MOTIVATING EXAMPLES
Traditional LLM selection methods (Ding et al., 2024; Chen et al., 2023; Dai et al., 2024) often only use ID or name information to model LLM information, which does not effectively utilize the contextual information (introduced in Sec 2.1) generated from the interaction between LLM and query. Here, we use some examples to illustrate the importance of contextual information. (1) We first argue that contextual information is important because it captures the variance in the ability of different LLMs to respond to diverse queries. We can first observe from Figure 3 that the performance of different LLMs in response to queries can differ significantly. Therefore, understanding the performance patterns of how LLMs handle queries is crucial for LLM selection, and these patterns embedded in the contextual information between the queries and the LLMs. In addition, from Figure 2, we can also observe that smaller LLMs sometimes outperform larger LLMs on certain queries. Even if we have unlimited budgets and can blindly rely on the largest LLM at high cost, we still may not achieve optimal performance. This also emphasizes the importance of capturing the varying capabilities of LLMs in handling queries based on contextual information. (2) We also claim that the importance of contextual information lies in its ability to capture the differences in how single L.LM responds to queries across different tasks. Through Figures 3 and 4, we can observe that certain LLMs exhibit significant differences in their performance across two different tasks, such as Mixtral-8x7B (Jiang et al., 2024). These two examples indicate that the performance may vary greatly on different LLMs and tasks. This suggests that in addition to understanding the capabilities of LLMss. the router must also understand the differences and similarities of each task. However, as these critical attributes of LLMs and tasks could not be sufficiently represented by their names IDs, an effective use of the contextual information that encompasses the interaction between task, query, and L.M is needed.
3

---

Published as a conference paper at ICLR 2025
1.2
x g 2 PRREE
Performance
Porformance
class 1770)
mutrrl
Hoathe (FN)
Leadtt 1780)
LLose, Mottr (86)
Haans brarte Orins
Unomna Seae 170%
Alaat, (76)
Be
Laakte (767
LLatE Shutto ran
Llaen Muree (ney
120
Figure 2: The probability dis- tribution of a small LLM (LLaMA-3 (7b)) having a bet- ter performance value than large LLM (LLaMA-3 (70b)) by t on the Alpaca dataset, where means the difference in performance between the small LLM and the large LLM and & E [-1,
Figure 3:
Distribution of Figure 4: Distribution of the performance of different LLMs responding to queries on the SQUAD task. In par- ticular, the performance of LLMs of varying sizes is dis- violin plot and the dot in each distribution is the me- dian performance.
the performance of different LLMs in response to queries on the Alpaca task. Specifi- cally, we present violin plot illustrating the performance of ten LLMs of varying sizes and the dot in each distribution is the median performance.
played
Train Stage Query, Xode
1-th GNN Layer
Task
Query
LIM
Performaner
Coot
Task Nade
LLM Node (observed)
All
Alpaca
91
LLAMN- (7b)
0.19
103.8
Concal
Cost Edge)
Alpacs
4:
11 aMMMA-1 (708)
0.11
4611
Transtate
[
Graph
AGG
Softmax
GSMAK
4a
3 (7b)
0.00
230,1
Fest Stage
Aurt (new)
peoduct
GSM8K
Figure Overview of GraphRouter methodology. GraphRouter first converts the interaction data among tasks, queries, and LLMs into graph. Specifically, illustrated on the right side, tasks, queries, and LLMs from the left table are represented as task nodes, query nodes, and LLM nodes, respectively. Moreover, their relationships derived from the interaction data modcled as edge features. With this structure, we leverage a GNN to embed both node and edge features, ultimately producing the probability distribution of the selected LLM. 2.3
GRAPHROUTER FRAMEWORK
Method Overview. As shown in Figure 5, GraphRouter first transforms the interaction data among tasks, queries, and LLMs into a graph. Specifically, as shown the right side of Figure 5, we model tasks, queries, and LLMs in the left table task nodes, query nodes, and LLM nodes, while the relationships derived from the interaction data are represented by edge features. We apply GNN to embed the node and edge features and use them for training and testing.
Initialize node/edge features. As shown Figure 5, we have three types of nodes (task node h!(). query node (1) and LLM node h(ll) and two types of edges (task-query edge Wiy and LLM-query edge Wmg)).
Given the inherent differences of task, query, and LLMs, during the initialization of their nodes, we adopt different strategies. For the initialization of task nodes, we utilize an additional LLM, such GPT-4o, to generate descriptions of the tasks, and then encode the description to obtain its embedding e,. More specifically, we take the average output token embedding after feeding the description into the moderate-size pre-trained language model (PLM), such BERT (Devlin, 2018). The initialization of query nodes is also obtained by embedding the query eg through the same PLM. As for the initialization of LLM nodes. the traditional approach often initializes directly using the name or ID the LLM, which not only limits its ability to generalize to new LLMs but also omits important background information. Here, we still adopt prompt- based approach. We design prompts
4

---

Published as a conference paper at ICLR 2025
for LLM describe the capabilities of each LLM. In addition, also add information about the cost of each LLM after the description. Then, similar to how we obtain the task's 's embedding, we use the same PLM to compute the initial embeddings e, of the different LLMs. All the descriptions generate for different tasks and LLMS can be found in Appendix A.
As for the task-query edges, assign value of for their initialization. For the initialization of LLM--uury edges, we jointly consider the performance and cost information in the interaction data, and assign the concatenation of performance and cost their initial features.
Predict via a heterogeneous GNN. We implement the predictive model fa over task nodes, query nodes, and LLM nodes using a heterogeneous GNN, shown in Figure 5. For aggregating different types of nodes and edges, we use heterogeneous aggregation with different learnable weights. The objective of the GNN is learn expressive node embeddings h through an iterative weighted aggregation of the local network neighborhoods. The /-th iteration of the GraphConv()). or the node embeddings update of the [-th layer, is represented as:
h (1) = U.") CONCAT MEAN MMEA((REEL(WW) (RELU hy (ihh""")))).
(I)
h CONCAT COONCAMMMA(LLLLUWWWEeEEe.... (RELU((I(EV..EW) w Il (1-1) u
E ((o))) h h((11)
h _ CONCAT MEEAN (RELU(W". Wh
D). M((m)) h b)
(3)
where h is the node embedding after iterations, h hg hhm have been initialized as h(0)) hg (0) (0) Cc, Cm respectively explained above, and (v)) denotes the direct neighbors of 1 ly E Va, u E Vi] indicates the message type (whether from task to query, or LLM to query), and U(), (( are learnable parameters. In addition, WIjrEV..me V.] indicates that different edge types correspond to different edge features; specifically, if from task to query, represented as Wey, and from LLM to query, it is represented as Wmq-
Following the update of the task, query, and LLM node embeddings, we obtain the query-task
combined embedding as h = (LP((CNCAT((D)"" h We model the LLM selection problem edge prediction problem and generate training data in the following way. We first determine the best LLM for each query in the training set based on the performance (best reward described in See 3.4) achieved by different LLMs and then set the edge labels of the query other LLMs and the edge label of the query to the best LLM to 1. As such, LLM prediction can be made through EdgePred(()) in the form of Mogite = MEAN DoT(h h )n We have summarized the detailed training process of GraphRouter in Algorithm whose details shown follows. In addition, in the testing of Graph we identify the L.M node that has maximum edge logits with the query node as the best LLM, which can be computed as
" = arg max m
(EdgePred(haps Hm))
GraphRouter for new LLMs setting. Traditional routers cannot generalize to new LLMs directly under few-shots settings, they require retraining through interactions with the query for each new LLM. This is inadequate for keeping with the rapidly evolving changes in LLMs in the real world. To test our framework and baselines under this real-world setting, following (Cao al., 2023; Fey et al., 2023), we construct an auxiliary dataset with the interaction data of the new LLMs on queries sampled uniformly from the training set. This auxiliary dataset is not involved in the training process but serves as few-shot examples during the testing phase.
3 EXPERIMENTAL SETUP
3.1 DATASETS AND LLM DESCRIPTIONS
We select data from four different types of tasks, whose statistics are summarized in Table 2.
5

---

Published as a conference paper at ICLR 2025
Algorithm 1 Training of Graph
Require:
Dataset Durin {(x. y) A parameterized heterogeneous GNN fa. Task-query edge weights Wey and LLM--uuer edge weights Wmg- Number of GNN layers L. 1: Initialize the embeddings of task nodes, query nodes, and LLM nodes, 11( nig ,.mm using PLM.
2: for each iteration do
3: M <- SampleMiniiddgeBachhh (Priin)
4: Mask the edges in Disain that are in M and obtain the labels of the edges in Thl
5: for / 1 to L do
hg hg hm + -GraphCon(hh" (1-1) hgg
((-1)) h Weg: Wmgg with So
6:
7: Diggtts < EdgePred(h..) hng) hm with f&
8:
Backward (Criterion (Nlggss: {(V) jETTo } E eM)
Alpaca (Taori et al., 2023) is a hybrid question-answer (QA) dataset containing 52k samples used for fine-tuning the Alpaca model. The dataset is automatically generated by the self-instruct (Wang et al., 2022) framework, which iteratively prompts the language model to generate new training instances given a few manually written instructions.
GSM8K (Cobbe et al., 2021) evaluates the model's s ability for multi-step mathematical reasoning with 8.5k linguistically diverse grade school math word problems.
SQUAD (Rajpurkar, 2016) is crowdsourced reading comprehension dataset based on Wiki articles. It contains over 100k QA pairs connected to over 500 articles.
Multi-News (Fabbri et al., 2019) is a benchmark on multi-document summarization. It consists of 56k news articles and summary pairs where the news articles are extracted from newser.com and the summary is written by professional editors.
Furthermore, we introduced 10 LLMs of varying sizes into problem, with their statistics shown in Table 3. All of the LLMs and their token costs here are accessed through the Together API
Table 2: Overview of Datasets.
3.2 DATA PREPROCESSING AND SPLITTING
Given the above dataset and LLMs, construct multi-task interaction dataset described in Sec 2.1. Specifically, we combine all the datasets of four tasks together first For each query, we utilize ten LLMs in Sec 3.1 to answer and obtain the corresponding response. Then the response is compared with its ground truth to get its performance using the metric of each task described in Table 2. Furthermore, the cost calculated with the number of input tokens and output tokens and the cost of different LLMs in Table 3. Here we utilize GPT-2 as in (Chen et al., 2023) to calculate the number of tokens.
Dataset
Task Type
Metric
Cases
Alpaca
Hybrid QA
Fi
600
GSM8K
Reasoning
Accuracy
600
SQUAD
Reading Comprehension
Multi-News
Summary
After obtaining the multi-task interaction dataset, we split the dataset according to different experi- mental settings. We mainly have two major settings. The first is the standard setting, where all LLMs visible in both the training and test sets, with some is divided into training, validation, and test sets in a ratio of 70% 10%: 20%, based on different queries. In the case of new LLM setting, assume that the first six LLMs in Table 3 are observable, while the remaining four are new LLMs. Therefore, based on the standard setting, we first remove data related to the latter four LLMs from the training and validation sets, while keeping the test as it is in the standard setting. Furthermore, following (Cao et al., 2023; Fey et al., 2023), we construct an auxiliary dataset with the interaction data of the four new LLMs on 80 queries sampled uniformly from the training set. This auxiliary dataset is not involved in the training process but serves as few-shots during the testing phase.
queries appearing in the test set. The data

---

Published as a conference paper at ICLR 2025
3.3 BASELINE METHODS
We compare GRAPHROUTER model with the following baselines. GraphRoutez with two rule-based baselines:
We first compare
Largest LLM always selects the largest LLM available.
Smallest LLM always selects the smallest LLM available.
Then compare GraphRoute: with prompt-based baseline:
Prompt LLM incorporates the query, candidate models, and objectives (e.g., prioritizing effective- ness) directly into the prompt, and feeds into an external LLM (c.g., GPT-4) to select the most suitable LLM from pool of candidates.
Further, GraphRouter is compared with three representative routers for LLM selection:
Hybrid LLM (Ding et al., 2024), when given a small LLM and a large LLM, trains a pre-trained language model assign queries the small large model. We use LLaMA-2 (7b) and Llama-
3.1-Turbo (70b) as our small and large LLM respectively, as they are the smallest and the largest LLM available. We also replace DeBERTa (He al., 2020) with RoBERTa (Liu, 2019) as the pre-trained language model and observe better performance.
FrugaIGPT (Chen al., 2023) utilizes a pre-trained language model predict the score of the generation result of all LLMs given query, and then selects the LLM with the highest score within given cost. We also use RoBERTa (Liu, 2019) as the router's backbone model.
C2MAB-V (Dai et al., 2024) uses bandit- based model for LLM selection, which regards each LLM as an arm and implements an exploration mechanism to search for a better solution.
Finally, we set up gold baseline the optimal solution for LLM selection. The purpose of setting up the baseline to see how far GraphRouter from the optimal solution.
Oracle defines the theoretical upper bound of the reward, where each query has been routed to the optimal LLM via oracle information.
3.4 METRICS
Table 3: Statistics of different LLMs and their costs on Together API.
utilize three metrics to evaluate the performance of GraphRouter and baselines.
LLM
Size
Cost per 1M tokens
Performance is to evaluate the average quality of the responses across different queries given by each method, which is introduced in Sec 3.2 and Table 2.
LLaMA-3 (7b)
0.2
Mixtral-x7BB
56b
0.6
NousReseurch
34b
0.8
Cost evaluates the average LLM inference cost when responding to the queries, which is described in Sec 3.2.
I.aMA (7h)
7h
0.2
Mistrul-7b
0.2
I.LaMA-3 (70b)
LLaMA-3-Tubbo (8b)
0.2
Reward used measure how well method bal- ances performance and cost. Different users may have varying levels of emphasis on performance and cost. Therefore, we define three scenarios: Per- formance First, Balance, and Cost First, to cor- respond to situations where users prioritize high performance, value both high performance and low cost equally, or prioritize low cost, respectively. Specifically, to eliminate the influ- ence of scale, Reward = a Per formance B Cost. For the three scenarios, we set the values of a and B to (1, 0), (0.5, 0.5), and (0.2, 0.8), respectively.
LLaMA-3-Turbo (70b)
706
0.9
Llama-3.1 3.1-ubbb (70b)
706
0.9
Qwen- 1.5 (72h)
72b
0.9
first normalize both performance and cost. Then, define the score as
3.5
IMPLEMENTATION DETAILS
In the training stage, set the graph neural network two-layer graph attention network, with 32-dim hidden dimension. The batch size is 32, and the max training epoch is set to 1000. We use Adam optimizer (Diederik, 201 4) for model training and gradually decay the learning rate from 1e-3
7

---

Published as a conference paper at ICLR 2025
Table 4: Comparison of Various Methods Multi-task Interaction Dataset across Three Distinct Performance-Cost Weight Scenarios. Bold and underline denote the best and second-best results, All datasets are evaluated Performance, Cost, and Reward. Each number is the average of multiple rounds.
Scenario
Performance First
Balance
Cost First
Performance
Cost
Reward Performance
Cost
Reward
Performance
Cost
Reward
Largest LLM
0.431
0.871
0.431
0.431
0.871
-0.220
0.431
0.871
-0.701
Smallest LLM
0.279
0.031
0.279
0.279
0.031
0.124
0.279
0.031
-0.009
Prompt LLM
0.474
0.812
0.474
0.285
0.0551
0.115
0.283
0.108
-0.03
Hybrid LLM
0.510
0.871
0.510
0,470
0,451
0.009
0.276
0.151
-0.066
FrugalGPT
0.517
0.671
0.517
0,400
0.072
0.164
0.411
0.031
0.057
C2MAB-V
0.479
0.871
0.479
0.423
0.031
0.196
0.279
0.031
0.031
GraphRouter 0.539
0.725
0.539
0.448
0.031
0.209
0.446
0.031
0.064
Oracle
0.588
0.586
0.588
0.504
0.040
0.231
0.483
0.031
0.072
to O with LambdaLR scheduler. We implement our proposed method using PyTorch2 and PyG', and all the experiments are conducted on a single NVIDIA A100 Tensor Core GPU. As for LLMs, we rely on API calling from Together AI* to obtain responses.
4 EXPERIMENTAL RESULTS
4.1 COMPARISON WITH EXISTING BASELINES.
We compare GraphRouter with seven baselines in three scenarios in Table 4. We can observe that GraphRouter consistently and substantially surpasses existing routers, delivering minimum effect improvement of 12.28% on metric Reward compared to the strongest baselines. Additionally, observe that GraphRouter achieves at least 88.89% of the optimal solution (Table 4, row Oracle), further demonstrating the superiority of our framework. On the other hand, compared with the two rule-based LLM, GraphRouter achieves a better trade-off between Performance and Cost, therefore achieving a higher effect on Reward. Analyzing the effect of Prompt LLM, Hybrid LLM (Ding al., 2024), and FrugalGPT (Chen et al., 2023), we demonstrate that without sufficient contextualized information, even LLM and trained moderate-size LM struggle to understand the query and candidates L.M effectively, even if we ignore their high inference costs. These results validate that effective usage of contextual information is crucial for selecting the optimal LLM.
4.2 GENERALIZATION ABILITY TO NEW LLMS
To validate the generalization ability of GraphRouter when facing new LLMs, we conduct experiments described in Sec 3.2 in scenario Balance. To compare with other baselines, we add the auxiliary dataset into their training dataset. Specifically, we compare GraphRouter (few-shots) with HybridLLM, FrugalGPT, C2MAB-V, V, and GraphRouter (trained) on Reward and time cost (training time + inference time). We report our results in Table 5. We can observe that in comparison to the most costly C2MAB-V (Dai et al., 2024), GraphRouter (few-shots) not only achieves substantial performance improvements in Reward by almost 10% but also greatly reduces Time Cost by over 99%. The amount of Reward Improvement and Time Cost Reduction has significantly surpassed those of other baselines. Additionally, compared to GraphRoutes (trained), our approach significantly reduces time cost with only a slight performance loss. These observations demonstrate that GraphRouter e is both effective and efficient in generalizing to new LLMs.
4.3 ABLATION STUDIES
How does GraphRouter perform with varying sizes of GNN? The size of GNN is an important factor to consider when designing GNN algorithms. It not only affects the performance of the GNN but also introduces additional computational overhead if the size is too large. To find an optimal GNN size for GraphRouter, we explored sizes ranging from 16 to 80, as shown in Figure 6. As depicted
https//pytorch..org/
http://ppww..ocmmmmm////cctt/tce hups://www.togetheraai/
8

---

Published as a conference paper at ICLR 2025
Table 5: Comparison of methods in the few-shot setting on Reward, Time Cost, and the corresponding percentage Reward improvements and Time Cost reduction rate, relative to the most costly method (C2MAB-V (Dai al., 2024)).
Method
Reward
Reward Improvement %)
Time Cast
Time Cost Reduction %6)
HybridLLM
0.01
-94.11
273.45
49.57
PrugalGPT
0.171
-9.52
63.15
88.35
C2MAB-V
0.189
0.00
542.25
0.00
GraphRouter (few-shots)
0.207
9.52
3.00
99.45
GraphRouter (Trained)
0.219
15.87
30.00
94.47
0.20
0.18
0.18
Reward
0.14
0.16
Reward
0.16
0.12
0.14
0.10
0.12
0.08
0.10
16
32
48
64
80
Size
Number of Layer
Figure 6: The Reward of GraphRouter varies with the input size the GNN. Here, we set in-channels equal to out- channels. As observed in the figure, the Re- ward of GraphRouter initially increases and then decreases, achieving the highest Re- ward when the size is 32.
Figure 7: The Reward of GraphRouter varies with the number of GNN layers. As shown in the figure, the Reward of GraphRouter initially increases and then de- creases, reaching its highest Reward when the number of layers is 2.
in the figure, the Reward of GraphRouter initially improves as the size increases, reaching its peak size of 32, after which it starts to decline.
What is the impact of different numbers of GNN layers on GraphRouter's effectiveness? The number of GNN layers has significant impact on the expressiveness the GNN. A shallow GNN struggles to learn deep contextual information, whereas an overly deep GNN can lead to issues such as over smoothing and overfitting. Moreover, increasing the number of layers also raises the computational cost. To determine the optimal number of GNN layers for GraphRouter, conduct an exploration with the number of layers ranging from 0 to as shown in Figure 7. As depicted in the figure, the Reward of GraphRouter initially improves with more layers but then declines, achieving its highest Reward when the number of layers
ADDITIONAL RELATED WORKS
LLM selection. With the scaling of the number of parameters in LLMs, their inference cost also rapidly increasing. To optimize inference cost, several works have proposed using model switching to mitigate this issue. When given small and a large LLM, Zhu et al. (2023) fine-tune pre-trained language model the model router to predict whether using a small LLM is sufficient. HybridLLM (Ding et al., 2024) improves on this by transforming training data to compensate for the influence of imbalanced labels. Other approaches move beyond choosing between two LLMS and introduce settings where multiple LLMs present. Chen et al. (2023) train the router to predict the reliability score given query and an LLM index. Šakota et al. (2024) generalize the router's s training objective to accommodate the scenario where additional cost or performance constraints exist. Stripelis et al. (2024) examine the effectiveness of lighter routers built upon the k-nearest neighbors algorithm or Multilayer Perceptrons. C2MAB-V (Dai et al., 2024) employs a bandit-based and incorporating exploration mechanism
model for LLM selection, treating each LLM as to find an improved solution. Different from previous approaches, where the router only learns from query-model interaction, our GRAPHROUTER fully utilizes the information in the training data by

---

Published as a conference paper at ICLR 2025
jointly modeling the query-moder query-quer and model-model relationship. This allows to learn effective representations for tasks, queries, and models, enabling better generalizability.
Graph for modeling relationships. Graphs have demonstrated great potential in modeling complex relationships (Fey et al.. 2023; Cao et al., 2023; Gao & Xu, 2020; Chen et al., 2022: Wu et al., 2022b; Yang et al., 2021). Solving relational data with graphs often involves extracting nodes and edges from the data, and then modeling their relationships with embeddings. 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. With the advancement of deep learning, graph neural networks (GNNs) have become the more popular approach for researchers to model relationships within data. They are also found to have widespread application in fields such as recommendation systems (Min al., 2022) and social networks (Wu et al., 2020). GNNs can be broadly classified into Message Passing Neural Networks (MPNNs), which include models like GCN (Kipf & Welling, 70, GraphS AGE (Hamilton et al., 2017), and GAT (Veličković et al., 2017), as well as non-MPNN architectures (Wu et al., 2022a; Ying et al., 2021). Additionally, Heterogeneous Graph Neural Networks (HeterGNNs) (Peng et al., 2019; Hu ct 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. In recent years, researchers have begun exploring the zero- shot or few-shot capabilities of GNNs (Fey al., 2023; Cao et al., 2023; Gao & Xu, 2020; Chen et al., 2022), aiming to address more complex real-world challenges, such as the cold start problem in recommendation systems. Building on these studies, we incorporate GNNs' powerful ability to represent contextual heterogeneous relationships and their zero-shot capabilities into the LLM selection problem.
6 CONCLUSION AND DISCUSSION
(1) Conelusion. We present GraphRouter, graph inductive framework for LLM routing during inference with multiple LLMs. This work is the first to address the LLM routing problem by reframing it as an edge prediction task between query nodes and LLM nodes. Using graph structure, we fully capture contextual information from prior interaction data to learn effective task, query, and graph representations. Through our experiments on the combined dataset from four open- domain QA datasets, and with three different application scenarios, we demonstrated the superiority of GraphRoutes er over competitive LLM selection baselines and showed that our framework is on par with the ideal "God's-eye view" solution. Beyond traditional settings, we also tested our framework in a more challenging setting where new LLMs were introduced during test time, and we demonstrated GraphRouter's strong generalization ability compared previous baselines. We hope that GraphRouter, along with our approach of incorporating interaction data through graphs, will facilitate future research on LLM routing. (2) Limitations. This work primarily serves as exploratory work to validate the idea of how modeling past interaction data in graph could enhance the process of LLM selection. We acknowledge that leveraging more complex graph signals, such as paths, or the taxonomy of L.LMs (e.g., family trees like LLaMA2 - LLaMA3 - LLaMA3.1 (Touvron et al., 2023)), could further improve GraphRout ertoutr''s performance, but that is beyond the scope of this paper, and we leave it for future work. (3) Future Work. Some other interesting questions to explore in the future include: 1) When answering complex queries, prompting methods like Chain-of- Thouughh (Wei et al., 2022), Tree-of- Tree-of-Thought Thought (Yao et al., 2024) are also widely used to enhance the reasoning ability of LLMs. Given the vast number of these methods, predicting the generation result and selecting the best one remains a great challenge. On the other hand, selecting the prompting method is similar to selecting the best LLM for inference, as we are both aiming to predict the generation result based on past interaction data. As a result, it interesting to explore whether Gr GraphRout er could also be adapted to this task. 2) In a multi-agent system, it is critical to choose the appropriate L.LM for each module on a specific query and task. It would be valuable to conduct experiments on whether an inductive graph framework like GraphRouter could also excel on this challenging task, where multiple LLMs are being selected simultaneously. 3) How to enable LLMs to better understand numerical differences is direction for future consideration in modeling LLM features more effectively. Using current LLMs to model and understand numerical information like token pricing and context length (Romera- Paredes et al., 2024; Ahn et al., 2024; Imani et al.. 2023; Lewkowycz et al., 2022) is still an open question. These numerical details are significant factors affecting the performance of L.M selection.
Koomera-Paredes
10

---

Published as a conference paper at ICLR 2025
REFERENCES
Toufique Ahmed, Christian Bird, Premkumar Devanbu, and Saikat Chakraborty. Studying Ilm performance on closed-and open-source data. arXiv preprint arXi::2402.15100 2024.
Janice Ahn, Rishu Verma, Renze Lou, Di Liu, Rui Zhang, and Wenpeng Yin. Large language models for mathematical reasoning: Progresses and challenges. arXiv preprint arXi::2402.00157 2024.
Andrew Arnold, Ramesh Nallapati, and William W Cohen. A comparative study of methods for transductive transfer learning. Seventh IEEE international conference on data mining workshops (ICDMW 2007), pp- 77-82. IEEE, 2007.
Fu Bang. Gptcache: An open-souree source semantic cache for Ilm applications enabling faster answers and cost savings. Proceedings of the 3rd Workshop for Natural Language Processing Open Source Software (NLP-OSS 2023), pp. 212-218, 2023.
Kaidi Cao, Jiaxuan You, and Jure Leskovec. Relational multi-task learning: Modeling relations between data and tasks. arXiv preprint 1xXi:2333.076666 2023.
Lihu Chen and Gaäl Varoquaux. What is the role small models in the IIm era: survey. arXiv preprint arXiv:2409.068: 2024.
Lingjiao Chen, Matei Zaharia, and James Zou. Frugalgpt: How large language models while reducing cost and improving performance. arXiv preprint arXiv:2305. 05/76. 2023.
Mark Chen, Jerry Tworek, Heewoo Jun, Qiming Yuan, Henrique Ponde De Oliveira Pinto, Jared Kaplan, Harri Edwards, Yuri Burda, Nicholas Joseph, Greg Brockman, et al. Evaluating large language models trained on code. arXiv preprint 1rXiv:2107.033744 2021.
Shiming Chen. Ziming Hong, Guosen Xie, Qinmu Peng, Xinge You, Weiping Ding, and Ling Shao. Gndaa: Graph navigated dual attention network for zero-shot learning. IEEE transactions on neural networks and learning systems, 35(4):4516-457)) 2022.
Karl Cobbe, Vineet Kosaraju, Mohammad Bavarian, Mark Chen, Heewoo Jun, Lukasz Kaiser, Matthias Plappert, Jerry Tworek, Jacob Hilton, Reichiro Nakano, et al. Training verifiers to solve math word problems. arXiv preprint arXiv:21 20.:2210.1111 10.14168, 2021.
Xiangxiang Dai, Jin Li, Xutong Liu, Anqi Yu, and John Lui. Cost-effective online multi-llm selection with versatile reward models. arXiv preprint urXiv:2405.1658) 2024.
Jacob Devlin. Bert: Pre-training of deep bidirectional transformers for language understanding. arXiv preprint arXiv:1810.0450 04805, 2018
P Kingma Diederik. Adam: A method for stochastic optimization. (No Title), 2014.
Dujian Ding, Ankur Mallick, Chi Wang, Robert Sim, Subhabrata Mukherjee, Victor Ruhle, Laks VS Lakshmanan, and Ahmed Hassan Awadallah. Hybrid lmm: Cost-efficient and quality-aware query routing. arXiv preprint wxXhv:24041166188 2024.
Alexander R Fabbri, Irene Li, Tianwei She, Suyi Li, and Dragomir Radev. Multi-news large- scale multi documet summarization dataset and abstractive hierarchical model. arXiv preprint 1:Xi::/9660017440 2019.
Matthias Fey, Weihua Hu, Kexin Huang, Jan Eric Lenssen, Rishabh Ranjan, Joshua Robinson, Rex Ying, Jiaxuan You, and Jure Leskovec. Relational deep learning: Graph representation learning on relational databases. arXiv preprint arXiv:2312.04611 2023.
Junyu Gao and Changsheng Xu. Ci-gnn: Building a -ategory-instance graph for zero-shot video classification. IEEE Transactions on Multimedia, 32800)3308- 2020.
Will Hamilton, Zhitao Ying, and Jure Leskovec. Inductive representation learning large graphs. Advances in neural information processing systems, 30, 2017.
Pengcheng He, Xiaodong Liu, Jianfeng Gao, and Weizhu Chen. Deberta: Decoding-enhancee bert with disentangled attention. arXiv preprini arXiv:2006.03654, 2020.
11

---

Published as a conference paper at ICLR 2025
Ziniu Hu, Yuxiao Dong, Kuansan Wang, and Yizhou Sun. Heterogeneous graph transformer. In Proceedings of the web conference 2020, pp. 2704-2710, 2020.
Shima Imani, Liang Du, and Harsh Shrivastava. Mathprompter: Mathematical reasoning using large language models. arXiv preprint arXiv:2303.0533
2023.
Albert Q Jiang, Alexandre Sablayrolles, Antoine Roux, Arthur Mensch, Blanche Savary, Chris Bamford, Devendra Singh Chaplot, Diego de las Casas, Emma Bou Hanna, Florian Bressand, et al. Mixtral of experts. arXiv preprint arXiv:2401 0400088 2024.
Thorsten Joachims. Transductive learning via spectral graph partitioning. In Proceedings of the 20th international conference on machine learning (ICML-03), pp. 290-297, 2003.
Thomas N. Kipf and Max Welling. Semi-supervised classification with graph convolutional networks. In ICLR (Poster). OpenReview.net, 2017.
Aitor Lewkowyc7, Anders Andreassen, David Dohan, Ethan Dyer, Henryk Michalewski, Vinay Ra- masesh, Ambrose Slone, Cem Anil, Imanol Schlag, Theo Gutman-Solo, et al. Solving quantitative reasoning problems with language models. Advances in Neural Information Processing Systems, 35:3843-3857, 2022.
Junyi Liu, Liangzhi Li, Tong Xiang, Bowen Wang, and Yiming Qian. Tera-llm: Token com- pression retrieval augmented large language model for inference cost reduction. arXiv preprint arXiv:2310.15556 2023.
Yinhan Liu. Roberta: A robustly optimized bert pretraining approach. arXiv preprint
1907 1692, 2019.
Renqian Luo, Liai Sun, Yingee Xia, Tao Qin, Sheng Zhang, Hoifung Poon, and Tie- Yan Liu. Biogpt: generative pre-trained transformer for biomedical text generation and mining. Briefings in bioinformatics, bbb6)bbb 2022.
Erxue Min, Yu Rong, Tingyang Xu, Yatao Bian, Peilin Zhao, Junzhou Huang, Da Luo, Kangyi Lin, and Sophia Ananiadou. Masked transformer for neighhourhood-awwan click-through rate prediction. CoRR, 000//2001.1331..... 2022.
Hao Peng, Jianxin Li, Qiran Gong, Yangqiu Song, Yuanxing Ning, Kunfeng Lai, and Philip Yu. Fine-grained event categorization with heterogeneous graph convolutional networks. arXiv preprint 1rXiv:1906.04580 2019.
Rajpurkar. Squad: 100,000+ questions for machine comprehension of text. arXiv preprint arXiv:1606.05250, 2016.
Bernardino Romera-Paredes, Mohammadamin Barekatain, Alexander Novikov, Matej Balog, M Pawan Kumar, Emilien Dupont, Francisco JR Ruiz, Jordan S Ellenberg. Pengming Wang, Omar Fawzi, et al. Mathematical discoveries from program search with large language models. Nature, 625(7995):468-473) 2024.
Marija Šakota, Maxime Peyrard, and Robert West. Fly-swat or cannon? cost-effective language model choice via meta-modeling. In Proceedings of the 17th ACM International Conference on Web Search and Data Mining, pp. 606-615, 2024.
M Schlichtkrull, TN Kipf, P Bloem, R Van Den Berg, Titov, and M Welling, Modeling relational data with graph convolutional networks. arxiv. arXiv preprint arXiv:1703.06103. 2017.
Karan Singhal, Shekoofeh Azizi, Tao Tu, S Sara Mahdavi, Jason Wei, Hyung Won Chung, Nathan Scales, Ajay Tanwani, Heather Cole-Lewis. Stephen Pfohl, et al. Large language models encode clinical knowledge. arXiv preprint arXiv:2212.13138. 2022.
Charlie Snell, Jaehoon Lee, Kelvin Xu, and Aviral Kumar. Scaling Ilm test-time compute optimally be more effective than scaling model parameters. arXiv preprint arXiv:2408.00 03314, 2024.
Dimitris Stripelis, Zijian Hu, Jipeng Zhang, Zhaozhuo Xu, Alay Shah, Han Jin, Yuhang Yao. Salman Avestimehr, and Chaoyang He. Polyrouter: A multi-lim querying system. arXiv preprint ar0iv:2408. 12320, 2024.
12

---

Published as a conference paper at ICLR 2025
Rohan Taori, Ishaan Gulrajani, Tianyi Zhang, Yann Dubois, Xuechen Li, Carlos Guestrin, Percy Liang, and Tatsunori B Hashimoto. Stanford alpaca: An instruction followiinn llama model, 2023.
Hugo Touvron, Thibaut Lavril, Gautier Izacard, Xavier Martinet, Marie-Anne Lachaux, Timothée Lacroix, Baptiste Rozière, Naman Goyal, Eric Hambro, Faisal Azhar, et al. Llama: Open and efficient foundation language models. arXiv preprint arXiv:2302. 13971, 2023.
Petar Veličković, Guillem Cucurull, Arantxa Casanova, Adriana Romero, Pietro Lio, and Yoshua Bengio. Graph attention networks. arXiv preprint 17X.v:17001190000 2017.
Xiao Wang, Houye Ji, Chuan Shi, Bai Wang, Yanfang Ye, Peng Cui, and Philip S Yu. Heterogeneous graph attention network. In The world wide web conference, pp. 0122-0032,2 2019.
Yizhong Wang, Yeganeh Kordi, Swaroop Mishra, Alisa Liu, Noah A Smith, Daniel Khashabi, and Hannanch Hajishirzi. Self- inffrinstruct: Aligning language models with self-generate instructions. arXiv preprint arXiv:2212. 2022.
Jason Wei, Xuezhi Wang, Dale Schuurmans, Maarten Bosma, Fei Xia, Ed Chi, Quoc V Le, Denny Zhou, et al. Chain-of-thought prompting elicits reasoning in large language models. Advances in neural information processing systems, 35:24824 248377, 2022.
Qitian Wu, Wentao Zhao, Zenan Li, David Wipf, and Junchi Yan. Nodeformer: A scalable graph structure learning transformer for node classification. Advances in Neural Information Processing Systems, 35:27387-27401, 2022a.
Shiwen Wu, Fei Sun, Wentao Zhang, Xu Xie, and Bin Cui. Graph neural networks in recommender systems: survey. ACM Computing Surveys, 55(5):1-377 2022b.
Yongji Wu, Defu Lian, Yiheng Xu, Le Wu, and Enhong Chen. Graph convolutional networks with markov random field reasoning for social spammer detection. In Proceedings of the AAAI conference on artificial intelligence, volume 34, pp. 1054-1061, 2020.
Tian Xie, Bin Wang, and C-C Jay Kuo. Graphhop: An enhanced label propagation method for node classification. IEEE Transactions on Neural Networks and Learning Systems, 34(11):9287-931, 2022.
liangwei Yang. Zhiwei Liu, Yingtong Dou, Jing Ma, and Philip S Yu. Consisrec: Enhancing gnn for social recommendation via consistent neighbor aggregation. In Proceedings of the 44th international ACM SIGIR conference on Research and development in information retrieval, pp. 2141-2145, 2021.
Zhilin Yang, Peng Qi. Saizheng Zhang, Yoshua Bengio, William W Cohen, Ruslan Salakhutdinov, and Christopher D Manning. Hotpotqa: A dataset for diverse, explainable multi-hop question answering. Xrviv preprint /we.ivvv899.... 09000, 2018
Shunyu Yao, Dian Yu, Jeffrey Zhao, Izhak Shafran, Tom Griffiths, Yuan Cao, and Karthik Narasimhan. Tree of thoughts: Deliberate problem solving with large language models. Advances in Neural Information Processing Systems, 36, 2024.
Chengxuan Ying, Tianle Cai, Shengjie Luo, Shuxin Zheng, Guolin Ke, Di He, Yanming Shen, and Tie- Yan Liu. Do transformers really perform badly for graph representation? Advances in Neural Information Processing Systems, 34:28877-28888 2021.
Yaolun Zhang, Yinxu Pan, Yudong Wang, Jie Cai, Zhi Zheng, Guoyang Zeng, and Zhiyuan Liu. Py- bench: Evaluating 1lm agent on various real-world coding tasks. arXiv preprint urXiv:2407.1677. 2024.
Banghua Zhu, Ying Sheng. Lianmin Zheng, Clark Barrett, Michael Jordan, and Jiantao Jiao. On op- timal caching and model multiplexing for large model inference. arXiv preprint arXiv:2306.02003, 2023.
Xiaojin Zhu and Zoubin Ghahramani. Learning from labeled and unlabeled data with label propaga- tion. ProQuest number: information to all users, 2002.
13

---

Published as a conference paper at ICLR 2025
A DESCRIPTION FOR TASKS AND LLMS
Using descriptions generated by LLM and embeddings derived from BERT as initial node embeddings for GNNs enhances their expressiveness and generalization capabilities. Here, have listed descriptions of different tasks and various LLMs obtained using GPT-4o. Specifically, GPT-4o provides insights into the unique characteristics and challenges of different tasks, well as the size, cost, and particular strengths of different LLMs. The detailed descriptions are shown in the table below.
Table 6: Description of Alpaca task.
The Alpaca dataset is designed for instruction- followinn- tasks, where the model is required to generate coherent and contextually appropriate responses given instructions prompts. It focuses on understanding diverse user requests and providing informative and accurate outputs based on those instructions.
Table Description of GSM8K task.
The GSM8K dataset is tailored for mathematical problem-sol tasks. It consists of natural language math problems that require the model to comprehend the problem statement, apply the correct mathematical operations, and provide the solution. The primary challenge lies both parsing complex language and performing accurate calculations.
Table 8: Description of SQUAD task.
The SQuAD dataset is focused question-amswering tasks, where the model is given passage of text and needs to extract or generate a precise answer to a question based on the content of the passage. The dataset emphasizes comprehension, retrieval of relevant information, and concise answer generation.
B EXPERIMENTS ON DIFFERENT DATASETS AND API SETTINGS
B.1 GENERALIZATION CAPABILITY ON LARGER DATASETS
To validate whether method is applicable to additional tasks and LLM models, expanded upon the dataset from Section 3.1 by adding two new datasets. The first, HumanEyal (Chen et al., 2021), is 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 basic software interview questions. The second dataset HotpotQA (Yang et al., 2018 2018), 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 under our experimental settings, as shown in Table 20. Additionally, we incorporated four more LLM models from the Together AI API, based on the LLM models in Section 3.1: Qwen-2 (72b), Code Llama (34b), Mixtral-8x22B, and Upstage. Likewise, the size and cost information of these LLMs are summarized in Table 21. Further, similar to section 3.2, we constructed dataset based the extended dataset and the interaction data from LLMs, which was then divided into training, validation, and test sets in ratio of 70% 10% 20%, based different queries. We compared the performance of G GraphRoutez with other baselines on this extended dataset and reported the experimental results in Table 22. We observed that GraphRouter improved the Reward by at least 12.3% compared to
14

---

Published as a conference paper at ICLR 2025
Table 9: Description of Multi-News task.
The Multi-News dataset is aimed at text summarization tasks. contains multiple news articles the same topic, and the model's objective is to generate concise and comprehensive summary that integrates information from all the articles. The challenge is to distill key points while maintaining coherence and avoiding redundancy.
Table 10: Description of LLaMA-3 (7b).
This is relatively small-sized model (7 billion parameters) designed for general purpose language tasks. Its low cost per million tokens (0.2) makes it an affordable option for many applications requiring quick responses with moderate accuracy.
the baselines, confirming that GraphRouter's approach can be generalized to more datasets and other LLMs.
B.2 DISCRIMINATIVE ABILITY FOR SIMILAR LLMs
To explore whether GraphRouter can effectively model and differentiate between similar LLMs, we extracted data related to the LLaMA series of LLMs from the interaction dataset introduced in section 3.2 for training and prediction. Specifically, we extracted interaction data for LLMs including LLaMA-3 (7b), LLaMA-2 (7b), LLaMA-3 (70b). LLaMA-3- --rbbo (8b), LLaMA-3- -urbo (70b), and Llama-3.1-Turin (70b). We compared the performance of GraphRouter and the best-performing baseline, FrugalGPT, on this dataset and reported the specific results in Table 23. We observed that, compared to FrugaIGPT, GraphRouter improved the Reward by at least 10.8%. These observations demonstrate that GraphRouter effectively capture differences in the capabilities of different LLMs through interactions, achieving good results.
15

---

Published as a conference paper at ICLR 2025
Table 11: Description of Mixtral-8x7B.
With a combined size of 56 billion parameters, this model aims to provide stronger language modeling capabilities. Its cost per million tokens is 0.6, reflecting its balance between performance and affordability for more complex tasks.
Table 12: Description of NousResearch (34b).
A mid-sized model with 34 billion parameters, suitable for handling moderately complex language tasks. Its cost higher 0.8 per million tokens, indicating greater computational demand, likely due to its enhanced capabilities over smaller models.
Table 13: Description of LLaMA-2 (7b).
A compact model at 7 billion parameters, it offers similar capabilities and pricing to LLaMA-3 (7b) cost of 0.2 per million tokens. It's an efficient choice for tasks requiring decent performance without high computational costs.
Table 14: Description of Mistral-7b.
With 7 billion parameters, Mistral-7b is optimized for lightweight tasks, balancing speed and efficiency. Its cost per million tokens is 0.2, making it cost-effective for standard use cases without the need for complex computations.
Table 15: Description LLaMA-3 (70b).
A larger variant of LLaMA-3, this model has 70 billion parameters, providing advanced capa- bilities for complex tasks. Its cost per million tokens is 0.9, indicating its higher computational demand and enhanced performance.
Table 16: Description of LLaMA-3-Turbo (8b).
A variant optimized for speed and efficiency with 8 billion parameters. Its cost million tokens is only 0.2, suggesting that it is designed handle tasks quickly while being highly cost-effective.
Table 17: Description of LLaMA-3-Turbee (70b).
This model, at 70 billion parameters, is tailored for high performance with emphasis efficiency. The cost is 0.9 per million tokens, reflecting its advanced capabilities for a broad range of tasks requiring more computation.
16

---

Published as a conference paper at ICLR 2025
Table 18: Description of Llama-3.1-Turbo (70b).
Large model with 70 billion parameters, likely to offer strong capabilities for various language tasks. cost also 0.9 per million tokens, suggesting similar performance and computational needs as other 70b models.
Table 19: Description of Qwen-1.5 (72b).
With 72 billion parameters, Qwen-1.5 is among the largest models in the list, designed for high-complexity tasks. Its cost per million tokens is 0.9, making it comparable to other high-performance models in terms of both capability and expense.
Table 20: Overview of extended datasets.
Dataset
Task Type
Metric
Cases
Alpaca
Hybrid QA
F1
600
GSM8K
Reasoning
Accuracy
600
SQUAD
Reading Comprehension
F1
600
Multi-News
Summary
F1
600
HumanEval
Code
Pass
600
HotpotQA
Multi-hop QA
EM
600
Table 21: Statistics of larger LLMs set and their costs on Together API.
LLM
Size
Cost per 1M tokens
LLaMA-3 (7b)
7b
0.2
Mixtral-8x7B
56b
0.6
NousResearch
34b
0.8
LLaMA-2 (7b)
7b
0.2
Mistral-7b
7b
0.2
LLaMA-3 (70b)
70b
0.9
LLaMA-3-Turee (8b)
8b
0.2
LLaMA-3 3-rM-o (70b)
70b
0.9
Llama-3.1 1-1-uorbo (70b)
70b
0.9
Qwen- 1.5 (72b)
72b
0.9
Qwen-2 (72b)
72b
0.9
Code Llama (34b)
34b
0.8
Mixtral-8x22B
176b
1.2
Upstage
11b
0.3
17

---

Published as a conference paper at ICLR 2025
Table 22: Comparison of various methods large multi-task interaction dataset across Three Distinet Performance-Cost Weight Scenarios. Bold denotes the best results. Each metric reflects average values from multiple evaluation rounds.
Model
Performance First
Balance
Cost First
Performance
Cost
Reward Performance
Cost
Reward Performance
Cost
Reward
Largest L.LM
0.321
0.611
0.321
0,321
0.611
-0.145
0.321
0.611
-0.425
Smallest LLM
0.180
0.018
0.180
0.180
0.018
0.081
0.180
0.018
0.022
Prompt LLM
0.260
0.654
0.260
0.180
0.026
0.077
0.184
0.024
0.018
Hybrid LLM
0.350
0.611
0.350
0.311
0.256
0.028
0.184
0.103
-0.046
PrugalGPT
0.277
0.357
0.277
0,259
0.143
0.058
0,272
0.034
0.027
C2MAB-V
0.254
0.366
0.254
0.295
0.122
0.087
0.261
0.036
0.023
GraphRouter
0.393
0.220
0.393
0.297
0,052
0.122
0.299
0.018
0.046
Oracle
0.432
0.429
0.432
0.432
0.398
0.171
0.330
0.018
0.052
Table 23: Comparison of various methods on LLaMA-series dataset. Bold denotes the best results. Each metric reflects average values from multiple evaluation rounds.
Model
Performance First
Balance
Cost First
Performance Cost
Reward Performance
Cost
Reward Performance
Cost
Reward
FrugalGPT
0.382
0.299
0.382
0.367
0.043
0.162
0.372
0.030
0.050
GraphRouter
0.422
0.307
0.422
0.416
0.032
0.192
0.416
0.032
0.058
Oracle
0.489
0.376
0.489
0.459
0.051
0.204
0.436
0.032
0.062
18


Reviews:
[{"review_id": "EC8IWhbFCa", "reviewer": "Reviewer_3cRz", "summary": "This paper studies the problem of LLM selection for specific tasks and proposes a graph-based routing framework to select suitable LLM for the input query. By modeling the contextual information among tasks, queries, and LLMs as a heterogeneous graph, query, and LLM nodes, with interactions represented as edges, which efficiently captures the contextual information between the query’s requirements and the LLM’s capabilities. From the experiments, the authors compared the proposed method with different baselines and show that it can outperform regarding the reward. However, there are some issues that need to be addressed for better quality.", "questions": "See weaknesses.", "limitations": "", "rating": 6, "confidence": 4, "soundness": 2, "presentation": 2, "contribution": 2, "strengths": "1. This paper studies the challenging LLM model selection problem, which has been well addressed.\n\n2. This paper considers leveraging graph learning to incorporate more contextual information for LLM model selection.\n\n3. The experiments show that the proposed method achieves good performance compared to baselines.", "weaknesses": "1. 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.\n\n2. 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.\n\n3. Some details are unclear. For example, in table 5, what is the specific reward used for evaluation? And what is the purpose of adding task nodes? Can we just incorporate the information of task nodes into the query nodes?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "fYp4hXfqc9", "reviewer": "Reviewer_c8tQ", "summary": "This paper introduces GraphRouter, a novel graph-based approach for selecting appropriate LLMs for different queries. The authors construct a heterogeneous graph comprising task, query, and LLM nodes, with interactions represented as edges to capture contextual relationships. Through an innovative edge prediction mechanism, GraphRouter can adapt to both existing and newly introduced LLMs without requiring retraining. The work demonstrates significant performance improvements over baseline methods across multiple experimental settings, achieving at least 12.3% improvement in standard scenarios and 9.5% improvement in new LLM scenarios. The framework is evaluated on four distinct tasks (Alpaca, GSM8K, SQUAD, Multi-News) using ten different LLMs under various performance-cost tradeoff conditions.", "questions": "1．\tHow does the performance of GraphRouter change when handling LLMs with similar architectures but different sizes (e.g., different versions of LLaMA)? Is the framework able to effectively distinguish between such similar models?\n2．\tCould 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?\n3．\tHave you considered incorporating more sophisticated edge features, particularly for task-query relationships? How might this affect the framework's performance?", "limitations": "", "rating": 6, "confidence": 5, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "The paper seems to be the first to reformulate LLM selection as a graph-based edge prediction problem, providing a fresh perspective on router design. The heterogeneous graph structure effectively captures the complex relationships between tasks, queries, and LLMs. The framework addresses real-world challenges in LLM deployment, particularly the ability to handle new LLMs without retraining and balance performance with computational costs. The evaluation across three different cost-performance scenarios demonstrates practical utility. The authors conduct thorough experiments using multiple datasets, LLMs, and evaluation metrics. The ablation studies on GNN layer count and size provide useful insights for implementation.", "weaknesses": "1.\tAuthors 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 explaine why the heterogeneous graph structure (Figure 5) is optimal for capturing LLM-query relationships.\n2.\tIn L. 219, task-query edges are initialized uniformly to 1, which seems overly simplistic given the rich task-query relationships that could be captured. LLM-query edge features only use performance and cost concatenation, ignoring other potentially valuable signals like response length or generation time. \n3.\tThe 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.\n4.\tL. 327 present results across different LLMs but don'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.\n5.\tWhile GNN layer count is analyzed, other hyperparameters and embedding dimension are not thoroughly explored, but I think this is a minor problem.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "v5skOAab9m", "reviewer": "Reviewer_pHun", "summary": "This paper studies the LLM-Selection task, i.e., given a user query, selecting the appropriate LLM for execution that achieves the balance between performance and cost. Motivated by the contextual relationships between tasks, queries and LLMs, this paper utilizes a graph structure to represent these information and reformulates the LLM-Selection problem as a link-prediction task within graph. The proposed GraphRouter framework leverages a GNN to perform this task and can easily adapt to an inductive setting, accommodating the rapid evolution of LLMs. Extensive experiments in both transductive and inductive settings demonstrate the effectiveness of GraphRouter.", "questions": "1. Can GraphRouter be extended to a training-free version? For example, could a training-free GNN, e.g., SGC or LightGCN, be adapted to perform the link prediction task? I am curious about the capability of a training-free GNN for the LLM-selection task, as well as the percentage of training samples required to achieve satisfactory performance. \n\n2. Can the empirical evaluation be conducted in a more complex scenario, such as by adding more LLMs or introducing a broader variety of tasks?\n\n3. (A Minor Point) LLM selection task is also dependent on specific querying times. For example, in June 2024, the system may recommend GPT-4o for solving a graph reasoning task, while GPT-o1 might become the best choice at a later date. Can such temporal information be incorporated into the methodology or evaluation? \n\nI understand that time is limited during the rebuttal phase, therefore, brief discussions would be greatly appreciated.", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "*S1* This paper is well-motivated and addresses an important problem. The rapid evolution of LLMs necessitates effective LLM selection, yet existing methods often over-simplify this task as failing to consider the contextual information between tasks, queries, and LLMs, and they struggle to accommodate newly emerged LLMs. The graph-based LLM selection method not only offers a new and effective approach to the LLM selection task but also contributes a novel application in LLM + GNN research. \n\n*S2* The proposed method, GraphRouter, is both reasonable and intuitive. Extensive experiments demonstrate its effectiveness across different settings.\n\n*S3* The paper is easy to follow, with illustrative figures that facilitate comprehension.", "weaknesses": "*W1* 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). \n\n*W2* Current methodology overlooks the inherent relationships among LLMs. For instance, some LLMs belong to the same series (e.g., LLaMA3-7B, LLaMA3-8B). Adding links to indicate such parent or sibling connections could provide a more comprehensive graph modeling. I appreciate that the authors discuss this point in the limitations section. Additionally, common prior knowledge could enhance the LLM descriptions: e.g., CodeLLaMA series are effective for code-related tasks, while BaiChuan performs better in Chinese contexts.\n\n*W3* (Minor Points) Figures 2-4 require further explanation. In Figure 2, what does the notation $t$ represent? For Figures 3 and 4, does \"effect\" refer to performance, and why do some LLMs have multiple dots? Additionally, there appears to be an unfinished sentence in Line 138.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "8wWu2byqBr", "reviewer": "Reviewer_z8Nx", "summary": "This paper tackles the problem of generalizable LLM selection, which requires making performance-cost trade-offs in LLM selection for a given query while being capable of making inductive selections for new LLMs and tasks unseen during training. To address this challenging scenario, inspired by previous studies in graph machine learning and recommender systems, it proposes GraphRouter, a graph-based approach for context aware inductive LLM selection. Technically, it constructs a heterogeneous graph comprising task, query, and LLM nodes, with interactions represented as edges. Evaluating candidate LLMs and making a selection with respect to a given query can then be modeled as an edge-level prediction problem and the authors train a heterogeneous graph neural network for this purpose. Empirical studies demonstrate the effectiveness of GraphRouter for transductive and inductive settings.", "questions": "See Weakness.", "limitations": "", "rating": 8, "confidence": 3, "soundness": 2, "presentation": 3, "contribution": 3, "strengths": "**S1.** Compared to the baselines, the proposed graph-based approach conceptually allows better effectiveness and generalizability in LLM selection due to the incorporation of context information.\n\n**S2.** Empirical studies demonstrate the effectiveness of the proposed approach for both transductive and inductive settings.\n\n**S3.** Overall, the paper is easy to follow and the story is appealing and convincing.", "weaknesses": "**W1.** 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.\n\n**W2.** 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.\n\n**W3.** The discussion of related works can be further improved for proper credit attribution. For example:\n\n- 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].\n\n- 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].\n\n- The authors employ an attention-based heterogeneous GNN and should also discuss the first few attention-based heterogeneous GNNs, such as [3].\n\n[1] Zhu & Ghahramani. Learning From Labeled and Unlabeled Data With Label Propagation. 2002.\n\n[2] Schlichtkrull et al. Modeling Relational Data with Graph Convolutional Networks. 2017.\n\n[3] Wang et al. Heterogeneous Graph Attention Network. 2019.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)