Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Large Language Model Routing with Benchmark Datasets

Tal Shnitzer, Anthony Ou, Mírian Silva, Kate Soule, Yuekai Sun, Justin Solomon, Neil Thompson, Mikhail Yurochkin

TL;DR

The paper tackles the problem of selecting the best LLM for a new task by learning routers from benchmark data. It formalizes per-model correctness predictors trained on embedded inputs from diverse benchmarks and introduces three routing scores, including an OOD-aware score that accounts for imperfect predictors. Empirical results on HELM and Mix-Instruct show that benchmark-informed routing can outperform the Best Model on Average and enable smaller models to compete with larger ones, with efficiency gains at test time. The work highlights the practical value of leveraging benchmark byproducts for cost-efficient, scalable LLM deployment and outlines paths for improving OOD generalization and dataset coverage.

Abstract

There is a rapidly growing number of open-source Large Language Models (LLMs) and benchmark datasets to compare them. While some models dominate these benchmarks, no single model typically achieves the best accuracy in all tasks and use cases. In this work, we address the challenge of selecting the best LLM out of a collection of models for new tasks. We propose a new formulation for the problem, in which benchmark datasets are repurposed to learn a "router" model for this LLM selection, and we show that this problem can be reduced to a collection of binary classification tasks. We demonstrate the utility and limitations of learning model routers from various benchmark datasets, where we consistently improve performance upon using any single model for all tasks.

Large Language Model Routing with Benchmark Datasets

TL;DR

The paper tackles the problem of selecting the best LLM for a new task by learning routers from benchmark data. It formalizes per-model correctness predictors trained on embedded inputs from diverse benchmarks and introduces three routing scores, including an OOD-aware score that accounts for imperfect predictors. Empirical results on HELM and Mix-Instruct show that benchmark-informed routing can outperform the Best Model on Average and enable smaller models to compete with larger ones, with efficiency gains at test time. The work highlights the practical value of leveraging benchmark byproducts for cost-efficient, scalable LLM deployment and outlines paths for improving OOD generalization and dataset coverage.

Abstract

There is a rapidly growing number of open-source Large Language Models (LLMs) and benchmark datasets to compare them. While some models dominate these benchmarks, no single model typically achieves the best accuracy in all tasks and use cases. In this work, we address the challenge of selecting the best LLM out of a collection of models for new tasks. We propose a new formulation for the problem, in which benchmark datasets are repurposed to learn a "router" model for this LLM selection, and we show that this problem can be reduced to a collection of binary classification tasks. We demonstrate the utility and limitations of learning model routers from various benchmark datasets, where we consistently improve performance upon using any single model for all tasks.
Paper Structure (36 sections, 2 theorems, 18 equations, 8 figures, 4 tables)

This paper contains 36 sections, 2 theorems, 18 equations, 8 figures, 4 tables.

Table of Contents

  1. Introduction
  2. Related work
  3. Benchmarking
  4. Model selection
  5. Out-of-distribution model selection
  6. Routing LLMs
  7. Learning from Benchmarks
  8. Supervised learning from benchmarks
  9. LLM routing with (imperfect) correctness predictors
  10. A simple OOD confidence model
  11. Connection to meta-learning
  12. Experiments
  13. Model routing on HELM
  14. Data
  15. Models
  16. ...and 21 more sections

Key Result

Lemma 4.1

Let $\ell(y_1,y_2) = \rho(y_1 - y_2)$ for some subadditive $\rho:\mathbf{R}\to\mathbf{R}$ (e.g. $\rho(x) = \frac{1}{2}x^2$ for the square loss). We have

Figures (8)

  • Figure 1: We learn the strengths of candidate LLMs (marked with corresponding colors) on various tasks (emojis: QA, reasoning, summarization, etc.) and domains (4 sections within each box: finance, legal, general knowledge, etc.) from benchmark datasets. We accomplish this by training a binary classifier per LLM (upper part of the figure). For a new task, we score each LLM with these binary classifiers and recommend an LLM for the user (lower part).
  • Figure 2: Using $\min(\alpha n^{d'}, 50)$ training samples from $d'$ to reduce OOD gap.
  • Figure 3: Average metrics on subsets of the MixInstruct test set, defined by limiting the maximal average distance between test instances and their closest neighbors in the reference (train) set.
  • Figure 4: LLM routing with $\leq13$B parameter models compared to Llama 2 70B.
  • Figure 5: Correlation($S_3$, Accs.) and $u(d')$.
  • ...and 3 more figures

Theorems & Definitions (3)

  • Lemma 4.1
  • Lemma D.1: Lemma \ref{['lem:adaptive-improvement']}
  • proof