Flexible Routing via Uncertainty Decomposition

TL;DR

This paper introduces a new uncertainty-aware routing method for machine learning systems that intelligently decides when to use a low-cost model, an expensive oracle, or abstain, based on decomposing uncertainty components.

Contribution

It presents a novel uncertainty decomposition technique enabling dynamic, cost-sensitive routing and abstention without retraining, with theoretical regret guarantees.

Findings

Effective in reducing unnecessary oracle calls on ambiguous queries.

Adapts to different loss functions and cost parameters easily.

Demonstrates improved routing performance on synthetic and real datasets.

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.