Optimized Deferral for Imbalanced Settings

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 LLM routing tasks.

• The paper presents mild, a novel deferral algorithm that leverages cost-sensitive margin losses to optimize expert routing in imbalanced settings.
• It provides theoretical H-consistency guarantees and explicit generalization bounds for cost-sensitive multi-class classification.
• Empirical evaluations show reduced deferral loss and improved expert allocation in applications like image classification and LLM routing.

Optimized Deferral for Imbalanced Settings: Technical Review and Implications

Problem Formulation and Motivation

The paper "Optimized Deferral for Imbalanced Settings" (2604.27723) addresses the two-stage learning to defer problem with multiple experts, focusing on scenarios exhibiting significant expert imbalance. The core objective is to devise algorithms that optimize both prediction accuracy and computational cost by intelligently routing instances to specialized experts, especially when some experts dominate the input domain, either due to superior coverage or lower cost. This setting is prevalent in medical diagnosis, image classification, and LLM routing, where a majority expert often overshadows specialized, costlier counterparts.

The authors reframe deferral loss minimization as a cost-sensitive learning problem over the input-expert domain, moving beyond traditional label-centric imbalance handling. The distinction is critical: expert imbalance is orthogonal to class imbalance, as illustrated by the toy 1-D scenario where the majority expert dominates independent of label distribution.

Figure 1: Four class densities and three experts with varying accuracies and costs, visually demonstrating input-expert domain imbalance.

Theoretical Contributions

Cost-sensitive Margin Losses and Guarantees

A significant advancement in the paper is the derivation of novel cost-sensitive margin losses explicitly tailored for the expert imbalance scenario. These losses leverage instance-dependent rewards inversely proportional to expert costs and are formalized via margin-based expressions. The analysis yields:

• Generalization bounds for imbalanced cost-sensitive multi-class classification, leveraging class-sensitive Rademacher complexity and vector contraction lemmas.
• Explicit regularization-based algorithms that select margin parameters based on distributional statistics, yielding improved generalization guarantees. The theoretical optima for margin parameters (ρ_j's) are computed as functions of the empirical expert allocation and feature mapping norms.

The foundational result is a hypothesis-set-dependent consistency guarantee (H-consistency) for the proposed surrogate losses. This guarantee proves that minimizing the surrogate deferral loss leads to minimization of the target deferral loss, tightly relating the excess estimation errors via $\sqrt{2}$-scaling bounds. The result strengthens previous Bayes-consistency guarantees, which are agnostic to hypothesis class constraints and asymptotic minimization.

Algorithmic Development

Margin-based Imbalanced Learning to Defer (mild)

Leveraging the theoretical framework, the authors introduce mild (“Margin-based Imbalanced Learning to Defer”), a new algorithm that minimizes a convex surrogate loss reflecting both error and expert-specific cost. The optimization objective is:

• Regularized minimization over the hypothesis set, with margins and loss penalties modulated according to expert cost and empirical distribution.
• Explicit routing of inputs to experts such that the expert allocation ratio approaches the distributional optimum, avoiding the collapse to majority experts.

Surrogate objectives are constructed to unify input-expert optimization within the standard input-label framework, facilitating practical learning with efficient solvers.

Empirical Evaluation

Image Classification and Real-world LLM Routing

The empirical evaluation rigorously compares mild to the tdef baseline (the only prior method with H-consistency for two-stage deferral), using both synthetic experts (with perfect domain accuracy) and real experts (trained on data partitions).

Key findings:

• Deferral loss: mild consistently attains lower deferral loss than tdef across several benchmarks (CIFAR-10, CIFAR-100, SVHN, Tiny ImageNet), with improvements ranging from ~0.01 to ~0.04 in absolute loss.
• Expert allocation: mild preserves the optimal allocation ratios, effectively deferring inputs to specialized experts even when their coverage is minimal.

On LLM routing for MMLU, mild achieves almost optimal proportions between strong (7B), medium (1.5B), and tiny (0.5B) models. Notably, tdef collapses to exclusive majority expert usage in imbalanced settings, while mild adapts according to cost and error structure.

Practical and Theoretical Implications

Theoretical implications include the formalization of cost-sensitive imbalance in expert routing as a margin-based learning problem, extending generalization theory to input-expert domains. The practical implication is that mild enables cost-aware, accuracy-preserving routing in systems with heterogeneous expert pools, avoiding wasteful overuse of generalist models and enabling efficient utilization of specialized, high-cost predictors.

mild’s robust empirical performance in severe imbalance conditions further demonstrates its applicability to real-world model ensembles, particularly for scalable LLM routing and automated decision-making systems with budget constraints.

Future Directions

Potential extensions include:

• Generalization to structured prediction settings with even more complex expert-label-cost dependencies, leveraging margin-based bounds.
• Adaptive margin parameter selection in an online regime, addressing non-stationary expert distributions.
• Exploration of theoretical guarantees under non-complete hypothesis sets and scenarios with privacy-restricted expert feedback.

Conclusion

"Optimized Deferral for Imbalanced Settings" presents rigorous theoretical and algorithmic advances for two-stage deferral under expert imbalance. By recasting the problem as a cost-sensitive margin-based classification task and providing strong H-consistency guarantees, the authors enable principled, efficient expert routing in domains where cost and coverage heterogeneity are pronounced. Empirical results substantiate mild’s superiority, especially in real-world LLM routing and severe imbalance settings, providing a scalable blueprint for future multi-expert systems.