PaperlayerView exact revision on arXiv

Source-linked AI summary

GraphRouter: A Graph-based Router for LLM Selections

Tao Feng, Yanzhen Shen, Jiaxuan You

arXiv:2410.03834v2cs.AI

TL;DR

Selecting among increasingly diverse LLMs requires balancing performance and computational cost while handling differences across tasks and queries. GraphRouter addresses this with an inductive heterogeneous graph over tasks, queries, and LLMs, using interaction edges for selection. Across three trade-off settings, it outperformed baselines by at least 12.3% and improved performance by at least 9.5% when new LLMs appeared in testing.

  • Problem

    LLM selection must balance performance and computational cost across diverse queries and tasks, but existing methods underuse contextual interactions and struggle with new LLMs.

  • Method

    GraphRouter constructs an inductive heterogeneous graph of task, query, and LLM nodes, representing their interactions as edges for contextual LLM selection.

  • Results

    12.3% minimum improvement over baseline models was achieved across three performance-cost settings, while new-LLM testing improved performance by at least 9.5%.

  • Takeaways & Limitations

    GraphRouter generalizes to new LLMs and adapts to diverse tasks without retraining, while improving selection outcomes over existing routers.

  • Takeaways & Limitations

    Existing comparison methods are limited by their reliance on transductive learning, requiring few-shot retraining when new LLMs are introduced.

Abstract

from arXiv · show

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 GraphRouter, which fully utilizes the contextual information among tasks, queries, and LLMs to enhance 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. Our codes for GraphRouter is released at https://github.com/ulab-uiuc/GraphRouter.

1 INTRODUCTION

GraphRouter addresses LLM selection as a contextual, inductive graph problem, using task, query, and LLM interactions to balance response effectiveness and cost. It is designed to generalize across new LLMs and tasks while improving over existing routers.

  • LLM selection must balance model performance and computational cost because models vary in capabilities, size, and demands across queries and tasks.
  • Existing routers use limited LLM sets, underuse task-query-LLM interaction context, require retraining for new LLMs, and often train separately for each task.
  • GraphRouter builds an inductive heterogeneous graph with task, query, and LLM nodes, representing their interactions as edges containing response performance and cost information.
  • GraphRouter reframes LLM selection as graph-based edge prediction to leverage contextual information among tasks, queries, and LLMs.
  • 12.3% minimum improvement over baseline models was achieved across three performance-cost settings, while new-LLM tests improved performance by at least 9.5% and saved training time.

2 GRAPHROUTER: GRAPH-BASED ROUTER FOR LLM SELECTION

GraphRouter frames LLM selection as choosing a model that balances query-specific performance with API cost. It represents tasks, queries, and LLMs in a heterogeneous graph and predicts query–LLM suitability through graph-based edge modeling.

  • LLM selection aims to choose a model that provides optimal performance while minimizing API cost.
  • Different LLMs vary in their ability to answer diverse queries, motivating contextual modeling beyond model names or IDs.
  • GraphRouter represents tasks, queries, and LLMs as three node types, with their interactions encoded as graph edges or edge features.
  • Task and query nodes are initialized from language-model embeddings, while LLM descriptions include capability and cost information for embedding.
  • LLM-query edges initialize with concatenated performance and cost, whereas task-query edges receive a value of 1.
  • A heterogeneous GNN updates node representations, and edge prediction identifies the best LLM by assigning the highest edge logit to the query.
  • In new-LLM tests, GraphRouter is evaluated using auxiliary interaction data that is excluded from training.

3 EXPERIMENTAL SETUP

GraphRouter is evaluated on multi-task interactions across diverse datasets and ten LLMs, using baselines, reward-oriented metrics, and standard and new-LLM test settings.

  • GraphRouter training: GraphRouter training uses a parameterized heterogeneous GNN with task-query and LLM-query edge weights, initialized node embeddings, mini-edge batches, and masked training edges.The listed algorithm initializes task, query, and LLM embeddings, samples mini-batches, and masks corresponding training edges.
  • Datasets and LLMs: The experiments combine four task datasets and responses from ten LLMs to form a multi-task interaction dataset.For each query, responses are evaluated against task ground truth, while cost uses token counts and model prices.
  • Data splitting: The new-LLM setting trains on six observable LLMs and evaluates all ten, using 80 auxiliary queries from four unseen LLMs only as testing-time few-shot data.The auxiliary dataset is excluded from training and validation.
  • Baselines: Baselines include largest-LLM, smallest-LLM, prompt-based, Hybrid LLM, FrugalGPT, other representative routers, and an oracle upper bound.The oracle routes each query to its optimal LLM using oracle information.
  • Metrics: Performance, cost, and reward are evaluated across Performance First, Balance, and Cost First scenarios with normalized metrics and different performance-cost weights.Reward is defined as Reward = α · Performance − β · Cost, using weights (1, 0), (0.5, 0.5), and (0.2, 0.8).

4 EXPERIMENTAL RESULTS

GraphRouter consistently outperforms existing routers across three scenarios and generalizes effectively to newly introduced LLMs. Its Reward is highest at a GNN size of 32 and two GNN layers, balancing effectiveness against computational overhead.

  • Comparison with Existing Baselines: 12.28% minimum Reward improvement over the strongest baselines was achieved across three scenarios, while GraphRouter reached at least 88.89% of the Oracle solution.It also achieved a better Performance–Cost trade-off than the two rule-based LLMs.
  • Generalization to New LLMs: Almost 10% Reward improvement and over 99% Time Cost reduction were achieved against C2MAB-V in the few-shot new-LLM setting.Compared with the trained GraphRouter, the few-shot approach substantially reduced time cost with only a slight performance loss.
  • Ablation Studies: Reward initially improves and then declines as GNN size increases, peaking at size 32 within the explored range of 16 to 80.The study treats GNN size as a performance and computational-overhead trade-off.
  • Ablation Studies: Reward initially improves and then declines as the number of GNN layers increases, reaching its maximum at two layers.The exploration covered zero to five layers; deeper networks can increase computational cost and introduce over-smoothing or overfitting.

5 ADDITIONAL RELATED WORKS

Related work frames LLM routing as a model-selection problem and highlights graphs as a way to represent complex relationships. GraphRouter differs by jointly modeling query–model, query–query, and model–model relationships for better generalizability.

  • LLM Selection: Prior LLM routers address inference-cost optimization by switching between models or selecting among a limited set of LLMs.These approaches include model routers for small-versus-large LLM choices and methods that balance exploration and exploitation.
  • LLM Selection: GraphRouter jointly models query–model, query–query, and model–model relationships rather than learning only from query–model interactions.This produces representations for tasks, queries, and models intended to improve generalizability.
  • Graph for Modeling Relationships: Graphs and GNNs provide established approaches for extracting nodes and edges and modeling relationships with embeddings.GNNs have also been applied in areas including recommendation systems and social applications.

6 CONCLUSION AND DISCUSSION

The paper presents GraphRouter as an inductive edge-prediction framework that uses graph structure to capture interaction context for multi-LLM routing. Experiments show superiority over competitive baselines and strong generalization when new LLMs appear at test time.

  • Conclusion: GraphRouter reframes multi-LLM routing as edge prediction between query nodes and LLM nodes.The framework learns task, query, and graph representations from prior interaction data.
  • Conclusion: Experiments on a combined dataset from four open-domain QA datasets and three application scenarios showed superiority over competitive LLM-selection baselines.The framework was also reported to be on par with the ideal “God’s-eye view” solution.
  • Discussion: GraphRouter showed strong generalization compared with previous baselines when new LLMs were introduced during testing.The conclusion positions graph-based incorporation of interaction data as a direction for future LLM-routing research.

A DESCRIPTION FOR TASKS AND LLMS

The paper supplies task and LLM descriptions for the datasets used in GraphRouter. These descriptions characterize task challenges and LLM size, cost, and strengths for initializing graph representations.

  • Task and LLM Descriptions: LLM-generated descriptions characterize the unique characteristics and challenges of different tasks.The descriptions are used as initial node embeddings for the GNNs.
  • Task and LLM Descriptions: LLM-generated descriptions characterize each LLM by its size, cost, and particular strengths.BERT-derived embeddings are also used as initial node embeddings.

B.1 GENERALIZATION CAPABILITY ON LARGER DATASETS

GraphRouter was evaluated on expanded datasets and additional LLMs to test generalization beyond the original setting. The results confirm that its approach generalizes to more datasets and other LLMs.

  • GraphRouter was tested with two additional datasets, HumanEval and HotpotQA, extending evaluation to coding and multi-hop question answering.HumanEval contains 164 programming problems, while HotpotQA contains 113k natural, multi-hop questions.
  • The expanded setting added four LLMs from the Together AI API: Qwen-2 (72b), Code Llama (34b), Mixtral-8x22B, and Upstage.
  • GraphRouter’s approach generalized to more datasets and other LLMs.

B.2 DISCRIMINATIVE ABILITY FOR SIMILAR LLMS

GraphRouter was evaluated on a dataset containing multiple closely related LLaMA models to test whether it could distinguish similar LLMs. It improved Reward over FrugalGPT by at least 10.8%, indicating effective interaction-based differentiation.

  • GraphRouter was tested on six closely related LLaMA models, spanning 7b, 8b, and 70b variants.The evaluated models included LLaMA-3, LLaMA-2, LLaMA-3-Turbo, and Llama-3.1-Turbo variants.
  • 10.8%: GraphRouter improved Reward by at least 10.8% compared with FrugalGPT on the LLaMA-series dataset.
  • GraphRouter effectively captured differences in the capabilities of similar LLMs through interactions.
Loading 2410.03834v2…