Source-linked AI summary

Scalable Prompt Routing via Fine-Grained Latent Task Discovery

Yunyi Zhang, Soji Adeshina, Sheng Guan, Ashwin Ganesh, Zhen Han, Vassilis N. Ioannidis, Huzefa Rangwala, George Karypis

arXiv:2603.19415v2cs.CLcs.AIcs.LG

TL;DR

Large model pools contain subtle, task-dependent capability differences that challenge manual taxonomies and monolithic routers. FineRouter discovers latent tasks automatically and applies task-aware quality estimation; across 10 benchmarks and 11 frontier models, it outperforms baselines and the strongest individual model at less than half its inference cost.

  • Problem

    Large frontier-model pools create fine-grained routing distinctions that manual taxonomies and monolithic routers struggle to capture.

  • Method

    FineRouter combines graph-based latent task discovery and prompt assignment with mixture-of-experts quality estimation using task-specific prediction heads.

  • Results

    Across 10 benchmarks and 11 frontier models, FineRouter consistently outperforms routing baselines and the strongest individual model while costing less than half as much for inference.

  • Takeaways & Limitations

    Explicitly modeling latent task structure enables more accurate and scalable routing decisions as model pools expand with subtle capability differences.

  • Takeaways & Limitations

    Task discovery depends on LLM-generated descriptions and reward-model scores, which may introduce bias or miss distinctions in specialized or creative domains.

Abstract

from arXiv · show

Prompt routing dynamically selects the most appropriate large language model from a pool of candidates for each query, optimizing performance while managing costs. As model pools scale to include dozens of frontier models with narrow performance gaps, existing approaches face significant challenges: manually defined task taxonomies cannot capture fine-grained capability distinctions, while monolithic routers struggle to differentiate subtle differences across diverse tasks. We propose a two-stage routing architecture that addresses these limitations through automated fine-grained task discovery and task-aware quality estimation. Our first stage employs graph-based clustering to discover latent task types and trains a classifier to assign prompts to discovered tasks. The second stage uses a mixture-of-experts architecture with task-specific prediction heads for specialized quality estimates. At inference, we aggregate predictions from both stages to balance task-level stability with prompt-specific adaptability. Evaluated on 10 benchmarks with 11 frontier models, our method consistently outperforms existing baselines and surpasses the strongest individual model while incurring less than half its cost.

1 Introduction

FineRouter addresses the difficulty of routing among many frontier models with subtle, task-dependent capability differences. It discovers fine-grained latent tasks, specializes quality estimation by task, and achieves stronger cost-performance outcomes than existing approaches.

  • Large language models have heterogeneous strengths, motivating prompt routing to select an appropriate candidate model per query while managing computational cost.
  • Existing routers scale poorly because manual coarse taxonomies miss nuanced capabilities, while monolithic estimators struggle with subtle distinctions across diverse tasks.
  • FineRouter discovers fine-grained task types through offline graph-based clustering and trains a classifier to assign prompts to those tasks.
  • Its second stage uses task-specific mixture-of-experts prediction heads, and inference combines both stages to balance task-level stability with instance-level adaptability.
  • 10 benchmarks and 11 frontier models show consistent gains over routing baselines, stronger performance than individual candidates, and less than half the strongest candidate's inference cost.

2 Preliminaries

Prompt routing can be formulated as quality-based model selection, but accurate estimation becomes harder as candidate pools grow and task-specific strengths diverge. The paper therefore introduces latent task types and task-aware routing to improve accuracy while preserving inference efficiency.

  • Prompt routing selects the most appropriate model from a candidate set for each input prompt according to a desired objective.
  • Quality-based routing treats model selection as regression over estimated response quality rather than brittle single-label classification.
  • Evaluating every candidate at inference is computationally prohibitive, so routing learns a quality estimator that predicts the best model without generating all responses.
  • As the candidate count grows, monolithic quality estimation becomes challenging because different models excel on different fine-grained task types.
  • Task-aware routing learns both a prompt-to-task assignment function and a task-aware routing function to use task-specific knowledge efficiently.

3 Methodology

FineRouter discovers fine-grained task types from prompt semantics and model-preference patterns, then uses task-aware quality estimation to route prompts. Its two-stage design combines specialized task signals with broader model coverage for efficient model selection.

  • 3.1 Stage 1: Task Type Discovery and Matching: FineRouter uses graph-based clustering to discover task types from semantic prompt similarity and similarity in preferred-LLM rankings.The prompt graph connects nearest neighbors by embedding similarity, filters edges with RBO, and weights retained edges using normalized semantic and preference signals.
  • 3.1 Stage 1: Task Type Discovery and Matching: Leiden community detection iteratively refines clusters into task types containing semantically similar prompts with similar preferred LLMs.Cluster centers summarize task-description embeddings, while rank fusion summarizes preferred-model rankings; prompts outside meaningful communities remain unassigned.
  • 3.1 Stage 1: Task Type Discovery and Matching: For each discovered task, FineRouter adaptively selects candidate LLMs to cover the cluster’s preferred models rather than using a fixed top-k.The candidate set is derived by combining ranked lists across prompts and maximizing coverage over the cluster’s preferred models.
  • 3.1 Stage 1: Task Type Discovery and Matching: A text classifier matches new prompts to discovered task types using a bi-linear architecture with prompt and task-description encodings.Multi-label training supports matching prompts to multiple task types with different confidence scores, or to no specific task when scores are low.
  • 3.2 Task-Aware Quality Estimation: Stage 2 replaces a monolithic router with a mixture-of-experts quality estimator that combines general model adapters and task-specific prediction heads.Task-specific adapters specialize in candidate models for each discovered task, while general adapters provide predictions for all models.
  • 3.2 Task-Aware Quality Estimation: Assigned prompts invoke task-specific adapters for selected candidates and general adapters for remaining models, preserving specialized estimates with comprehensive model coverage.The router follows the same inference structure during training, allowing both adapter types to contribute to model-quality prediction.

4 Experiments

FineRouter is evaluated against routing baselines and 11 frontier-model candidates across 10 benchmarks, where narrow capability gaps make model selection difficult. It consistently achieves the strongest overall performance, uses multiple models, and benefits from combining fine-grained task discovery with task-aware routing.

  • Experimental Setup: 10 benchmarks cover question answering, multiple-choice, mathematical reasoning, and code generation, while the candidate pool contains 11 recent frontier models.The setup includes models from multiple families and capability profiles.
  • Experimental Setup: Frontier-model routing is challenging because performance gaps are narrower and more nuanced than in settings centered on smaller open-source models.The baselines and candidate pool are designed to reflect real-world deployment scenarios involving high-performing models.
  • Experiment Results: FineRouter consistently outperforms routing baselines across the majority of tasks and achieves the strongest overall performance among routers and individual models.The comparison spans 10 diverse benchmarks and shows that no single LLM is best on every task.
  • Experiment Results: FineRouter routes 28% of prompts to Claude-Sonnet-4.5, 27% to DeepSeek-R1, 23% to Llama-4-Maverick, and 13% to Qwen3-235B.The distribution indicates substantial use of several high-performing candidates rather than one model alone.
  • Experiment Results: FineRouter’s cost-performance curve consistently dominates baseline routers across all cost points.The curve compares average quality based on Q∗ with estimated total API cost while varying a tolerance parameter.
  • Ablation Studies: The full two-stage architecture consistently outperforms Stage-1-only and Stage-2-only variants, while fine-grained task types outperform manually defined coarse categories.The ablations attribute complementary routing signals to task discovery and task-aware specialized prediction heads.
  • Ablation Studies: 332 discovered task types reduce the model pool to ∼32%, with an average of 3.55 candidates per task.Case studies show distinct model recommendations for symbolic mathematics, general word problems, and telephone area-code queries.

5 Related Work

Related work develops prompt routers that balance model quality and cost using cascades, embeddings, reward supervision, contrastive learning, and graph-based task relationships. These approaches motivate FineRouter’s focus on fine-grained distinctions among model capabilities.

  • Prior Prompt-Routing Methods: Prompt-routing methods dynamically select models to balance performance and cost, including cascaded selection, embedding-based routing, reward-supervised ranking, and contrastive learning.Examples include FrugalGPT, HybridLLM, Zooter, and RouterDC.
  • Graph-Based Routing: GraphRouter connects queries to pre-defined coarse-grained task categories and LLMs, whereas FineRouter targets automatically discovered fine-grained task structure.The contrast concerns how task structure is represented for routing.

6 Conclusion

FineRouter addresses scalable routing among frontier models with narrow performance gaps through automated latent-task discovery and specialized quality estimation. Across 10 benchmarks and 11 frontier models, it consistently outperforms baselines and the strongest individual model at less than half its inference cost.

  • Conclusion: FineRouter combines graph-based discovery of fine-grained latent tasks with a mixture-of-experts router using specialized quality estimates.The architecture is presented as a two-stage approach to scalable prompt routing.
  • Conclusion: 10 benchmarks and 11 frontier models show consistent improvements over routing baselines and superiority to the strongest individual model at less than half its inference cost.The conclusion links explicit latent-task modeling with more accurate and scalable routing decisions.

Limitations

The method has limitations involving supervision quality, clustering assumptions and tuning, training-data requirements, and text-only input scope.

  • Task discovery depends on LLM-generated descriptions and reward-model scores, which may bias or miss distinctions in specialized or creative tasks.The implementation uses Claude-Sonnet-4.5 and Skywork-Reward-V2 as supervision sources.
  • Graph-based clustering requires tuned hyperparameters and assumes semantic and model-preference similarity forms meaningful task communities.Optimal settings may vary across domains or model pools, and the community assumption may not hold for all prompt distributions.
  • Training requires both a task classifier and a mixture-of-experts router, creating computational overhead and requiring responses from all candidate models.These response requirements may be infeasible in some deployment scenarios.
  • The current approach handles text-only prompts and does not yet cover multimodal inputs.The paper identifies multimodal task discovery as future work.

Task Description Generation

The offline task-discovery prompt asks an LLM to produce short task descriptions for user prompts, using tagged outputs and broad task coverage instructions.

  • The prompt asks the LLM to generate a short description of the task intended by a user prompt.
  • The generated answer must place its final task description inside <task> tags.
  • The instructions require one short sentence covering different aspects of the task, including task types.
  • Figure 3 presents this prompt for offline task type discovery.

A Prompt for Task Description Generation

Figure 3 shows the prompt used to generate domain and task descriptions for training prompts during offline task type discovery.

  • The prompt generates domain and task descriptions for training prompts.
  • These descriptions are produced during the offline task type discovery process.
  • The process is described in Section 3.1.1.

B Implementation Details

The implementation specifies graph-clustering settings, pretrained encoders and adapter training, compute requirements, and inference-time aggregation and reward-model choices.

  • Stage 1 uses Claude-Sonnet-4.5 descriptions, k = 5 nearest neighbors, RBO threshold τ = 0.4, three Leiden iterations, and coverage threshold δ = 0.8.
  • The task classifier uses STELLA-EN-400M-V5 and is fine-tuned for 10 epochs at learning rate 2e-5 with batch size 8.
  • Stage 2 initializes its prompt encoder from QWEN3-EMBEDDING-0.6B and uses 512-dimensional embeddings with 2-layer MLP adapters.
  • The pipeline trains a base model for 10 epochs, then task-specific adapters for another 10 epochs while freezing other parts.Training uses batch size 8 on 8 NVIDIA A100 GPUs.
  • Inference uses aggregation weight α = 0.5, with SKYWORK-REWARD-V2-LLAMA-3.1-8B as the reference quality function.

C Task Classifier Analysis

The task classifier identifies discovered task types in a 332-class problem and assigns most test prompts to one. Assigned prompts achieve higher quality scores than those routed to the general “Others” category.

  • 0.643 macro F1 is achieved on held-out validation for the 332-class cluster-assignment problem.
  • 71% of test prompts are assigned to a discovered task type, covering 22,485 of 31,774 prompts.
  • 0.665 quality score for assigned prompts exceeds 0.619 for prompts in the “Others” category.
Loading 2603.19415v2…