Reducing Class-Wise Performance Disparity via Margin Regularization

TL;DR

MR$^2$ addresses substantial class-wise performance disparity even on balanced data by introducing margin regularization at two levels: per-class logit margins and an intra-class representation margin. Grounded in a class-sensitive generalization bound, the method selects per-class margins $\gamma_y$ proportional to the feature spread, while the representation-margin loss reduces intra-class variability to tighten the bound. Theoretical results include a $\bm{\gamma}$-margin risk bound and an optimal-margin corollary, along with a demonstration that reducing mean-squared deviation improves generalization. Empirically, MR$^2$ consistently improves hard-class accuracy across seven datasets and diverse backbones, while avoiding penalties on easy classes, thereby reducing disparity and improving overall performance; the approach generalizes to various norm settings and backbone models, making it a practical, principled tool for fairer, more reliable classification.

Abstract

Deep neural networks often exhibit substantial disparities in class-wise accuracy, even when trained on class-balanced data, posing concerns for reliable deployment. While prior efforts have explored empirical remedies, a theoretical understanding of such performance disparities in classification remains limited. In this work, we present Margin Regularization for Performance Disparity Reduction (MR$^2$), a theoretically principled regularization for classification by dynamically adjusting margins in both the logit and representation spaces. Our analysis establishes a margin-based, class-sensitive generalization bound that reveals how per-class feature variability contributes to error, motivating the use of larger margins for hard classes. Guided by this insight, MR$^2$ optimizes per-class logit margins proportional to feature spread and penalizes excessive representation margins to enhance intra-class compactness. Experiments on seven datasets, including ImageNet, and diverse pre-trained backbones (MAE, MoCov2, CLIP) demonstrate that MR$^2$ not only improves overall accuracy but also significantly boosts hard class performance without trading off easy classes, thus reducing performance disparity. Code is available at: https://github.com/BeierZhu/MR2

Figures (11)

• Figure 1: (a) Deep models exhibit severe class-wise accuracy disparity on class-balanced dataset such as ImageNet. (b) Feature diversity imbalance: feature distribution of "hard" classes is more diverse than that of "easy" classes. Classes are sorted in descending order of per-class accuracy.
• Figure 2: Plot of 0-1 loss, ramp loss and $\Phi_\gamma$.
• Figure 3: Weight-norm $\|\mathbf{w}_k\|_2$ on ImageNet.
• Figure 4: Relative improvements on fine-grained datasets: StanfordCars, OxfordPets, Flowers, Food and FGVCAircraft with CLIP ResNet-50. Numerical values are provided in Table \ref{['tab:other-datasets']}.
• Figure 5: Effect of $\bar{c}$.

Theorems & Definitions (18)

• Definition 1
• Definition 2
• Lemma 1
• Lemma 2
• Lemma 3
• Proposition 1
• Corollary 1
• Proposition 2: General bound on the $\bm{\gamma}$-margin Rademacher complexity
• proof
• proof