No Single Best Model for Diversity: Learning a Router for Sample Diversity

1. Problem Statement

Large Language Models (LLMs) are increasingly required to generate diverse, high-quality responses for open-ended tasks (e.g., creative writing, brainstorming, code generation). Existing approaches typically focus on improving the diversity of a single model by adjusting inference hyperparameters (temperature, top-p) or modifying prompts.

However, the authors identify a critical gap: No single LLM dominates in generating a comprehensive set of diverse and high-quality answers across all types of queries. Different models exhibit complementary strengths, excelling at different subsets of prompts. Furthermore, existing evaluation metrics often fail to jointly quantify the coverage (how much of the valid answer space is covered) and quality of a generated set, particularly for open-ended questions where ground truth is infinite.

The core problem addressed is: How can we efficiently select the optimal LLM (or set of LLMs) for a specific query to maximize the diversity coverage of the generated answer set?

2. Methodology

A. New Evaluation Metric: Diversity Coverage (Div-Cov)

To evaluate the performance of generating diverse answer sets, the authors introduce Diversity Coverage.

• Definition: It measures the total quality score assigned to all unique answers in a predicted set $A$, relative to the best possible answer set of the same size.
• Formula:
  $$\text{div-cov}(q, A) = \frac{\sum_{a \in \text{uniq}(q, A)} \text{quality}(q, a)}{\text{max-uniq-sum}(q, B)}$$
  Where:
  • $\text{uniq}(q, A)$ is the set of semantically distinct answers.
  • $\text{quality}(q, a)$ is a scalar score (via reward models or ground truth comparison).
  • $\text{max-uniq-sum}$ is the theoretical maximum score achievable with $B$ distinct answers.
• Significance: Unlike previous metrics that count unique outputs or measure pairwise similarity, Div-Cov accounts for both the number of unique answers and their individual quality, making it suitable for both fixed-answer and open-ended domains.

B. Pilot Study: The "Oracle" Upper Bound

Before training a router, the authors conducted a pilot study using 18 LLMs (from Llama, Qwen, OLMo, and Gemma families) across three datasets (Simple Questions, NB-Curated, NB-WildChat).

• Finding: No single model consistently achieves the highest diversity coverage across all queries.
• Oracle Result: If one could magically select the best model for every single query (Oracle), diversity coverage on NB-WildChat could reach 33.0%, compared to 23.8% when using the single best overall model. This reveals a significant performance gap (approx. 29% relative improvement) achievable through model selection.

C. The Router Architecture

Motivated by the Oracle results, the authors propose training a Router to predict the best model for a given query without exhaustively sampling from all models.

• Task Formulation: Framed as a classification task where the input is a query $q$ and the output is the index of the best model $i^*$.
• Training Data: Constructed by sampling answers from all 18 models, calculating their Div-Cov scores, and labeling the model with the highest score as the ground truth.
• Router Models:
  • KNN: Simple non-parametric classifier.
  • BERT: Fine-tuned with a classification head.
  • MLP (Multi-Layer Perceptron): The best-performing architecture. They experimented with:
    • Model-Agnostic (agn): Using a pre-trained retriever (inf-retriever-v1) to encode the query.
    • Model-Specific (spec): Using the hidden states of the candidate model $m_i$ itself to encode the query (hypothesized to capture how the model "decodes" the input).
• Classification Strategies:
  • Multi-way Classification: Predicting the single best model index directly.
  • Binary Classification: Training a separate binary classifier for each model to predict if it is the best for a given query, then selecting the model with the highest score.

D. Ensembling Strategy

The router can also be used to select two models per query. The generation budget is split (e.g., 25 answers from Model A, 25 from Model B), allowing for a hybrid approach that leverages complementary strengths.

3. Key Results

Performance on In-Domain (NB-WildChat)

• Router vs. Baseline: The best router (Binary MLP with model-specific encoding) achieved 26.3% diversity coverage, significantly outperforming the "Top Overall" single model baseline (23.8%).
• Oracle Gap: While the router does not reach the Oracle (33.0%), it closes a substantial portion of the gap with minimal computational overhead compared to exhaustive sampling.
• Ensembling: Using the router to select the top 2 models per query further improved performance to 26.7%.

Generalization (Out-of-Domain)

• The router trained on NB-WildChat generalized well to NB-Curated (OOD), achieving 39.3% coverage compared to the baseline's 38.6%.
• Scaling: Increasing the training data size (from 500 to 2k samples) consistently improved router performance, suggesting the approach scales well with data.

Efficiency Analysis

• Oracle Cost: Exhaustively sampling all 18 models to find the best one incurs a 19x computational overhead.
• Router Cost: The router adds only a fixed, small inference cost (2-3x slower than a single baseline model) because it predicts the model before sampling. This makes it highly practical for real-world deployment.

Prompting Strategies

The study explored different prompting templates ("Generate One," "Generate Two," "Generate All").

• Trade-off: "Generate All" (listing all answers in one go) yielded the highest diversity coverage but resulted in lower average answer quality for later-generated items.
• Router Robustness: Routers trained on specific prompt types did not generalize well to other prompt types, indicating that the "best model" is context-dependent on the generation strategy.

4. Key Contributions

1. Diversity Coverage Metric: Introduced a unified metric that jointly evaluates the diversity and quality of answer sets for both fixed and open-ended questions.
2. Router for Diversity: Proposed and validated a learning-based router that dynamically selects the optimal LLM for a query to maximize diversity coverage, demonstrating that heterogeneous ensembles outperform single models.
3. Empirical Insights:
   • Demonstrated that no single model dominates diversity generation across all queries.
   • Showed that model-specific query encodings provide better routing signals than model-agnostic ones.
   • Highlighted the trade-off between diversity and answer quality in different prompting strategies.

5. Significance

This work shifts the paradigm from optimizing a single model's generation to orchestrating a suite of models. It proves that:

• Complementarity is Key: Different models have distinct "diversity niches."
• Efficiency: A lightweight router can approximate the performance of an expensive Oracle (exhaustive search) with a fraction of the cost.
• Future Direction: The findings lay the foundation for "Multi-LLM Systems" where models are not just used for accuracy but are dynamically combined to solve problems requiring broad coverage of the solution space (e.g., creative brainstorming, scientific discovery, and synthetic data generation).

The authors release their code and dataset, facilitating further research into multi-model collaboration and diversity-aware routing.