Scalable Prompt Routing via Fine-Grained Latent Task Discovery
Organizations: Amazon Web Services
Abstract
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.
Figures & tables
| Models | Per-Task Evaluation | Overall | ||||||||||
| NQ | Triv-QA | Com-QA | mmlu | arc | OB-QA | gsm8k | math | hum-eval | mbpp | Ave | Qual | |
| Candidate Models | ||||||||||||
| Llama-3.3-70B | 57.6 | 23.2 | 76.0 | 87.4 | 90.5 | 90.7 | 89.2 | 94.0 | 21.1 | 40.0 | 67.0 | 0.446 |
| Llama-4-Maverick | 53.8 | 22.7 | 79.9 | 93.0 | 98.8 | 95.7 | 94.0 | 95.2 | 21.1 | 86.7 | 74.1 | 0.580 |
| Claude-Haiku-4.5 | 46.6 | 21.6 | 79.8 | 85.2 | 98.8 | 95.5 | 83.8 | 88.1 | 73.7 | 86.7 | 76.0 | 0.572 |
| Claude-Sonnet-4.5 | 59.4 | 24.7 | 85.7 | 94.7 | 97.6 | 97.4 | 88.7 | 91.1 | 63.2 | 93.3 | 79.6 | 0.621 |
| Models | Per-Task Evaluation | Overall | ||||||||||
| NQ | Triv-QA | Com-QA | mmlu | arc | OB-QA | gsm8k | math | hum-eval | mbpp | Ave | Qual | |
| Stage 1 Only | ||||||||||||
| Embedding | 58.5 | 24.8 | 84.5 | 92.5 | 95.2 | 95.5 | 88.7 | 94.3 | 42.1 | 86.7 | 76.3 | 0.609 |
| Coarse CLS | 60.3 | 25.9 | 80.4 | 92.5 | 98.8 | 95.0 | 92.8 | 93.0 | 63.2 | 93.3 | 79.5 | 0.623 |
| Fine CLS | 59.3 | 25.1 | 82.7 | 93.3 | 98.8 | 95.7 | 91.8 | 93.6 | 63.2 | 93.3 | 79.7 | 0.637 |
| Stage 2 Only | 57.4 | 24.3 | 81.8 | 92.2 | 97.6 | 95.5 | 90.7 | 94.0 | 52.6 | 83.3 | 77.0 | 0.647 |
| Class ID | Task Description | Selected Candidates |
| 109 | Answer multiple-choice questions requiring inference, detail identification, numerical calculation, and statement verification. | DeepSeek-v3 |
| 267 | Formal symbolic mathematics: Solve algebra, geometry, number theory, linear algebra, and calculus problems. | Llama-4-Maverick, Llama-3.3-70B |
| 161 | Telephone area code queries: geographic locations, regional identifiers, and telecom history. | Qwen3-235B, GPT-OSS-120B, DeepSeek-R1 |
LLM Routing
Large language model routing selects models, tools, experts, or inference strategies to balance answer quality, cost, latency, and deployment constraints.
See all 349 papers