PaperlayerView exact revision on arXiv

Source-linked AI summary

RouteLLM: Learning to Route LLMs with Preference Data

Isaac Ong, Amjad Almahairi, Vincent Wu, Wei-Lin Chiang, Tianhao Wu, Joseph E. Gonzalez, M Waleed Kadous, Ion Stoica

arXiv:2406.18665v4cs.LGcs.AIcs.CL

TL;DR

LLM deployment must balance response quality against the higher cost of stronger models, while routers should generalize across queries and model pairs. The paper trains preference-data routers with data augmentation to select between strong and weak LLMs. Across benchmarks, the routers reduce costs by over 2x with minimal quality impact and maintain performance across multiple model pairs without retraining.

  • Problem

    LLM deployment faces a quality-cost trade-off, and effective routers must generalize across domains and model pairs without retraining.

  • Method

    The paper trains binary routers from human preference data, using data augmentation to select between strong and weak LLMs.

  • Results

    Over 2x cost savings are achieved with minimal impact on response quality, while routers maintain strong performance across multiple strong/weak model pairs without retraining.

  • Takeaways & Limitations

    Dataset augmentation improves routing on benchmarks, especially when augmented data resembles the evaluation data, and the routers generalize to unseen data.

  • Takeaways & Limitations

    Real-world application distributions may differ substantially from the evaluated benchmarks, and this work focuses on routing between two models.

Abstract

from arXiv · show

Large language models (LLMs) exhibit impressive capabilities across a wide range of tasks, yet the choice of which model to use often involves a trade-off between performance and cost. More powerful models, though effective, come with higher expenses, while less capable models are more cost-effective. To address this dilemma, we propose several efficient router models that dynamically select between a stronger and a weaker LLM during inference, aiming to optimize the balance between cost and response quality. We develop a training framework for these routers leveraging human preference data and data augmentation techniques to enhance performance. Our evaluation on widely-recognized benchmarks shows that our approach significantly reduces costs-by over 2 times in certain cases-without compromising the quality of responses. Interestingly, our router models also demonstrate significant transfer learning capabilities, maintaining their performance even when the strong and weak models are changed at test time. This highlights the potential of these routers to provide a cost-effective yet high-performance solution for deploying LLMs.

1 INTRODUCTION

LLM deployment requires balancing model quality against serving cost, motivating routers that select between stronger and weaker models. The proposed preference-data framework reduces cost while preserving quality and generalizes across models and unseen data.

  • Larger LLMs generally provide greater capability at higher cost, whereas smaller models are cheaper but may produce lower-quality responses.
  • LLM routing sends simpler queries to smaller models and more complex queries to larger models to balance response quality with cost efficiency.
  • An ideal router should optimize quality with one model call, generalize to out-of-domain queries, and operate across LLMs without retraining.
  • The framework learns routers from preference data to route between strong and weak models, targeting a specified fraction of strong-model performance while reducing cost.The strong-model examples include GPT-4, while Mixtral-8x7B represents the weak-model class.
  • Over 2x cost savings are achieved on popular benchmarks with minimal impact on response quality through human preference data and data augmentation.
  • The routers generalize to unseen data and remain effective across multiple LLMs, while the framework and evaluation tools are open sourced.

2 RELATED WORK

Prior routing methods often use synthetic or reward-model labels, multiple inference calls, or fixed router and model configurations. This work emphasizes human preference data, multiple router architectures, data augmentation, and out-of-domain evaluation.

  • Reward modeling evaluates generated responses afterward, whereas routing selects the LLM beforehand based on query complexity and model capabilities.
  • LLM-Blender, Frugal-GPT, and AutoMix query multiple LLMs during inference, which can increase latency.
  • Hybrid-LLM uses synthetic preference labels and a single BERT-based router, whereas this work uses Chatbot Arena human preferences and explores multiple router architectures.
  • This work evaluates on several decontaminated public benchmarks to demonstrate out-of-domain generalization, unlike Hybrid-LLM’s evaluation on the MixInstruct test split.
  • Zooter uses QwenRM reward-model labels and a fixed set of LLMs, while this approach relies mainly on human preferences and demonstrates adaptability across model pairs.

3 LLM ROUTING

The routing formulation learns a binary decision from preference data and converts predicted strong-model wins into weak- or strong-model selection. Evaluation metrics quantify both call cost and recovered performance across routing thresholds.

  • 3.1 PROBLEM FORMULATION: The framework learns a binary routing function from preference data comparing strong- and weak-model responses for each query.
  • 3.1 PROBLEM FORMULATION: The win prediction model estimates the probability that the strong model will outperform the weak model for a query.
  • 3.1 PROBLEM FORMULATION: The cost threshold α routes to the weak model when P(wins | q) < α and to the strong model when P(wins | q) ≥ α.A higher α favors weak-model routing more often, while a lower α favors the stronger model.
  • 3.1 PROBLEM FORMULATION: The router model returns the response from whichever strong or weak model the routing function selects for the query.
  • 3.2 METRICS: Cost efficiency is measured by the percentage of calls sent to strong models, which are more expensive than weak models.
  • 3.2 METRICS: Performance gap recovered measures how much of the quality difference between weak and strong models the router recovers.
  • 3.2 METRICS: APGR summarizes performance-gap recovery over different cost constraints by varying the routing threshold α.The call-performance graph is evaluated over discretized strong-model call percentages.
  • 3.2 METRICS: Call-performance threshold CPT(x%) is the minimum percentage of strong-model calls needed to achieve a target performance-gap recovery of x%.For Figure 1, CPT(50%) is approximately 37% of GPT-4 calls.

4 METHODOLOGY

The methodology builds routing supervision from human and augmented preference data, then evaluates several approaches for predicting whether strong or weak models will win. These predictors support routing decisions based on query similarity, learned model–query scores, or contextual query representations.

  • Chatbot Arena Data: 80k Chatbot Arena battles provide human preference labels from user votes comparing anonymous model responses.The resulting dataset contains queries, responses, and pairwise comparison labels based on human judgment.
  • Chatbot Arena Data: Model-tier clustering addresses label sparsity by assigning top-tier models to the strong class and third-tier models to the weak class.The method retains only winner identities rather than model responses in the resulting preference dataset.
  • Data Augmentation: Two augmentation strategies supplement sparse human preferences: golden labels from MMLU and LLM-judge labels for open-ended chat.MMLU comparisons are derived by checking strong and weak model responses against known correct answers, while GPT-4 supplies pairwise judgments in the judge-labeled approach.
  • Routing Approaches: Similarity-weighted ranking uses query embeddings and training-query similarity to weight examples before learning Bradley-Terry coefficients for strong and weak models.The similarity score determines a weight scalar for each training query, and the learned coefficients estimate the strong model’s win probability.
  • Routing Approaches: Matrix factorization learns a bilinear model–query scoring function that represents response quality and is optimized on preference data.Model and query embeddings are projected into aligned dimensions before producing a scalar score.
  • Routing Approaches: Neural routing alternatives include a BERT classifier and a causal Llama 3 8B classifier that predicts win probabilities from contextual query representations.The BERT approach uses a contextualized query embedding, whereas the causal classifier predicts comparison labels as next tokens.

5 EXPERIMENTS

Experiments show that preference-data routers improve cost-quality trade-offs across MT Bench, MMLU, and GSM8K, especially after dataset augmentation. The routers also transfer across unseen model pairs, while routing overhead remains small relative to generation cost.

  • Experimental setup: 65k pairwise comparisons from Chatbot Arena, spanning 64 models and over 100 languages, form the primary router-training data.The data excludes prompts shorter than 16 characters and holds out 5k samples for validation.
  • Benchmark results: 8.8 at MT Bench CPT(50%) reaches 95% of GPT-4’s 9.3 score, while augmented routers reduce costs by up to 75% versus random routing.Matrix factorization requires half as many GPT-4 calls as random routing for a PGR of 50%.
  • Benchmark results: Augmented preference data substantially improves router performance across MT Bench, MMLU, and GSM8K.Augmentation uses GPT-4 judge data for MT Bench and GSM8K and golden-label MMLU data for MMLU.
  • Benchmark results: Approximately 1,500 additional MMLU samples, representing less than 2% of training data, enable routers to require about 20% fewer GPT-4 calls than random at CPT(50%).Routers trained only on Arena data perform poorly because most MMLU questions are out-of-distribution.
  • Benchmark results: Augmentation moves GSM8K routing from APGR below random to above random, with the best classifier requiring 17% fewer GPT-4 calls than random at CPT(50%) and CPT(80%).The causal LLM classifier performs best among the evaluated routers on the augmented GSM8K data.
  • Adaptability across models: The same routers generalize across Claude 3 and Llama 3.1 model pairs without retraining, with best routers requiring approximately half as many GPT-4 calls as random at CPT(80%).The replacement strong and weak models were absent from the training data, and performance remained comparable to the original pair.
  • Data similarity and cost analysis: Benchmark-dataset similarity correlates with router performance, while overall routing costs reach up to 3.66x savings and the most expensive router adds no more than 0.4% overhead.GPU-based routers are more efficient than CPU-based routers, although the authors note room for optimization.

6 CONCLUSION

The routers perform strongly across diverse benchmarks, reduce costs while preserving response quality, and retain performance across strong/weak model pairs without retraining. Dataset augmentation improves routing, but benchmark mismatch and the two-model setting remain important limitations.

  • Routers achieve strong performance across benchmarks spanning open-ended question answering, humanities, and math problems.
  • Dataset augmentation enables routers to outperform the random baseline across all benchmarks, with the largest gains when training data resembles evaluation data.
  • Real-world applications may differ substantially from evaluated benchmarks, motivating in-domain data collection for specific use cases.
  • The study focuses on two-model routing, while extending the approach to multiple models remains future work.

B DATA CONTAMINATION

The evaluation removes examples that may overlap with training preference data. Contamination is detected through embedding similarity against all training examples using a fixed threshold.

  • Evaluation and training examples are compared using embeddings generated by OpenAI’s text-embedding-3-small model.
  • For each evaluation example, similarity search across all training data uses a threshold of 0.95 to identify contaminated examples.
  • Contaminated evaluation examples are discarded before reporting uncontaminated scores.

C BENCHMARK-DATASET SIMILARITY

The benchmark-dataset similarity score measures how closely evaluation prompts match preference-data samples. It uses each prompt’s maximum similarity because a few highly similar examples are considered most valuable for routing.

  • S(B, Dpref) is defined as the average maximum similarity between each benchmark prompt and samples in a preference dataset.
  • The score compares embeddings of benchmark prompts with embeddings of preference-dataset samples.
  • Maximum similarity is used because a small number of highly similar preference examples is considered more valuable than many less-similar examples.

D COST CALCULATION

The cost analysis estimates routing economics under GPT-4 and Mixtral pricing assumptions for short, single-turn prompts. GPT-4 is substantially more expensive per million tokens than Mixtral 8x7B under these assumptions.

  • The analysis assumes routers mostly handle short prompts in a single-turn setting, with average inputs of 95 tokens and outputs of 264 tokens.
  • The estimated average cost of GPT-4 is approximately 24.7 USD per 1 million tokens.
  • For Mixtral 8x7B, the analysis assumes equal input and output pricing, yielding an average cost of $0.24 USD per 1 million tokens.

E INDEPENDENT BENCHMARKS

The routers are evaluated against commercial routing systems on MT Bench, with results indicating stronger performance at comparable cost and fewer GPT-4 calls.

  • Figure 2 reports that the proposed routers achieve stronger performance than existing routers for the same cost on MT Bench.
  • The evaluated comparison includes Unify AI and Martian, two existing commercial offerings for LLM routing.Figure 2 presents the best-performing routers on MT Bench against these systems.
  • The benchmark routes between gpt-4-turbo-2024-04-09 as the stronger model and either Mixtral-8x7B-Instruct-v0.1 or Llama-2-70B-Chat as the weaker model.
  • Up to 40% fewer calls are routed to GPT-4 while achieving the same performance as commercial routing systems.This result applies to matrix factorization and causal LLM routers trained on Darena +Djudge.

F ADDITIONAL PLOTS

This section provides additional plots for the Section 5.1 results, covering router performance on MT Bench, MMLU, and GSM8K.

  • Additional plots include MT Bench performance for all routers.
  • Additional plots include 5-shot MMLU performance for all routers.
  • Additional plots include 8-shot GSM8K performance for all routers.
Loading 2406.18665v4…