Multi-LLM Routing: From Feature Engineering to Model Building

Multi-LLM Routing is the process of selecting and combining different Large Language Models (LLMs) to handle various tasks effectively, as no single LLM can perform optimally across all applications. This approach leverages the strengths of multiple models to balance performance and cost. The problem is crucial because as LLM usage expands, finding an efficient way to route tasks between models is key to ensuring optimal results while controlling resource consumption and costs, which are major concerns in large-scale deployments of LLMs.

Feature Engineering

To train a machine learning model for Multi-LLM Routing, where the input consists of a prompt and a list of candidate LLMs, the following features can be used:

1. Prompt Features:

• Prompt Length: The number of tokens or words in the prompt.
• Prompt Complexity: Measured by syntactic depth, vocabulary richness, or readability score.
• Domain Specificity: Whether the prompt contains domain-specific terms (e.g., medical, technical, legal language).
• Task Type: Classification of the prompt task (e.g., question answering, summarization, code generation, creative writing).
• Sentiment or Tone: Sentiment analysis of the prompt (positive, neutral, negative), or tone (formal, casual).
• Prompt Structure: Presence of multi-part questions, instructions, or lists.

2. LLM Model Features:

• Model Size: Number of parameters in the model (e.g., 7B, 13B, 175B).
• Training Data Domain: The type of data the LLM was trained on (e.g., open-domain, medical, legal, conversational).
• Model Latency: Response time or speed of inference for the LLM.
• Fine-tuning Specialization: Whether the LLM has been fine-tuned for specific tasks (e.g., customer support, summarization, code generation).
• Cost per Query: Computational cost or API pricing per inference for the LLM.
• Inference Accuracy: Historical performance or accuracy on similar prompts/tasks (if known).

3. Prompt vs. LLM Model Intersection Features:

• Task-Model Match: Whether the task type (e.g., summarization) aligns with the model’s strengths (e.g., a model fine-tuned for summarization).
• Vocabulary Alignment: Overlap between domain-specific terms in the prompt and the model's training data (measuring how well the model is expected to understand the prompt’s vocabulary).
• Model Response Complexity: Predicted output complexity (e.g., does a large prompt with complex reasoning align better with a larger, more powerful model?).
• Expected Latency-Task Fit: Balancing model inference speed with the complexity of the task (e.g., selecting faster models for simple prompts).
• Prior Task Success Rate: The historical success rate of the LLM model on similar prompts (if available), indicating how well it has performed on related tasks.

These features help in predicting which LLM from the candidate list is most suitable for handling a specific prompt, balancing performance and cost effectively.

To design a model for predicting the best LLM for a given prompt, considering response accuracy, running time, and cost, we can build a three-step pipeline. Here's a detailed design for model training and pipeline building, with focus on feature engineering, data collection, and deployment.

Pipeline Overview:

1. Step 1: Predict if an LLM can generate an acceptable/correct response Use a classification model to predict whether a given LLM will produce an accurate or acceptable response for a prompt.
2. Step 2: Predict the running time Train a regression model to estimate the expected running time of the LLM on the given prompt.
3. Step 3: Calculate the cost Based on the predicted running time and the cost-per-query, compute the total cost for running the LLM on the prompt.
4. Final Selection The system selects the LLM with an acceptable response, minimal running time, and lowest cost.

Model Training Design

1. Feature Engineering:

• Prompt Features: Length, complexity, domain specificity, task type, sentiment, structure, etc.
• LLM Features: Model size, training domain, response time (historical), fine-tuning details, cost per query, historical accuracy.
• Prompt vs LLM Intersection Features: Task-model match, vocabulary alignment, historical task success rates, expected latency-task fit.

2. Label Generation:

For the three-step prediction, generate labels as follows:

• For Step 1 (Correctness Prediction):Label: Binary (1 = correct/acceptable response, 0 = incorrect/unacceptable response). Generate labels by evaluating historical responses from LLMs. Use human annotations, automated evaluation metrics (e.g., BLEU, ROUGE, GPT score) to assess the quality of each response.
• For Step 2 (Running Time Prediction):Label: Numeric value representing the actual running time for the prompt on each LLM. Running time can be measured during inference, collecting the latency for each LLM.
• For Step 3 (Cost Calculation):Label: Cost based on the formula: Cost = Running Time * Cost per second (or API call).This cost can be derived from the known pricing models of LLM APIs.

3. Data Collection:

• Prompt-LLM Interaction Data: Collect prompts and responses across different LLMs by running prompts through multiple LLMs in real-world scenarios. Store the actual response, running time, and cost.
• Human Annotations: Get human evaluators to rate responses on correctness/acceptability for labeled data in Step 1.
• System Logs: Use logs from past interactions with LLMs to record response times and query costs, adding historical data to improve model training.

4. Model Training:

• For Step 1 (Correctness Model):
• For Step 2 (Running Time Model):
• For Step 3 (Cost Calculation):

5. Model Deployment:

• Deployment Strategy: Deploy the models as microservices, where each step (correctness prediction, running time prediction, and cost calculation) runs in sequence.Use containerization (e.g., Docker) to package the models, and orchestrate them using a service like Kubernetes for scalability.
• API Integration: Create an API endpoint where new prompts are fed into the pipeline. The system will return the selected LLM, along with expected running time and cost.Ensure that the system can access real-time cost updates from LLM providers to stay current on pricing.

Pipeline Flow:

1. Input: A prompt and a list of candidate LLMs.
2. Step 1: Pass the prompt and candidate LLMs to the Correctness Model. Predict whether each LLM can produce an acceptable response. Filter out LLMs predicted to give incorrect responses.
3. Step 2: For the remaining LLMs, pass the prompt and LLM features to the Running Time Model. Predict the running time for each model.
4. Step 3: Calculate the cost for each remaining LLM by multiplying the predicted running time with the cost per query.
5. Final Selection: Select the LLM that: Can provide an acceptable response.Has the lowest combination of response time and cost.
6. Output: Return the best LLM model, along with the expected time and cost.

Summary of Steps:

1. Feature Engineering: Extract prompt, LLM, and prompt-LLM intersection features.
2. Label Generation: Collect historical data on response accuracy, running time, and cost.
3. Data Collection: Gather data from real interactions with LLMs, annotate response correctness, and log response times and costs.
4. Model Training: Train classification and regression models for correctness prediction and running time estimation.
5. Deployment: Deploy the models as microservices, building an API for real-time inference and LLM selection based on correctness, time, and cost.