GraphRouter: A Graph-based Router for LLM Selections

## Abstract

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 dependence on a transductive learning framework. To address these shortcomings, we introduce a novel inductive graph framework, named as \method , which fully utilizes the contextual information among tasks, queries, and LLMs to enhance the LLM selection process. \method 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, \method 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 \method 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. Our codes for \method will soon be released at https://github.com/ulab-uiuc/GraphRouter .

## 1 Introduction

The field of Large Language Models (LLMs) is advancing quickly, offering an increasingly diverse range of models that vary in size, functionality, and computational demands (Bang, ; Liu et al., ). Although larger models tend to deliver better performance, their high computational costs make them inefficient for many less complex tasks (Snell et al., ; Chen & Varoquaux, ). Additionally, LLMs demonstrate varied performance across different types of queries and tasks (Ahmed et al., ; Zhang et al., ), especially with the development of domain-specific LLMs (Singhal et al., ; Luo et al., ). These challenges make it difficult to recommend the optimal LLM services to 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 LLMs, 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 et al., ) trains a binary score router function to determine whether to select a small LLM or a large LLM for a specific query. Although it balances cost and performance, it is limited to just two LLMs, which falls short of addressing the real-world demand for a wide range of LLMs. Some other studies (Dai et al., ; Chen et al., ) have introduced more advanced router models to address the challenge of selecting among a limited set of LLMs (typically 3 to 5). More specifically, FrugalGPT (Chen et al., ) proposes a router model based on BERT (Devlin, ) to determine whether to switch to a larger LLM or not, and C2MAB-V (Dai et al., ) constructs a bandit-based router to balance between the exploration and exploitation when choosing LLM for the user. However, as shown in Table 1, 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., ; Joachims, ), 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-shot interaction data before they can be used – a process that is impractical for recommending LLMs to a 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 \method, a graph-based router for LLM selection. \methodutilizes an 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, \methodconstructs 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 LLM node. Then, we are able to transform the task of predicting the cost and performance of an LLM-query 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, so an effective framework should also have the ability to accommodate these evolving models. In order to make \methodgeneralizable to new LLMs, we 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 on this, we derive an initial embedding for each LLM using a moderate-size pre-trained language model (e,g, BERT (Devlin, )). This approach offers an advantage over directly using one-hot encoding, as it enables us to generate inductive and more informative initial embeddings for new LLMs.
For the \methodmodel, we further develop a heterogeneous GNN that aggregates information from neighboring nodes of different types; given few-shot data, we verified that a trained \methodcan generalize to new LLM nodes without retraining.

In summary, our main contributions are as follows:

- •
To the best of our knowledge, we are the first work to build router for LLM 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 a variety of tasks without retraining.
- •
In three experiment settings with different performance and cost tradeoffs, \methodoutperformed 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.

**Table 1: Comparison of \methodwith existing methods from three perspectives: contextual info, generalization to new LLMs, and multi-task support. Compared to other approaches, \methodintroduces 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., ) | Indices | ✗ | ✗ |
| FrugalGPT (Chen et al., ) | LLM name | ✗ | ✗ |
| C2MAB-V (Dai et al., ) | One-hot embedding | ✗ | ✗ |
| \method | Graph-based contexts | ✓ | ✓ |

## 2 GraphRouter: Graph-based Router for LLM Selection

Figure: Figure 1: Overview of \method’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_{n}$ for generation. Then, $LLM_{n}$ is used to generate the response. In the end, this response, along with the measured effectiveness and cost, is returned to 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 to train the router.
Refer to caption: x1.png

### 2.1 Preliminaries

We introduce the LLM selection problem in this section. As shown in the left part of Figure 1, 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 to 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 a suitable $LLM_{n}$ to answer the user query. Finally, the response is returned to the user with its performance and cost. Such an interaction process generates rich contextual data, which contains information on tasks, queries, selected LLM, response, performance, and cost. We organize the data in a table, as shown on the right.

### 2.2 Motivating Examples

Traditional LLM selection methods (Ding et al., ; Chen et al., ; Dai et al., ) 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 4 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 are embedded in the contextual information between the queries and the LLMs. In addition, from Figure 4, 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 a 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 a single LLM responds to queries across different tasks. Through Figures 4 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., ). 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 LLMs, 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 or IDs, an effective use of the contextual information that encompasses the interaction between task, query, and LLM is needed.

Figure: Figure 2: The probability distribution of a small LLM (LLaMA-3 (7b)) having a better performance value than a large LLM (LLaMA-3 (70b)) by $t$ on the Alpaca dataset, where $t\in[-1,1]$.
Refer to caption: x2.png

### 2.3 GraphRouter Framework

Figure: Figure 5: Overview of \methodmethodology. \methodfirst converts the interaction data among tasks, queries, and LLMs into a graph. Specifically, as 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 are modeled 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.
Refer to caption: x5.png

Method Overview. As shown in Figure 5, \methodfirst transforms the interaction data among tasks, queries, and LLMs into a graph. Specifically, as shown in the right side of Figure 5, we model tasks, queries, and LLMs in the left table as 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 in Figure 5, we have three types of nodes (task node $h^{(l)}_{t}$, query node $h^{(l)}_{q}$, and LLM node $h^{(l)}_{m}$) and two types of edges (task-query edge $w_{tq}$ and LLM-query edge $w_{mq}$).

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 as GPT-4o, to generate descriptions of the tasks, and then encode the description to obtain its embedding $e_{t}$. More specifically, we take the average output token embedding after feeding the description into the moderate-size pre-trained language model (PLM), such as BERT (Devlin, ). The initialization of query nodes is also obtained by embedding the query $e_{q}$ through the same PLM. As for the initialization of LLM nodes, the traditional approach often initializes directly using the name or ID of the LLM, which not only limits its ability to generalize to new LLMs but also omits important background information. Here, we still adopt a prompt-based approach. We design prompts for an LLM to describe the capabilities of each LLM. In addition, we also add information about the cost of each LLM after the description. Then, similar to how we obtain the task’s embedding, we use the same PLM to compute the initial embeddings $e_{l}$ of the different LLMs. All the descriptions we generate for different tasks and LLMs can be found in Appendix A.

As for the task-query edges, we assign a value of 1 for their initialization. For the initialization of LLM-query edges, we jointly consider the performance and cost information in the interaction data, and assign the concatenation of performance and cost as their initial features.

Predict via a heterogeneous GNN. We implement the predictive model $f_{\phi}$ over task nodes, query nodes, and LLM nodes using a heterogeneous GNN, as 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 to learn expressive node embeddings $\mathbf{h}$ through an iterative weighted aggregation of the local network neighborhoods. The $l$-th iteration of the $\text{GraphConv}(\cdot)$, or the node embeddings update of the $l$-th layer, is represented as:

$$ $\mathbf{h}_{t}^{(l)}=\mathbf{U}^{(l)}_{t}\textsc{Concat}\Big{(}\textsc{Mean} \Big{(}\{\textsc{ReLU}(\mathbf{W}^{(l)}_{t}\mathbf{h}_{q}^{(l-1)}),q\in \mathcal{N}(t)\}\Big{)},\mathbf{h}_{t}^{(l-1)}\Big{)},$ (1) $$

$$ $\mathbf{h}_{q}^{(l)}=\mathbf{U}^{(l)}_{q}\textsc{Concat}\Big{(}\textsc{Mean} \Big{(}\{\textsc{ReLU}(\mathbf{w}_{\mathbf{1}[t\in V_{t},m\in V_{t}]}\mathbf{W }_{\mathbf{1}[t\in V_{d},m\in V_{t}]}^{(l)}\mathbf{h}_{u}^{(l-1)},u\in\mathcal {N}(v)\}\Big{)},\mathbf{h}_{q}^{(l-1)}\Big{)},$ (2) $$

$$ $\mathbf{h}_{m}^{(l)}=\mathbf{U}^{(l)}_{m}\textsc{Concat}\Big{(}\textsc{Mean} \Big{(}\{\textsc{ReLU}(\mathbf{w}_{mq}^{T}\mathbf{W}^{(l)}_{m}\mathbf{h}_{q}^{ (l-1)}),q\in\mathcal{N}(m)\}\Big{)},\mathbf{h}_{m}^{(l-1)}\Big{)},$ (3) $$

where $\mathbf{h}^{(l)}$ is the node embedding after $l$ iterations, $\mathbf{h}_{t}^{(l)},\mathbf{h}_{q}^{(l)},\mathbf{h}_{m}^{(l)}$ have been initialized as $\mathbf{h}_{t}^{(0)},\mathbf{h}_{q}^{(0)},\mathbf{h}_{m}^{(0)}=e_{t},e_{q},e_{m}$ respectively
as explained above, and $\mathcal{N}(v)$ denotes the direct neighbors of $v$. $\mathbf{1}[v\in V_{d},u\in V_{t}]$ indicates the message type (whether from task to query, or LLM to query), and $\mathbf{U}^{(l)},\mathbf{W}^{(l)}$ are learnable parameters. In addition, $\mathbf{w}_{\mathbf{1}[t\in V_{t},m\in V_{t}]}$ indicates that different edge types correspond to different edge features; specifically, if it is from task to query, it is represented as $w_{tq}$, and from LLM to query, it is represented as $w_{mq}$.

Following the update of the task, query, and LLM node embeddings, we obtain the query-task combined embedding as $\mathbf{h}_{qt}^{(l)}=\textsc{MLP}(\textsc{Concat}(\mathbf{h}_{t}^{(l)},
\mathbf{h}_{q}^{(l)}))$. We model the LLM selection problem as an 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 Sec 3.4) achieved by different LLMs and then set the edge labels of the query to other LLMs to 0 and the edge label of the query to the best LLM to 1. As such, LLM prediction can be made through $\text{EdgePred}(\cdot)$ in the form of $\hat{y}_{logits}=\textsc{MEAN}\Big{(}\textsc{Dot}(\mathbf{h}_{qt}^{(l)},
\mathbf{h}_{m}^{(l)})\Big{)}$. We have summarized the detailed training process of \methodin Algorithm 1, whose details are shown as follows. In addition, in the testing of \method, we identify the LLM node that has maximum edge logits with the query node as the best LLM, which can be computed as

$$ $\hat{y}=\arg\max_{m}\Big{(}\text{EdgePred}(h_{qt},h_{m})\Big{)}.$ (4) $$

Figure: Algorithm 1 Training of \method

for new LLMs setting. Traditional routers cannot generalize to new LLMs directly under few-shots settings, as they require retraining through interactions with the query for each new LLM. This is inadequate for keeping up with the rapidly evolving changes in LLMs in the real world. To test our framework and baselines under this real-world setting, following (Cao et al., ; Fey et al., ), 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 a 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.

- •
Alpaca (Taori et al., ) 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., ) framework, which iteratively prompts the language model to generate new training instances given a few manually written instructions.
- •
GSM8K (Cobbe et al., ) evaluates the model’s ability for multi-step mathematical reasoning with 8.5k linguistically diverse grade school math word problems.
- •
SQUAD (Rajpurkar, ) is a crowdsourced reading comprehension dataset based on Wiki articles. It contains over 100k QA pairs connected to over 500 articles.
- •
Multi-News (Fabbri et al., ) 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 our problem, with their statistics shown in Table 3. All of the LLMs and their token costs here are accessed through the Together API (^1^11https://docs.together.ai/docs/inference-models).

### 3.2 Data Preprocessing and Splitting

**Table 2: Overview of 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 |

Given the above dataset and LLMs, we construct a 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 it 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 is 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., ) to calculate the number of tokens.

After obtaining the multi-task interaction dataset, we split the dataset according to different experimental settings. We mainly have two major settings. The first is the standard setting, where all LLMs are visible in both the training and test sets, with some new queries appearing in the test set. The data 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, we 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 set as it is in the standard setting. Furthermore, following (Cao et al., ; Fey et al., ), 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 a few-shots during the testing phase.

### 3.3 Baseline Methods

We compare our GraphRouter model with the following baselines. We first compare \methodwith two rule-based baselines:

- •
Largest LLM always selects the largest LLM available.
- •
Smallest LLM always selects the smallest LLM available.

Then we compare \methodwith a prompt-based baseline:

- •
Prompt LLM incorporates the query, candidate models, and objectives (e.g., prioritizing effectiveness) directly into the prompt, and feeds it into an external LLM (e.g., GPT-4) to select the most suitable LLM from a pool of candidates.

Further, \methodis compared with three representative routers for LLM selection:

- •
Hybrid LLM (Ding et al., ), when given a small LLM and a large LLM, trains a pre-trained language model to assign queries to the small or 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 et al., ) with RoBERTa (Liu, ) as the pre-trained language model and observe better performance.
- •
FrugalGPT (Chen et al., ) utilizes a pre-trained language model to predict the score of the generation result of all LLMs given a query, and then selects the LLM with the highest score within a given cost. We also use RoBERTa (Liu, ) as the router’s backbone model.
- •
C2MAB-V (Dai et al., ) uses a 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 a gold baseline as the optimal solution for LLM selection. The purpose of setting up the baseline is to see how far \methodis 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.

**Table 3: Statistics of different LLMs 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-Turbo (8b) | 8b | 0.2 |
| LLaMA-3-Turbo (70b) | 70b | 0.9 |
| Llama-3.1-Turbo (70b) | 70b | 0.9 |
| Qwen-1.5 (72b) | 72b | 0.9 |

### 3.4 Metrics

We utilize three metrics to evaluate the performance of \methodand baselines.

- •
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.
- •
Cost evaluates the average LLM inference cost when responding to the queries, which is described in Sec 3.2.
- •
Reward is used to measure how well a method balances performance and cost. Different users may have varying levels of emphasis on performance and cost. Therefore, we define three scenarios: Performance First, Balance, and Cost First, to correspond to situations where users prioritize high performance, value both high performance and low cost equally, or prioritize low cost, respectively. Specifically, to eliminate the influence of scale, we first normalize both performance and cost. Then, we define the score as $Reward=\alpha\cdot Performance-\beta\cdot Cost$. For the three scenarios, we set the values of $\alpha$ and $\beta$ to (1, 0), (0.5, 0.5), and (0.2, 0.8), respectively.

### 3.5 Implementation Details

In the training stage, we set the graph neural network as a two-layer graph attention network, with a 32-dim hidden dimension. The batch size is 32, and the max training epoch is set to 1000. We use Adam optimizer (Diederik, ) for model training and gradually decay the learning rate from 1e-3 to 0 with LambdaLR scheduler. We implement our proposed method using PyTorch(^2^22 https://pytorch.org/) and PyG(^3^33 https://pytorch-geometric.readthedocs.io/en/latest/), 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(^4^44 https://www.together.ai/) to obtain responses.

## 4 Experimental Results

### 4.1 Comparison with existing baselines.

As shown in Table 4, we compare \methodwith seven baselines in three scenarios. We can observe that \methodconsistently and substantially surpasses existing routers, delivering a minimum effect improvement of 12.28% on metric Reward compared to the strongest baselines. Additionally, we observe that \methodachieves 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, \methodachieves a better trade-off between Performance and Cost, therefore achieving a higher effect on Reward. Analyzing the effect of Prompt LLM, Hybrid LLM (Ding et al., ), and FrugalGPT (Chen et al., ), we demonstrate that without sufficient contextualized information, even LLM and trained moderate-size LM struggle to understand the query and candidates LLM effectively, even if we ignore their high inference costs. These comparisons and results validate our claim that effective usage of contextual information is crucial for selecting the optimal LLM.

**Table 4: Comparison of Various Methods on 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 on 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 |

### 4.2 Generalization ability to new LLMs

To validate the generalization ability of \methodwhen facing new LLMs, we conduct experiments as 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 \method(few-shots) with HybridLLM, FrugalGPT, C2MAB-V, and \method(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., ), \method(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 \method(trained), our approach significantly reduces time cost with only a slight performance loss. These observations demonstrate that \methodis both effective and efficient in generalizing to new LLMs.

**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 et al., )).**
| Method | Reward | Reward Improvement(%) | Time Cost | Time Cost Reduction(%) |
| --- | --- | --- | --- | --- |
| HybridLLM | 0.01 | -94.71 | 273.45 | 49.57 |
| FrugalGPT | 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 |

Figure: Figure 6: The Reward of \methodvaries with the input size of the GNN. Here, we set our in-channels equal to out-channels. As observed in the figure, the Reward of \methodinitially increases and then decreases, achieving the highest Reward when the size is 32.
Refer to caption: x6.png

### 4.3 Ablation Studies

How does \methodperform with varying sizes of GNN? The size of a 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 \method, we explored sizes ranging from 16 to 80, as shown in Figure 7. As depicted in the figure, the Reward of \methodinitially improves as the size increases, reaching its peak at a size of 32, after which it starts to decline.

What is the impact of different numbers of GNN layers on \method’s effectiveness? The number of GNN layers has a significant impact on the expressiveness of 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 \method, we conduct an exploration with the number of layers ranging from 0 to 5, as shown in Figure 7. As depicted in the figure, the Reward of \methodinitially improves with more layers but then declines, achieving its highest Reward when the number of layers is 2.

## 5 Additional Related Works

LLM selection. With the scaling of the number of parameters in LLMs, their inference cost is also rapidly increasing. To optimize inference cost, several works have proposed using model switching to mitigate this issue.
When given a small and a large LLM, Zhu et al. () fine-tune a pre-trained language model as the model router to predict whether using a small LLM is sufficient. HybridLLM (Ding et al., ) 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 are present. Chen et al. () train the router to predict the reliability score given a query and an LLM index. Šakota et al. () generalize the router’s training objective to accommodate the scenario where additional cost or performance constraints exist. Stripelis et al. () examine the effectiveness of lighter routers built upon the k-nearest neighbors algorithm or Multilayer Perceptrons. C2MAB-V (Dai et al., ) employs a bandit-based model for LLM selection, treating each LLM as an arm and incorporating an exploration mechanism 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 jointly modeling the query-model, query-query, and model-model relationship. This allows us 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., ; Cao et al., ; Gao & Xu, ; Chen et al., ; Wu et al., ; Yang et al., ). 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., ), 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 et al., ) and social networks (Wu et al., ). GNNs can be broadly classified into Message Passing Neural Networks (MPNNs), which include models like GCN (Kipf & Welling, ), GraphSAGE (Hamilton et al., ), and GAT (Veličković et al., ), as well as non-MPNN architectures (Wu et al., ; Ying et al., ). Additionally, Heterogeneous Graph Neural Networks (HeterGNNs) (Peng et al., ; Hu et al., ) 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 et al., ; Cao et al., ; Gao & Xu, ; Chen et al., ), 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

#### Conclusion.

We present \method, a 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 \methodover 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 \method’s strong generalization ability compared to previous baselines. We hope that \method, along with our approach of incorporating interaction data through graphs, will facilitate future research on LLM routing.

#### Limitations

This work primarily serves as exploratory work to validate the idea of how modeling past interaction data in a graph could enhance the process of LLM selection. We acknowledge that leveraging more complex graph signals, such as paths, or the taxonomy of LLMs (e.g., family trees like LLaMA2 → LLaMA3 → LLaMA3.1 (Touvron et al., )), could further improve \method’s performance, but that is beyond the scope of this paper, and we leave it for future work.

#### Future Work

Some other interesting questions to explore in the future include: 1) When answering complex queries, prompting methods like Chain-of-Thought (Wei et al., ), Tree-of-Thought (Yao et al., ) 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 is interesting to explore whether \methodcould also be adapted to this task. 2) In a multi-agent system, it is critical to choose the appropriate LLM for each module on a specific query and task. It would be valuable to conduct experiments on whether an inductive graph framework like \methodcould also excel on this challenging task, where multiple LLMs are being selected simultaneously.

## Appendix 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 of. Here, we 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, as 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-following tasks, where the model is required to generate coherent and contextually appropriate responses to given instructions or prompts. It focuses on understanding diverse user requests and providing informative and accurate outputs based on those instructions. |
| --- |

**Table 7: Description of GSM8K task.**
| The GSM8K dataset is tailored for mathematical problem-solving 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 in both parsing complex language and performing accurate calculations. |
| --- |

**Table 8: Description of SQUAD task.**
| The SQuAD dataset is focused on question-answering tasks, where the model is given a 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. |
| --- |

**Table 9: Description of Multi-News task.**
| The Multi-News dataset is aimed at text summarization tasks. It contains multiple news articles on the same topic, and the model’s objective is to generate a 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 a 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. |
| --- |

**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 is higher at 0.8 per million tokens, indicating a greater computational demand, likely due to its enhanced capabilities over smaller models. |
| --- |

**Table 13: Description of LLaMA-2 (7b).**
| Compact model at 7 billion parameters, it offers similar capabilities and pricing to LLaMA-3 (7b) at a 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 of LLaMA-3 (70b).**
| A larger variant of LLaMA-3, this model has 70 billion parameters, providing advanced capabilities 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 per million tokens is only 0.2, suggesting that it is designed to handle tasks quickly while being highly cost-effective. |
| --- |

**Table 17: Description of LLaMA-3-Turbo (70b).**
| This model, at 70 billion parameters, is tailored for high performance with an emphasis on efficiency. The cost is 0.9 per million tokens, reflecting its advanced capabilities for a broad range of tasks requiring more computation. |
| --- |

**Table 18: Description of Llama-3.1-Turbo (70b).**
| Large model with 70 billion parameters, likely to offer strong capabilities for various language tasks. Its cost is 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. |
| --- |

## Peer Reviews

confidence: 4
contribution: 2
initial_rating: 3
presentation: 2
questions: See weaknesses.
soundness: 2
strengths: 1. This paper studies the challenging LLM model selection problem, which has been well addressed.

2. This paper considers leveraging graph learning to incorporate more contextual information for LLM model selection.

3. The experiments show that the proposed method achieves good performance compared to baselines.
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.
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.

2. 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.

3. 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?

---

confidence: 5
contribution: 3
initial_rating: 6
presentation: 3
questions: 1．	How 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?
2．	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?
3．	Have you considered incorporating more sophisticated edge features, particularly for task-query relationships? How might this affect the framework's performance?
soundness: 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.
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.
weaknesses: 1.	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 explaine why the heterogeneous graph structure (Figure 5) is optimal for capturing LLM-query relationships.
2.	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. LLM-query edge features only use performance and cost concatenation, ignoring other potentially valuable signals like response length or generation time. 
3.	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.
4.	L. 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.
5.	While GNN layer count is analyzed, other hyperparameters and embedding dimension are not thoroughly explored, but I think this is a minor problem.

---

confidence: 4
contribution: 3
initial_rating: 6
presentation: 3
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. 

2. Can the empirical evaluation be conducted in a more complex scenario, such as by adding more LLMs or introducing a broader variety of tasks?

3. (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? 

I understand that time is limited during the rebuttal phase, therefore, brief discussions would be greatly appreciated.
soundness: 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. 

*S2* The proposed method, GraphRouter, is both reasonable and intuitive. Extensive experiments demonstrate its effectiveness across different settings.

*S3* The paper is easy to follow, with illustrative figures that facilitate comprehension.
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.
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). 

*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.

*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.

---

confidence: 3
contribution: 3
initial_rating: 6
presentation: 3
questions: See Weakness.
soundness: 2
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.

**S2.** Empirical studies demonstrate the effectiveness of the proposed approach for both transductive and inductive settings.

**S3.** Overall, the paper is easy to follow and the story is appealing and convincing.
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.
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.

**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.

**W3.** 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.

Meta-review: This paper proposes GraphRouter, a graph-based routing system which is used to select appropriate LLMs to handle input queries.  The authors propose a graph ML angle on modeling the context information between queries, tasks, and LLMs as a heterogeneous graph, with interactions as edges.  The authors introduce a reward metric which offsets cost from performance  and demonstrate their proposal outperforms others in terms of net reward.

Overall, reviewers consistently evaluated this work positively.  There were a few themes of comments which I encourage the authors to fold into the revision:

- Improved clarifications about the introduction of various graph metadata (e.g task nodes as 3cRz raised, graph definition choices which c8tQ raised, node features as z8Nx raised)

- Discussion and rationale around the practicality of the evaluation setting and proposed approach (pHun, 3cRz, z8Nx) -- some reviewers still have doubts on this point and some further text in the manuscript to justify the choices made would be helpful.

Despite these limitations, the work proposes an interesting interplay conceptually between graph ML models and their use in the space of LLM selection.

Decision: Accept (Poster)