Optimized Deferral for Imbalanced Settings
| Authors | Corinna Cortes et al. |
| Year | 2026 |
| Field | Machine Learning |
| arXiv | 2604.27723 |
| Download | |
| Categories | cs.LG, stat.ML |
Abstract
Learning algorithms can be significantly improved by routing complex or uncertain inputs to specialized experts, balancing accuracy with computational cost. This approach, known as learning to defer, is essential in domains like natural language generation, medical diagnosis, and computer vision, where an effective deferral can reduce errors at low extra resource consumption. However, the two-stage learning to defer setting, which leverages existing predictors such as a collection of LLMs or other classifiers, often faces challenges due to an expert imbalance problem. This imbalance can lead to suboptimal performance, with deferral algorithms favoring the majority expert. We present a comprehensive study of two-stage learning to defer in expert imbalance settings. We cast the deferral loss optimization as a novel cost-sensitive learning problem over the input-expert domain. We derive new margin-based loss functions and guarantees tailored to this setting, and develop novel algorithms for cost-sensitive learning. Leveraging these results, we design principled deferral algorithms, MILD (Margin-based Imbalanced Learning to Defer), specifically suited for expert imbalance settings. Extensive experiments demonstrate the effectiveness of our approach, showing clear improvements over existing baselines on both image classification and real-world Large Language Model (LLM) routing tasks.
Engineering Breakdown
Plain English
This paper addresses a critical problem in learning to defer systems: when multiple expert models (like different LLMs or classifiers) are available, the deferral mechanism often becomes biased toward majority experts, degrading overall performance. The authors conduct a comprehensive study of two-stage learning to defer in imbalanced expert settings, where you have access to existing pre-trained predictors and must decide which inputs to route to which expert. They reframe the deferral loss function to handle expert imbalance explicitly, enabling algorithms to make better routing decisions that don't just favor the most confident or most frequently used expert. The result is a framework that balances accuracy gains with computational efficiency, particularly valuable in high-stakes domains like medical diagnosis and NLG where routing to the right specialist matters more than using the cheapest expert.
Core Technical Contribution
The core novelty is formalizing and solving the expert imbalance problem in two-stage learning to defer. Prior work assumed roughly balanced access to experts or didn't account for skewed expert utilization patterns. The authors cast the deferral loss in a way that explicitly penalizes relying too heavily on any single expert, encouraging more equitable use of available specialists. This is a loss function and algorithmic innovation, not a new architecture—they provide a principled framework for weighting expert predictions and deferral decisions when experts have different reliability profiles or availability constraints. The key insight is that imbalance in expert coverage leads to suboptimal routing, and this can be mitigated by reformulating the learning objective itself.
How It Works
The system operates in two stages: first, you have a collection of pre-trained expert models (classifiers, LLMs, etc.) with varying performance characteristics. Second, you train a deferral policy that takes an input and decides to either route it to a specific expert or defer it to a human reviewer. The deferral loss function is reformulated to account for expert imbalance by imposing constraints or penalties on how often each expert is selected—if one expert is being overused, the loss increases. During training, the algorithm learns a scoring or confidence function that estimates which expert should handle each input, adjusted for the empirical usage distribution of experts. At inference, for each input, the system either selects the highest-confidence expert recommendation or defers to human judgment, with the routing decision optimized to minimize error while respecting computational budgets and expert capacity constraints.
Production Impact
In production, this approach lets you effectively combine multiple models without one dominant expert drowning out others. If you're building a system with multiple LLM APIs (OpenAI, Anthropic, open-source), a local classifier, and human review, you'd use this framework to learn which model to query for each input type, avoiding the trap of always using the most accurate or fastest model. This is directly applicable to medical diagnostic systems where you might have a pathology AI, a radiologist AI, and human experts—the algorithm ensures each is used appropriately rather than over-relying on one. The trade-off is that you need to retrain the deferral policy when your expert pool changes, and you incur compute cost by calling multiple experts to build confidence estimates. Latency improves if deferral reduces human review rates, but integration complexity increases because you must manage fallbacks, expert availability, and possibly batch requests across multiple services.
Limitations and When Not to Use This
The paper assumes you have access to fixed, pre-trained expert models and doesn't address how to handle expert failures, latency variability, or dynamic expert availability. It also doesn't fully address the data requirements—you need labeled examples of when each expert should be used, which is expensive to collect in many domains. The framework may not scale well if you have dozens or hundreds of experts, as the optimization complexity could explode. Additionally, the paper likely doesn't deeply explore failure modes where all experts are weak on a particular input type, or where human reviewers themselves introduce bias; deferral can only help if at least one expert or human is reliable.
Research Context
Learning to defer builds on decades of work in selective prediction and abstention, where classifiers learn to say 'I don't know.' This paper extends that to multi-expert settings, which has become urgent with the rise of large pretrained models that teams want to ensemble or specialize. It likely benchmarks on standard datasets (CIFAR, ImageNet, medical imaging) and possibly proprietary NLG or medical diagnosis datasets. The work sits at the intersection of ensemble learning, cost-sensitive learning, and fairness—it's related to problems of equitable classifier usage and resource allocation. This opens research directions in online deferral (adapting expert routing as performance drifts), expert selection under latency constraints, and theoretical guarantees on deferral loss in imbalanced regimes.
:::tip Subscribe Get weekly breakdowns of papers like this in AI Letters - the newsletter for engineers building production AI systems. :::