Reducing Class-Wise Performance Disparity via Margin Regularization

AI summary

Overview

Research area: Machine learning — classification theory, generalization bounds, and class-wise fairness/performance disparity in deep neural networks.

Technical level: Intermediate overall, with Advanced theoretical sections. The method itself is a small modification to the training loss and is easy to apply, but the paper's justification is built on Rademacher complexity and margin-based generalization bounds, which require a statistical learning theory background.

Scope: The paper proposes MR², a regularization scheme that sets per-class logit margins from measured feature spread and adds an intra-class compactness penalty, and shows both theoretically and empirically that this narrows the accuracy gap between "easy" and "hard" classes without sacrificing the easy ones.

What This Paper Is About

Deep networks trained on balanced datasets still produce wildly uneven accuracy across classes — the paper reports that on ImageNet, a ResNet-50 reaches 100% top-1 accuracy on its best class but only 16% on its worst. Prior work documented this disparity and tried empirical fixes, but lacked a theory explaining where it comes from. This paper argues that the root cause is unequal feature variability: hard classes have more spread-out feature distributions, and derives a margin-based learning guarantee that says exactly how to set per-class margins to compensate.

Key Contributions

1. A class-sensitive generalization bound. The authors derive Proposition 1, a margin-based bound on the true risk of a classifier that depends on per-class feature means and per-class mean squared deviations, showing that classes with larger feature variability contribute more to the error.

2. A closed-form optimal margin rule (Corollary 1). Under a fixed average margin budget c̄, the complexity term is minimized when each class margin γ_y is proportional to (‖μ̂_y‖₂² + ‖ŝ_y‖₂²)^(1/3), normalized across classes. This turns "give hard classes bigger margins" into a specific formula.

3. MR², a two-part regularizer. A logit-space loss ℓ_γ,ce (class-dependent temperature/margin) and a representation-space loss ℓ_s̄ that penalizes pairwise distances within a class beyond twice the average mean squared deviation, combined with weight λ as in Eq. 4.

4. A general L_p extension (Proposition 2 and Eq. 17). The analysis is extended beyond L₂ so that it stays informative for cosine-classifier models such as CLIP, where L₂-normalized features would otherwise make the bound class-agnostic; the authors use p = 3 in the CLIP experiments.

Main Findings

• Large improvements on hard classes on CIFAR-100 (ResNet-32). MR² reaches 73.9 overall, 85.9 easy, 73.8 medium, and 61.9 hard, versus ERM at 70.9, 84.5, 71.0, and 56.7. The hard-class gain is +5.2 points.

• Large hard-class gains on ImageNet (CLIP ResNet-50). MR² reaches 76.9 overall, 91.5 easy, 79.7 medium, 59.6 hard, versus ERM at 75.2, 91.1, 78.3, 56.4 — a +3.2 point hard-class gain.

• Gains hold across architectures and pre-trained backbones. On CIFAR-100, ResNet-20 improves 68.7 → 70.9, PreAct RN-20 69.1 → 71.3, PreAct RN-32 71.2 → 73.3, and WideResNet-22-10 78.4 → 81.2. On ImageNet, ResNet-50 trained from scratch improves 71.7 → 74.2 (hard 52.6 → 55.9), MoCov2 ResNet-50 71.1 → 72.4 when using only ℓ_γ,ce, CLIP ViT-B/32 75.6 → 77.1, and MAE ViT-B/16 80.4 → 82.0 (hard 62.7 → 66.1).

• No easy-class trade-off. The authors emphasize that easy-class accuracy also rises slightly in their runs, in contrast to reweighting, margin-based long-tail methods, and distributionally robust optimization, which they say entail performance trade-offs across hard and easy groups.

• Weight norms are already balanced. Figure 3 shows ‖w‖₂ values from a ResNet-50 trained on ImageNet, supporting the claim that the disparity is not caused by classifier bias — which is why the bound uses a uniform weight-norm bound Λ.

• Disparity-linked feature variability. Figure 1(b) shows feature spread increasing from easy to hard classes, consistent with the earlier empirical finding the paper builds on.

Methodology in Plain English

The authors start from a hypothesis that is checked empirically: underperforming classes have more diverse features, not just worse classifiers. They then write down a standard style of generalization bound for a multi-class classifier with per-class margins, and work out what happens to that bound when each class is given its own margin.

The result is a recipe: a class whose features are more spread out should get a larger margin. Applying Lagrange optimization under a fixed average margin budget gives the cube-root formula in Corollary 1. The margin is implemented softly by dividing the logits by a class-specific value γ_y before the softmax, which is equivalent to scaling the logit gaps for each class differently — similar in spirit to a temperature, but chosen per class from feature statistics rather than globally.

Because large feature spread also inflates the bound's complexity term, the authors add a second loss that pulls same-class examples together. This loss compares each example to others in its class and penalizes squared distances that exceed twice the dataset-average mean squared deviation, a differentiable relaxation of a hinge-like max penalty. Both class statistics (‖μ̂_k‖₂² and ‖ŝ_k‖₂²) are tracked with an exponential moving average during training.

The combined objective is the γ-margin cross-entropy plus λ times the representation margin loss, trained with standard augmentation, with c̄ chosen from {1, 2, 3}, λ from {0.1, 0.3, 0.5, 0.7, 0.9}, EMA decay 0.9, and p = 3 for cosine-classifier models. Experiments cover seven datasets (CIFAR-100, ImageNet, StanfordCars, OxfordPets, Flowers, Food, FGVCAircraft), CNN and ViT backbones, train-from-scratch and fine-tuning/linear-probing settings, against 14 baselines, with all runs repeated three times.

Why This Matters

The paper converts an empirically observed fairness problem into a theory-driven fix. Instead of treating class-wise imbalance as a data or optimizer artifact, it ties the disparity to a measurable quantity (per-class feature variance) and gives a formula for how much margin each class deserves. That matters because most prior remedies in this space were described by the authors as largely empirical.

Real-world applications where uneven per-class accuracy is costly:

• Safety-critical perception, such as autonomous driving, where a rarely-seen but important object category must not be the worst-performing class.
• Medical imaging, where a diagnostic model that is accurate on average but poor on a specific condition can cause harm.
• Content moderation and security, where certain categories of harmful content may fall into the "hard" tail.
• Large-scale visual recognition and e-commerce tagging built on foundation models such as CLIP or MAE, where the method can be added to existing fine-tuning pipelines.

Industry relevance: Because MR² is a modification to the training loss rather than a new architecture, it can be layered onto existing pre-trained backbones (MAE, MoCov2, CLIP) with either full fine-tuning or linear probing, and the authors provide code. That makes it a relatively low-cost intervention for teams already fine-tuning foundation models and needing to report per-class rather than only average accuracy.

Future Directions

• Choosing p systematically. The paper shows that the L₂ bound degenerates for cosine classifiers and adopts p = 3 for CLIP, but how to select the best p in general is left open.
• Interaction with advanced augmentation. The paper mentions evaluating MR² under RandAugment and AutoAug with results in Table 5, which is not included in the truncated content; a fuller picture of how augmentation and margin regularization interact is not reported here.
• Per-class accuracy without partitioning. Section C.3 is said to report per-class accuracies without the easy/medium/hard split, but those numbers are not in the content provided.
• Sensitivity and variance. Standard deviations are reported in Tab. 15(b), but the content does not report them, so the stability of the gains across seeds cannot be assessed from this material.
• Extension beyond classification accuracy disparity. The framework is developed for balanced-prior classification; whether the same margin-from-variance logic transfers to detection, segmentation, or domain-shift settings is not reported.

Target Audience

Researchers and practitioners working on class-wise fairness, robust or balanced classification, and margin-based losses will get the most from this paper. It suits readers who have some familiarity with generalization bounds and Rademacher complexity, since the core argument is theoretical. Practitioners fine-tuning vision foundation models on balanced datasets with uneven per-class accuracy — and engineers who need a drop-in loss-level fix rather than a new architecture — are the most direct beneficiaries.

Authors’ 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