Task-Aware LLM Routing with Multi-Level Task-Profile-Guided Data Synthesis for Cold-Start Scenarios

1. Problem Statement

Large Language Models (LLMs) vary significantly in performance and computational cost depending on the specific task and query. LLM Routing aims to select the optimal model from a candidate pool for a given query to maximize a utility function balancing performance and cost.

However, existing routing systems face a critical Cold-Start Problem:

• Data Scarcity: Real-world deployments (e.g., new startups, niche domains) often lack in-domain labeled training data (query-performance-cost tuples) required to train routers.
• Poor Generalization: Pre-trained routers trained on one domain often fail to generalize to new domains due to distribution shifts.
• Limitations of Current Methods:
  • Rule-based/Adaptive methods are too rigid.
  • Classification-based routers struggle with fine-grained preference control.
  • Regression-based routers (predicting cost/performance directly) often overfit to surface-level query features and fail in cold-start scenarios without sufficient training data.
  • Direct LLM selection is computationally expensive and difficult to calibrate without ground truth.

2. Methodology

The authors propose a two-pronged solution: a Multi-Level Task-Profile-Guided Data Synthesis Framework to generate training data, and TRouter, a task-type-aware routing model.

A. Multi-Level Task-Profile-Guided Data Synthesis

To solve the cold-start data scarcity, the authors introduce a framework that synthesizes diverse Question-Answer (QA) pairs approximating the test-time query distribution using only high-level domain descriptions as input. The framework consists of three components:

1. Task Type Generator:
   • Constructs a hierarchical taxonomy with three levels: Domain $\rightarrow$ Subcategory $\rightarrow$ Difficulty.
   • Iteratively expands the taxonomy by generating child nodes based on parent definitions, ensuring non-overlapping and diverse categories.
2. Task Type Quality Evaluator:
   • Performs self-critique on generated task types to ensure specificity, completeness, and minimal redundancy.
   • Iteratively refines the set until stability is reached.
3. Question-Answer Pair Generator:
   • Uses the "Task Profile" (Domain + Subcategory + Difficulty definition) as a prompt to generate QA pairs.
   • Employs semantic similarity filtering (using sentence transformers) to discard near-duplicates, ensuring high diversity.
   • Output: A synthetic dataset ($D_{syn}$) where candidate LLMs are evaluated to generate ground-truth cost and performance metrics for training.

B. TRouter: Task-Type-Aware Routing

Instead of directly modeling $p(\text{metric} | \text{query}, \text{model})$, TRouter introduces a latent task-type variable ($t$) to disentangle task semantics from surface query features.

• Theoretical Formulation:
  The conditional distribution is decomposed as:
  $$p(h | q, m) = \sum_{t \in T} p(h | t, m) p(t | q)$$
  Where $h$ is the metric (cost or performance), $q$ is the query, $m$ is the model, and $t$ is the latent task type.
• Model Architecture:
  1. Task Recognition Module ($q_\phi(t|q)$): Predicts the probability distribution of task types for a given query. It is trained using Cross-Entropy loss against a prior derived from the synthesized taxonomy.
  2. Metric Prediction Module ($p_\theta(h|t, m)$): Predicts cost and performance conditioned on the task type and the specific model.
  3. Training Objective: The model optimizes an Evidence Lower Bound (ELBO), combining a reconstruction loss (MSE for metrics) and a regularization term (KL-divergence) to align the predicted task distribution with the synthesized taxonomy prior.
• Inference: For a new query, the router marginalizes over the task space to predict performance/cost for all candidate models and selects the one maximizing the user-defined utility function.

3. Key Contributions

1. First Identification of Cold-Start in LLM Routing: The paper formally addresses the lack of in-domain data for router training and proposes a solution.
2. Novel Data Synthesis Framework: Introduces a hierarchical, task-profile-guided approach to generate synthetic training data that approximates real-world query distributions without manual annotation.
3. TRouter Architecture: Proposes a regression-based router enhanced with latent task-type variables and prior regularization, which captures deep semantic structures beyond surface-level query features.
4. Comprehensive Validation: Demonstrates effectiveness across multiple LLM pools (Qwen, Gemini, Doubao), datasets (Math, Coding, Reasoning), and evaluation protocols (Traditional metrics vs. LLM-as-a-Judge).

4. Experimental Results

Experiments were conducted in both Cold-Start (trained only on synthetic data) and In-Domain (trained on real labeled data) settings.

• Cold-Start Performance:
  • TRouter trained on synthetic data ($D_{syn}$) significantly outperforms baselines like Adaptive LLM (rule-based) and Prompt LLM.
  • It even surpasses cross-domain baselines (e.g., RouterDC, MetricRouter trained on other tasks) which suffer from severe domain shift.
  • Under "LLM-as-a-Judge" evaluation, TRouter trained on synthetic data achieved utility scores comparable to or better than models trained on real in-domain data.
• In-Domain Performance:
  • TRouter consistently outperforms state-of-the-art baselines (including GraphRouter, MetricRouter, and C2MAB-V) across diverse user preferences (Cost-first, Balanced, Performance-first).
  • It shows higher robustness and lower variance compared to competing methods.
• Ablation Studies:
  • Taxonomy Level: Using Difficulty-level task profiles for data synthesis yields the best results, as it provides the most granular control and reduces redundancy compared to Domain or Subcategory levels.
  • Task Regularization: Removing the task-type prior degrades TRouter to a standard MetricRouter, confirming the importance of the latent variable.
  • Sample Efficiency: TRouter achieves near-optimal performance with as few as 5 shots per task type, demonstrating high data efficiency.

5. Significance and Impact

• Democratizing LLM Routing: By eliminating the need for expensive, manual in-domain data collection, this framework enables startups and researchers to deploy efficient, cost-aware LLM routing systems immediately upon product launch.
• Interpretability: Unlike black-box embedding routers, TRouter's reliance on a hierarchical taxonomy allows for human inspection, bias correction, and incremental updates (e.g., adding new task branches without retraining the whole system).
• Cost-Performance Trade-off: The method provides a scalable way to model complex cost-performance frontiers across different model sizes and task difficulties, ensuring users get the best value for their specific needs.
• Robustness: The approach proves that synthetic data, when guided by structured task profiles, can effectively simulate real-world distributions, overcoming the limitations of traditional data augmentation.

In conclusion, this work bridges the gap between theoretical routing optimization and practical deployment constraints by providing a robust, data-efficient, and interpretable framework for cold-start LLM routing.