Regularized Discriminative Alignment for Deep Representations under Label Shift
Abstract
Label shift refers to the distribution shift scenario where the marginal label distribution changes while the class-conditional distribution remains invariant. To address this challenge in complex real-world settings, we propose Regularized Discriminative Alignment for Label Shift (RDALS), a novel framework that adapts to target domains by aligning distributions within the deep latent space. By shifting the focus from raw inputs to learned representations, RDALS effectively operates under a weaker and more practical invariance assumption. Specifically, we construct a moment-matching linear system using Linear Discriminant Analysis (LDA) and show that this choice maximizes numerical stability. We further provide rigorous theoretical analysis, establishing finite-sample error bounds for the importance weight estimation and the generalization bounds for the adapted classifier. Extensive experiments on standard benchmarks demonstrate that RDALS significantly outperforms state-of-the-art baselines, achieving superior robustness and accuracy in both data-scarce and extreme-shift regimes.
Lay Summary
Machine learning models often fail in real-world deployment when the frequency of observed categories changes—a challenge known as "label shift." For instance, a disease might be rare in training data but common in a new hospital. Existing methods to correct this are fragile; they rely heavily on noise-sensitive raw inputs or perfect initial predictions, which falter when data is scarce. To tackle this, we introduce Regularized Discriminative Alignment for Label Shift (RDALS). Instead of analyzing raw data, RDALS leverages the robust "deep representations" extracted by pre-trained models. We use Linear Discriminant Analysis (LDA) to project these features into a specialized subspace that maximizes the separation between classes, effectively filtering out irrelevant noise. We then apply a stabilized mathematical solver to precisely estimate the new category proportions. Our research provides a highly stable, "plug-and-play" solution that helps AI adapt to extreme distribution shifts without requiring complex calibrations. By maintaining high accuracy even in data-scarce scenarios, RDALS ensures AI systems remain reliable for critical applications like healthcare and fraud detection.