RouteLLM: Learning to Route LLMs with Preference Data

Original Paper: https://arxiv.org/abs/2406.18665v2

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

Abstract:

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.

Summary Notes

Figure: Routing performance/cost trade-off between GPT-4 and Mixtral-8x7B. (left) We demonstrate several routers that outperform the random baseline on OOD eval GSM8K. (center) We demonstrate improvement in router performance through data augmentation, denoted by (A), on MT Bench. (right) We display the main metrics we consider: call-performance threshold (CPT, denoted in green) and average performance gain recovered (APGR, denoted by the blue shaded region).

Introduction

Large Language Models (LLMs) have revolutionized natural language processing with their impressive capabilities in tasks ranging from open-ended conversation to code generation. However, deploying these models involves a significant trade-off between performance and cost. Powerful models like GPT-4 can deliver high-quality responses but come with prohibitive costs, while smaller models such as Mixtral-8x7B are more economical but may fail to handle complex queries effectively.

In this context, the research paper "RouteLLM: Learning to Route LLMs with Preference Data" presents an innovative approach to optimize the balance between cost and response quality. This blog post delves into the methodology, findings, and implications of this research, providing insights into how efficient routing between LLMs can be achieved.

Key Methodologies

The core idea behind RouteLLM is to develop router models that dynamically select between a stronger and a weaker LLM during inference. These router models are trained using human preference data and data augmentation techniques. Here’s a breakdown of the methodologies used:

Problem Formulation

The researchers framed the problem as a binary routing issue, where each user query is processed by a router model that decides whether to route the query to a strong model (e.g., GPT-4) or a weak model (e.g., Mixtral-8x7B). The objective is to minimize costs while maintaining a high level of response quality.

Router Training Framework

1. Human Preference Data: The training data consists of 80k battles from the Chatbot Arena, where users vote for the better response between two models. This data is clustered into tiers to reduce label sparsity and enhance training efficiency.
2. Data Augmentation: To overcome the issue of label sparsity, two data augmentation methods were employed:
• Golden-labeled datasets: These datasets, such as the MMLU benchmark, provide predefined metrics to calculate preference labels.
• LLM-judge-labeled datasets: GPT-4 was used as a judge to generate pairwise comparison labels, significantly enhancing the training dataset.

Routing Approaches

Several models were explored for learning the win prediction function:

1. Similarity-weighted (SW) Ranking: Utilizes a Bradley-Terry model to predict the winning probability based on query similarity.
2. Matrix Factorization: Leverages a bilinear function to capture the low-rank structure of query-model interactions.
3. BERT Classifier: Uses a BERT-base architecture for contextualized query embeddings.
4. Causal LLM Classifier: Employs an instruction-following paradigm to predict win probabilities via next-token prediction.

Main Findings and Results

The research demonstrated significant cost savings without compromising response quality across various benchmarks, including MT Bench, MMLU, and GSM8K. Key results include:

Performance on MT Bench

• Matrix Factorization and SW Ranking: These routers performed significantly better than the random baseline, achieving up to 60% cost savings.
• Augmented Data: When trained with augmented data from GPT-4 judge, the BERT and causal LLM classifiers showed substantial improvements, with the BERT classifier achieving over 50% improvement in average performance gain recovered (APGR).

Performance on MMLU and GSM8K

• Golden-labeled Data: Augmenting training data with golden labels from MMLU led to significant performance improvements, with routers requiring approximately 20% fewer GPT-4 calls for a given performance level.
• LLM-judge-labeled Data: Similar improvements were observed on GSM8K, with the causal LLM classifier achieving the best performance when trained on this augmented dataset.

Implications and Applications

The research presents a robust framework for optimizing LLM deployment in real-world applications, balancing cost and performance effectively. Key implications include:

• Cost-effective Deployment: By routing simpler queries to cheaper models and reserving complex queries for powerful models, significant cost savings can be achieved. This is particularly beneficial for applications with high query volumes.
• Scalability: The framework can be extended to multi-model routing, providing a scalable solution for diverse applications.
• Customizability: The benchmark-dataset similarity score offers a quantitative measure to enhance router performance for specific use cases, making the approach highly customizable.

Conclusion

RouteLLM represents a significant advancement in the efficient deployment of large language models. By leveraging human preference data and innovative training methodologies, the researchers have developed a cost-effective yet high-performance solution for real-world applications. As LLMs continue to evolve, such routing frameworks will be crucial in maximizing their potential while managing costs effectively.

Quote from the Research Paper: "Our approach significantly reduces costs—by over 2 times in certain cases—without compromising the quality of responses."

Future research can explore extending this framework to multiple models and further refining the routing algorithms for even greater efficiency and effectiveness.

Potential Areas for Future Research:

• Expanding the routing framework to handle more than two models.
• Investigating the reasons behind the varied performance of different router models on the same dataset.
• Exploring the application of this routing framework in other domains such as image processing and speech recognition.

By intelligently routing queries, RouteLLM paves the way for more efficient and cost-effective use of large language models, ensuring that high-quality responses are delivered without breaking the bank.