Technical Summary: From Sampled Outcomes to Capability Distributions: Rethinking Supervision for LLM Routing

1. Problem Statement

Current Large Language Model (LLM) routing methods typically rely on single-shot supervision, where a model's capability for a specific query is labeled based on a single generated response and its corresponding score. The authors argue that this approach is fundamentally flawed because LLM generation is inherently stochastic. A single sampled outcome provides only a noisy observation of a query-model pair's behavior rather than a reliable estimate of its true capability.

This "point-estimate" view introduces systematic noise into the routing supervision signal. Because LLM outputs vary due to sampling randomness (output-side uncertainty) and sensitivity to prompt phrasing (input-side uncertainty), a single sample may lead to an unstable or misleading label regarding which model is best suited for a query. Consequently, routers trained on such noisy labels may learn policies that reflect incidental generation noise rather than stable differences in model capability, leading to unreliable model selection in production.

2. Methodology: DARS (Distribution-Aware Routing Supervision)

To address these limitations, the authors propose DARS, a framework that constructs routing supervision from a distributional view of model behavior rather than isolated samples. DARS operates by capturing uncertainty from two distinct dimensions:

A. Data Construction

Instead of generating one response per query-model pair, DARS constructs a set of repeated observations:

1. Input-Side Uncertainty: The original query is rewritten multiple times using a strong language model to generate semantically equivalent variations (preserving the gold answer and task semantics). This captures sensitivity to query formulation.
2. Output-Side Uncertainty: For each rewritten query, the candidate model performs multiple independent stochastic decoding runs. This captures variability in reasoning paths, factuality, and final answers.

These steps result in an observation matrix (e.g., $N$ rewrites $\times$ $M$ decoding runs) for each query-model pair, recording both performance scores and inference costs for every observation.

B. Distributional Capability Estimation

From these repeated observations, DARS summarizes three key capability signals for each query-model pair $(x, m)$:

• Expected Quality ($\mu_q$): The mean performance score across all observations.
• Expected Cost ($\mu_c$): The mean inference cost across all observations.
• Performance Risk ($\sigma_q$): The standard deviation of the performance scores, representing the aggregate variability caused by both input rewriting and stochastic decoding.

C. Risk-Aware Supervision

DARS constructs a risk-aware utility function to serve as the supervision target:
$$U(x, m) = \mu_q(x, m) - \lambda \mu_c(x, m) - \beta \sigma_q(x, m)$$
Where:

• $\lambda$ balances quality and cost.
• $\beta$ penalizes unstable behavior (high variance).

The preferred model for a query is defined as $m^*(x) = \arg\max_{m \in \mathcal{M}} U(x, m)$. This utility can be used as a hard label, a soft target (via softmax), or a continuous regression target, making DARS compatible with various router architectures (regression, classification, ranking, retrieval, etc.).

3. Key Contributions

1. Diagnosis of Single-Shot Limitations: The authors systematically analyze the single-shot label assumption, demonstrating through diagnostic metrics (outcome instability, winner flip rate, and disagreement with distribution-aware labels) that single-shot supervision is unstable and can directly alter model selection decisions.
2. DARS Framework: The introduction of a distribution-aware supervision framework that accounts for both input-side (prompt rewriting) and output-side (stochastic decoding) uncertainties to construct reliable routing labels.
3. Empirical Validation: Extensive experiments across diverse tasks (scientific reasoning, mathematical problem solving, reading comprehension) and heterogeneous model pools show that DARS yields more stable routing labels and improves the performance of learned routing policies compared to single-response supervision.

4. Experimental Results

The paper evaluates DARS on three datasets (GPQA, MATH-500, DROP-800) using a pool of six heterogeneous LLMs.

• Instability of Single-Shot Labels: Diagnostic analysis reveals high instability in single-shot supervision. For example, on GPQA, the "winner flip rate" (the frequency with which the best model changes between samples) is 0.970, indicating that a single sample is a poor predictor of the optimal model.
• Performance Gains: DARS improves routing performance across almost all tested router families (including MLP, kNN, EmbedLLM, GraphRouter, MIRT, etc.). The gains are particularly pronounced on tasks with high output-side uncertainty (e.g., GPQA).
• Ablation Studies:
  • Both input-only and output-only variants of DARS outperform the single-shot baseline, confirming that both sources of uncertainty are informative.
  • The full DARS framework (combining both) achieves the best results, suggesting the two sources are complementary.
  • Including the risk term ($\sigma_q$) significantly improves performance; removing it leads to lower average scores.
• Sample Efficiency: The study finds that a moderate number of observations (e.g., a $3 \times 3$ protocol) is sufficient to achieve performance close to the full $5 \times 5$ setting, suggesting DARS is practical for deployment without excessive sampling costs.
• Cost-Quality Trade-off: Routers trained with DARS achieve a superior Pareto frontier compared to single-shot routers and often outperform the best fixed-model baseline under various cost constraints.

5. Significance and Claims

The paper claims that reliable LLM routing must move beyond single-response observations. By grounding routing supervision in query-level model capability distributions, DARS provides a more robust signal for training routers. The authors emphasize that this approach does not require a new router architecture but rather a rethinking of the supervision signal itself. The results suggest that treating routing labels as stochastic samples from an underlying capability distribution leads to more stable, effective, and cost-efficient model selection policies.