FlexRouter: Learning Complementary Model Sets for Flexible LLM Routing

AI summary

Overview

Research area: Machine learning / large language model serving — specifically LLM routing, the problem of choosing which model(s) to invoke for a given query.

Technical level: Advanced. The paper builds on Determinantal Point Processes (DPPs), submodular optimization, and Schur complements, and assumes familiarity with LLM serving pipelines and benchmark evaluation.

Scope in one sentence: The paper proposes FlexRouter, a DPP-based router that selects complementary (rather than individually top-ranked) sets of LLMs to maximize the probability that at least one selected model answers correctly, and evaluates it on the RouterEval benchmark with 3811-model and 5000-model candidate pools.

What This Paper Is About

Most LLM routers score each candidate model independently and then take the top-k models. This ignores correlations between models: models trained on similar data or with similar architectures tend to fail on the same queries, so a top-k set can be redundant and fail as a group. FlexRouter instead treats routing as a coverage-oriented subset selection problem — it selects a set of models that is individually strong but mutually diverse, so that at least one model in the set is likely to be correct. The paper's motivating example (Figure 2) shows top-k routing returning three near-identical answers for coin-flip, car-wash, and palindrome questions, while the proposed method returns three different reasoning paths or implementations.

Key Contributions

1. Problem reformulation. LLM routing is cast as coverage-oriented subset selection: the reward is the indicator that the selected subset intersects the query's "correct set" of models, R(x, S) = 1 if S ∩ C_x ≠ ∅.
2. DPP-based routing framework (FlexRouter). A query-dependent kernel L_x = diag(q(x)) · K · diag(q(x)) is built from query-conditioned competence scores q_i(x) = σ(v_xᵀ u_i) and a model-similarity matrix K, so diagonal entries favor strong models and off-diagonal entries penalize redundant pairs.
3. A coverage training objective. Because no unique ground-truth subset exists, the authors derive a tractable loss by marginalizing over failure sets: the probability a sampled subset lies entirely inside the failure set F_x is det(I + (L_x)_{F_x}) / det(I + L_x), giving the "hit" loss L_hit(x) = −log(1 − that quantity), combined with an auxiliary binary cross-entropy term weighted by λ.
4. Adaptive greedy inference. Instead of a fixed budget, a greedy MAP procedure adds the model with the largest marginal log-determinant gain and stops when that gain falls below a threshold τ relative to the first gain, yielding variable-size subsets per query.

Main Findings

• Higher coverage on the medium-pool setting (3811 candidate LLMs). FlexRouter reaches an average Success@10 of 0.8632, versus 0.8329 for EmbedLLM and 0.8320 for BaRP; the single-model reference average is 0.619. Per-task, FlexRouter scores 0.9471 on BBH, 0.8075 on MATH, 0.8487 on GPQA, 0.9390 on IFEval, and 0.7738 on MuSR.
• Higher coverage on the large-pool setting (5000 candidate LLMs). FlexRouter reaches an average Success@10 of 0.9914, versus 0.9802 for EmbedLLM and 0.9142 for BaRP; the reference average is 0.855. It obtains the highest score on five of six tasks: 0.9907 (MMLU), 0.9851 (HellaSwag), 0.9886 (GSM8K), 0.9915 (ARC), 0.9951 (TruthfulQA), 0.9976 (WinoGrande).
• One task exceptions. EmbedLLM achieves the highest MATH score in the medium-pool setting (0.8264 vs 0.8075 for FlexRouter), which the authors attribute to MATH containing a small number of strong specialist models; EmbedLLM also leads on ARC in the large-pool setting (0.9957 vs 0.9915).
• Substantially more diverse selected subsets. On the medium-pool setting, FlexRouter's average ILD@10 is 0.7620, versus 0.2112 for EmbedLLM and 0.1982 for BaRP. On the large-pool setting the average is 0.6802, versus 0.2810 for EmbedLLM and 0.5521 for BaRP. The authors note BaRP's diversity is nearly constant across tasks (e.g., 0.5520–0.5521 across all six large-pool tasks), suggesting a fixed model set rather than query-adaptive selection.
• Out-of-domain generalization. FlexRouter achieves the best Success@10 on the unseen tasks in both settings (IFEval and MuSR in the medium pool; TruthfulQA and WinoGrande in the large pool) and selects more diverse subsets there as well.
• Kernel choice matters mainly for diversity. With a Gaussian RBF kernel instead of cosine similarity, large-pool in-domain Success@10 rises slightly (0.9920 vs 0.9854) but in-domain ILD falls (0.6019 vs 0.6474) and out-of-domain ILD falls (0.6923 vs 0.7456). On the medium pool, cosine wins on in-domain success (0.9058 vs 0.8804) and both diversity figures (0.7392 vs 0.7109 in-domain; 0.7961 vs 0.7883 out-of-domain), while RBF is slightly better out-of-domain in success (0.8757 vs 0.8486). Cosine is adopted as the default.
• Flexible cost via the stopping threshold. On the combined in-domain test set of the large-pool setting with k_max = 10: at τ = 0.000 to 0.050 the average subset size is 10.00 and Success@10 is 0.9885; at τ = 0.100 the size is 9.95 with Success@10 0.9883 (0.5% size reduction, 0.02% coverage drop); at τ = 0.200 the size drops to 6.41 with Success@10 0.9772 (35.9% size reduction, 1.13% coverage drop); at τ = 0.300 the size is 4.25 with Success@10 0.9562 (57.5% reduction, 3.24% drop); at τ = 0.500 the size is 1.41 with Success@10 0.8828 (85.9% reduction, 10.57% drop). The authors identify τ ≈ 0.2 as a favorable operating point.
• Theoretical guarantees. Propositions in the appendix establish that log det((L_x)_S) is submodular but not monotone, that greedy selection with the stopping rule produces a set no superset can improve upon in determinant, and that minimizing the hit loss equals maximizing the total DPP probability mass on subsets containing at least one correct model.

Methodology in Plain English

Each query is encoded with a shared text encoder into a 128-dimensional vector. Every candidate LLM has its own learnable embedding. Multiplying the query vector by a model's embedding (and passing through a sigmoid) gives that model's predicted competence for the query. Comparing model embeddings by cosine similarity gives a static picture of which models are alike. These two signals are combined into a matrix whose diagonal says "how good is this model here" and whose off-diagonal says "how similar are these two models."

That matrix parametrizes a Determinantal Point Process — a probability distribution over subsets that automatically favors sets whose members are both high quality and dissimilar. The problem is that the training data does not say which subset is "correct"; it only records, for each query and each model, whether that model got the answer right. So instead of fitting a target subset, the authors compute the probability that a randomly drawn subset lands entirely inside the set of models that failed, and train the router to make that failure probability small — equivalent to maximizing the chance that at least one selected model is right. A supplementary binary cross-entropy loss on the per-model correctness scores keeps training stable.

At test time, exact MAP inference is NP-hard, so the router greedily adds the model with the largest marginal gain to the subset's determinant (computed efficiently with a Schur complement). It stops when the best remaining candidate's gain falls below τ times the first gain. Because the log-determinant objective is non-monotone, adding a highly correlated model can actually reduce the score, so the stopping rule naturally halts when candidates offer nothing new. Training uses RoBERTa-base as the encoder, Adam for up to 100 epochs, λ = 1.0, a maximum subset size k = 10, and NVIDIA A100 80GB GPUs.

Why This Matters

Impact on research. The paper challenges the dominant independent-scoring design in LLM routing and supplies both a principled objective (answer coverage) and a concrete learning algorithm that does not require ground-truth target subsets. It connects LLM routing to DPP-based diverse subset selection and to results on submodular maximization, and it reframes the evaluation question from "rank models well" to "build a pool that contains a correct answer."

Real-world applications.

• Multi-model inference pipelines that generate several candidate answers and pass them to a verifier, reward model, LLM-as-judge reranker, or aggregation scheme.
• Cost-aware serving where easy queries should use few models and hard queries should get more, controlled by a single threshold τ.
• Human-in-the-loop assistants where a user picks among several suggested answers, benefiting from genuinely different reasoning paths rather than three paraphrases of the same output.
• Routing over large, fast-changing model pools (thousands of candidates), where the per-query subset size is decided automatically rather than fixed in advance.

Industry relevance. Serving stacks that host many models need exactly this trade-off: coverage per unit of compute. The paper's reported operating point — at τ = 0.200, a 35.9% reduction in average subset size for a 1.13% coverage drop — is the kind of knob an operator could tune. The adaptive stopping rule means no fixed top-k budget has to be chosen per task.

Caveats the authors themselves stress: Success@k is a routing-stage coverage metric, not final deployed accuracy; a lone correct answer is not guaranteed to be recovered by a downstream selector. Diversity is measured with fixed model embeddings, which does not guarantee output-level semantic diversity. Cost is proxied by average subset size, ignoring model-specific latency, token price, output length, batching, and hardware constraints.

Future Directions

• End-to-end evaluation with a concrete selector. Pair FlexRouter with a specific verifier, reranker, aggregation method, or judge and measure final answer accuracy rather than coverage.
• Response-level diversity and agreement analysis. RouterEval supplies correctness labels but not full generated responses for all query–model pairs, so output-level semantic diversity remains unmeasured.
• Cost-aware routing. Extend the objective to account for model-specific latency, token pricing, output length, batching, and hardware, instead of treating subset size as a proxy.
• Controlling cost on hard queries. The ethics statement notes that adaptive routing can increase cost for difficult queries if not properly controlled, leaving open how best to bound that risk.

Also deferred to appendices but not reported in the available content: the Success@k curves, Avg-Correct@10, and Zero-Correct Rate analyses (Appendix D.1), the Random-k, MaxDiversity, and MMR baselines and their results (Appendix D.2), full dataset details (Appendix C.1), additional baseline details (Appendix C.2), further metric details (Appendix C.3), and the computational complexity analysis (Appendix C.4).

Target Audience

Researchers and engineers working on LLM serving, model routing, and inference-cost optimization; practitioners building multi-model or ensemble-style generation pipelines with downstream verification; and readers with a background in submodular optimization, determinantal point processes, or diverse subset selection who want to see those tools applied to LLM routing at the scale of thousands of candidate models.

Authors’ abstract

Existing Large Language Model (LLM) routing methods score LLMs independently to select top-$k$ models. However, this ignores model correlations and enforces a rigid computational budget. Consequently, routers often select redundant models that share failure modes, limiting the overall probability of success. To address this, we propose FlexRouter, a routing framework that explicitly models model complementarity. FlexRouter optimizes for \textit{answer coverage}, maximizing the probability that at least one selected model yields a correct response. This objective aligns with practical inference pipelines where multiple candidate outputs are generated and a downstream verifier or user selects the final one. We formulate routing as a coverage-oriented subset selection problem and model the routing policy using Determinantal Point Processes (DPPs), which naturally capture both model competence and redundancy. To directly optimize coverage without requiring a ground-truth target subset, we introduce a training objective based on marginalizing over failure sets. During inference, we employ a greedy strategy based on marginal log-determinant gains, enabling the router to adaptively determine subset sizes without a predefined budget. Extensive experiments on the large-scale RouterEval benchmark demonstrate that our proposed FlexRouter achieves higher coverage with lower redundancy across both in-domain and out-of-domain tasks than strong baselines while maintaining flexible inference cost.