Optimized Deferral for Imbalanced Settings

2026-04-30

AI Summary

This work addresses the bias toward majority experts in two-stage “learning-to-defer” systems caused by imbalanced expert categories. It formalizes this issue for the first time as a cost-sensitive learning problem over the input-expert domain and introduces MILD, a theoretically grounded margin-based loss function together with a tailored optimization algorithm, to achieve more balanced expert routing within the two-stage deferral framework. Theoretical analysis provides strong generalization guarantees for the proposed loss. Extensive experiments on both image classification and large language model routing benchmarks demonstrate that MILD significantly outperforms existing methods, confirming its effectiveness and robustness.

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.

Problem

• learning to defer
• expert imbalance
• imbalanced settings
• cost-sensitive learning
• two-stage deferral

Innovation

• learning to defer
• expert imbalance
• cost-sensitive learning
• margin-based loss
• MILD