Technical Summary: MR2 (Margin Regularization for performance disparity Reduction)

Problem Statement

Deep neural networks frequently exhibit substantial disparities in class-wise accuracy, even when trained on class-balanced datasets. Empirical studies (e.g., Cui et al., 2024) indicate that "hard" classes (those with low accuracy) often possess greater feature variability compared to "easy" classes. This imbalance threatens the reliable and fair deployment of deep models. While prior works have explored empirical remedies such as data augmentation or representation learning, a theoretical understanding of why these disparities persist and how to mitigate them via principled regularization remains limited. Existing margin-based methods are primarily designed for long-tailed (class-imbalanced) distributions and often collapse to standard cross-entropy when class priors are uniform, failing to address disparities in balanced settings.

Methodology: MR2

The authors propose MR2, a theoretically grounded framework that reduces class-wise performance disparity by dynamically adjusting margins in both the logit space and the representation space. The method is built upon a novel class-sensitive generalization bound that links per-class feature variability to generalization error.

1. Theoretical Foundation

The authors derive a margin-based, class-sensitive learning guarantee (Proposition 1). This bound reveals that the generalization error is influenced by the ratio of per-class feature variability (specifically the sum of the squared norm of the feature mean and the mean squared deviation) to the square of the class-specific margin.

• Key Insight: Classes with larger feature variability ("hard" classes) contribute more to the generalization error bound. To minimize this bound, larger margins should be assigned to these classes.
• Optimal Margin: Theoretical analysis (Corollary 1) suggests that the optimal margin $\gamma_y$ for a class $y$ should be proportional to $(\|\hat{\mu}_y\|_2^2 + \|\hat{s}_y\|_2^2)^{1/3}$, where $\hat{\mu}_y$ is the empirical feature mean and $\hat{s}_y$ is the empirical mean squared deviation.

2. Logit Margin Regularization ($\ell_{\gamma,ce}$)

MR2 introduces a class-dependent margin into the cross-entropy loss. The margin $\gamma_y$ is not a fixed hyperparameter but is dynamically calculated based on the feature spread of each class:
$$\gamma_y = \bar{c} \cdot K \frac{(\|\hat{\mu}_y\|_2^2 + \|\hat{s}_y\|_2^2)^{1/3}}{\sum_{k=1}^K (\|\hat{\mu}_k\|_2^2 + \|\hat{s}_k\|_2^2)^{1/3}}$$
This adaptive mechanism assigns larger margins to "hard" classes (high feature variability), thereby tightening their generalization bounds, while maintaining a balanced average margin $\bar{c}$.

3. Representation Margin Regularization ($\ell_{\bar{s}}$)

To further tighten the generalization bound, MR2 penalizes excessive intra-class feature spread. The representation loss encourages intra-class compactness by treating the average mean squared deviation $\bar{s}$ as a margin constraint:
$$\ell_{\bar{s}}(f, x, y) = \ln \left[ 1 + \sum_{x^+ \in D_y \setminus \{x\}} \exp(\|\phi(x) - \phi(x^+)\|_2^2 - 2\bar{s}) \right]$$
This loss acts as a differentiable relaxation of a constraint that forces pairwise distances within a class to remain below $2\bar{s}$, effectively reducing the numerator in the generalization bound.

4. Generalization to Cosine Classifiers

The framework extends to models using cosine classifiers (e.g., CLIP), where features are $L_2$-normalized. In such cases, the standard $L_2$ norm of features becomes constant. The authors propose using alternative $L_p$ norms (e.g., $p=3$) to restore class-dependent variability measures, allowing the margin regularization to remain effective.

Key Contributions

1. Theoretical Analysis: The paper provides a novel, non-asymptotic, class-sensitive generalization bound that explicitly connects per-class feature variability to generalization error, motivating the need for class-dependent margins even in balanced settings.
2. MR2 Framework: A dual-margin regularization approach that dynamically adjusts logit margins based on feature spread and penalizes representation variance, applicable to both standard and cosine classifiers.
3. Empirical Validation: Extensive experiments across seven datasets (including ImageNet, CIFAR-100, and fine-grained datasets) and diverse architectures (ResNet, ViT, MAE, MoCov2, CLIP) demonstrate that MR2 significantly reduces performance disparity.

Experimental Results

The authors evaluate MR2 against 14 baselines, including reweighting methods (LfF, JTT), margin-based methods (LDAM, EQL), contrastive learning (SupCon), and distributionally robust optimization (DRO, FairDRO).

• Disparity Reduction: MR2 consistently achieves the largest improvements on "hard" classes. For example, on ImageNet with ResNet-50, MR2 improves "hard" class accuracy by 3.2% over the baseline, whereas other methods often show negligible gains or trade-offs (improving hard classes at the expense of easy ones).
• Overall Performance: Unlike many debiasing methods that sacrifice "easy" class performance to boost "hard" classes, MR2 improves overall accuracy while simultaneously narrowing the gap between easy and hard classes.
• Generalization: The method is robust across different backbones (CNNs, Transformers) and pre-training paradigms (training from scratch, linear probing on frozen features, fine-tuning foundation models).
• Ablation Studies: Results confirm that both the logit margin (class-dependent $\gamma$) and the representation margin ($\ell_{\bar{s}}$) are necessary; using uniform margins or removing the representation loss yields inferior results.

Significance and Claims

The paper claims that MR2 offers a principled explanation for class-wise performance disparity, moving beyond empirical heuristics. By linking generalization bounds to feature variability, the authors demonstrate that "hard" classes require larger margins to achieve comparable generalization performance to "easy" classes.

The significance of MR2 lies in its ability to:

1. Reduce Disparity without Trade-offs: It improves "hard" class performance without degrading "easy" class performance, a common failure mode in existing debiasing techniques.
2. Theoretical Grounding: It provides a mathematical justification for why class-dependent margins are necessary even when class priors are uniform, distinguishing it from long-tail learning methods.
3. Broad Applicability: It functions effectively across a wide range of modern architectures, including foundation models like CLIP and MAE, suggesting it is a viable strategy for improving fairness in large-scale deep learning systems.

The authors conclude that while Neural Collapse (NC) suggests feature variance should vanish in the terminal phase, empirical evidence on large-scale datasets shows that significant, class-dependent variability remains. MR2 addresses this reality by regularizing these residual spreads to ensure balanced generalization.