Flexible Routing via Uncertainty Decomposition

Abstract: A key strategy for balancing performance and cost in modern machine learning systems is to dynamically route queries to either a low-cost model or a more expensive oracle (such as a large pretrained model or human expert), an approach known as model routing. In this work we present a new uncertainty-aware router that (1) avoids unnecessary oracle calls on inherently ambiguous queries, and (2) adapts dynamically to different loss functions and cost parameters through simple hyperparameter changes, without retraining. Our method, applicable to any classification setting where multiple independent annotations per input are available, is based on decomposing total uncertainty into irreducible and reducible components using higher-order predictors [Ahdritz et al., 2025]. This enables a unified approach to both routing and abstention: predict with the weak model when uncertainty is low, route to the oracle when reducible uncertainty is high, and abstain when irreducible uncertainty is high. Our router comes with strong theoretical guarantees bounding regret relative to optimal task-specific routers. We conduct experiments on both synthetic and real-world datasets that demonstrate the benefits of our approach in suitable regimes -- in particular, whenever reducible and irreducible uncertainty are not too correlated.

• The paper introduces higher-order calibration for uncertainty decomposition to differentiate between reducible and irreducible losses in ML query routing.
• It employs a bin-and-estimate approach to adapt routing decisions to arbitrary loss functions and cost parameters without retraining.
• Empirical results show the method minimizes oracle costs with effective three-way decisions, outperforming traditional baselines on synthetic and real datasets.

Flexible Routing via Uncertainty Decomposition: An Expert Analysis

Motivation and Background

Modern ML deployments increasingly require flexible orchestration of queries between low-cost, weak models and more costly, high-fidelity oracles. Conventional routing approaches rely either on model confidence or on supervised loss predictors, but they are fundamentally limited. Confidence-based routing typically amplifies routing under high uncertainty without discriminating between irreducible (aleatoric) and reducible (epistemic) uncertainty, resulting in wasteful oracle invocations for queries that remain ambiguous even for the oracle. Supervised routers lack zero-shot adaptability to loss and cost changes, as every such change requires expensive retraining.

This paper introduces a unified framework that leverages higher-order calibration to decompose uncertainty, enabling principled routing decisions and flexible adaptation to various downstream loss functions and penalties without retraining.

Uncertainty Decomposition via Higher-Order Calibration

A novel aspect of this work is the operationalization of higher-order predictors, which output distributions over label distributions. Higher-order calibration, as defined here, ensures that each predicted mixture aligns with the empirical distribution of ground truth label distributions (e.g., from multi-annotator data). This calibration allows precise decomposition of a weak model's loss into:

• Irreducible loss: The baseline error incurred by the oracle, equivalent to the entropy of the label distribution.
• Reducible loss: The excess error incurred by the weak model relative to the oracle across the same input.

This decomposition is central for fine-grained routing: abstention is preferred when irreducible uncertainty dominates, routing is optimal when reducible uncertainty exceeds the oracle penalty, and prediction is preferred otherwise. The policy landscape is visualized in the following figure.

Figure 2: Routing performance comparison on synthetic data, showing the effectiveness of the uncertainty-decomposition approach relative to optimal and total-uncertainty baselines.

Theoretical Guarantees

The framework provides strong regret bounds relative to the optimal prediction-only router, i.e., any router whose action is a deterministic function of the weak model's output. Specifically, given a weak model that can be higher-order calibrated, the expected excess cost—over any proper loss and any routing/abstaining penalties—is tightly bounded in terms of the calibration error. Notably, this guarantee holds distribution-free and for arbitrary downstream configurations, substantiating the flexibility claim.

Practical Construction: Bin-and-Estimate Approach

Construction of a practical higher-order predictor proceeds via empirical mixtures over $k$-label snapshot calibration data, with partitioning of inputs into buckets based on prediction vectors (typically top-class and confidence) or auxiliary features. The empirical distribution in each bucket is used to estimate both reducible and irreducible losses, enabling rapid decision-making for new loss functions and cost parameters. Notably, the calibration statistics are independent of the downstream task configuration, supporting zero-shot flexibility.

Empirical Evaluation

Three key benefits are demonstrated experimentally:

1. Routing with Irreducible Uncertainty

On synthetic data with decoupled reducible and irreducible uncertainty, the proposed method significantly outperforms total-uncertainty baselines, routing only when oracle improvement is possible. This separation avoids wasteful oracle calls. Synthetic environments demonstrate that reducible and irreducible uncertainty, when not strongly correlated, enable substantial gains.

Figure 4: Sinusoidal synthetic function illustrating variable irreducible uncertainty for each input.

Figure 6: Routing curve demonstrating the superiority of the method when reducible and irreducible loss are uncorrelated.

On real datasets such as CIFAR-10H, SNLI, and ChaosNLI, however, calibrations reveal a high correlation between reducible and irreducible losses across model output buckets, and thus total-uncertainty-based routing yields similar performance post-calibration. The authors emphasize that decomposition is most valuable in regimes with substantive uncertainty separation, and recommend verifying this regime before deploying advanced routers.

2. Adaptability to Arbitrary Loss Functions

Empirical results show that supervised loss predictors perform robustly only for the fixed loss used in training and degrade markedly for novel downstream losses. The proposed method maintains near-optimal routing performance across a range of proper losses—demonstrating genuine zero-shot flexibility.

Figure 1: Routing curve for SNLI across multiple loss functions; performance of the supervised loss predictor collapses on losses differing from its training loss, while the higher-order calibrated method remains robust.

3. Seamless Extension to Multi-way Decisions

The method is naturally extensible to three-way decisions: predict, route, or abstain. It achieves lower expected loss than all two-way baselines (predict/route or predict/abstain only) on cost-penalized routing tasks.

Figure 7: Cost-penalized performance for three-way decision tasks on SNLI. The three-way router traces the lower envelope of two-way alternatives across penalty sweeps.

Implementation Nuances

For partition-based higher-order calibration, sample size per bucket is a critical tradeoff: finer partitions require more calibration snapshots for stable estimation. Centroid calibration of outputs within buckets ensures that weak model predictions are constant across each bin, aligning with the theoretical framework. Where the weak model is already near-optimal, this may discard useful signal and can be bypassed. Feature-based binning is shown to improve decomposition estimates and routing efficacy in synthetic scenarios with more nuanced uncertainty structures.

Theoretical and Practical Implications

This framework advances the state of uncertainty-aware orchestration and selective prediction, integrating abstention, routing, and flexible loss adaptation in a unified system. Distribution-free guarantees and zero-shot adaptation to loss and cost parameters directly address operational constraints in real-world ML pipelines (including LLM routing contexts, cost-sensitive annotation pipelines, and human-AI complementarity scenarios).

Theoretically, the development of higher-order calibration and its integration with loss-agnostic omniprediction expands the frontier for flexible inference and post-hoc adaptation. It also provides guidance for cases where uncertainty decomposition conferred by higher-order calibration does and does not deliver practical routing gains, motivating future research in high-quality higher-order calibration under limited snapshot calibration resources.

Future Directions

Future research should target scalable methods for learning higher-order calibrated predictors without requiring extensive snapshot data, more sophisticated partition schemes for non-trivial multiclass settings, and extension to complementarity scenarios with heterogeneous oracles. Integration with downstream orchestration and mixture-of-expert architectures, and more fine-grained modeling of abstention penalties, remain promising avenues.

Conclusion

The paper presents a rigorous, practical, and theoretically sound framework for flexible query routing in classification problems, operationalized via uncertainty decomposition using higher-order calibration. It enables principled decisions among prediction, routing, and abstention, supports rapid adaptation to new losses and cost structures, and produces robust empirical results in suitable regimes. This contributes meaningfully to both the theory and practice of ML orchestration and selective prediction, with broad implications for the deployment of cost-conscious, adaptive AI systems.