[Paper Review] Learning Imbalanced Datasets with Label-Distribution-Aware Margin Loss

Kaidi Cao, Colin Wei+3 more|arXiv (Cornell University)|Jun 18, 2019

Imbalanced Data Classification Techniques61 references226 citations

TL;DR

The paper introduces LDAM loss to enforce class-dependent margins inverse to class frequencies and a deferred re-balancing schedule to improve generalization on minority classes in imbalanced datasets, showing strong gains across CIFAR, iNaturalist, and IMDB.

ABSTRACT

Deep learning algorithms can fare poorly when the training dataset suffers from heavy class-imbalance but the testing criterion requires good generalization on less frequent classes. We design two novel methods to improve performance in such scenarios. First, we propose a theoretically-principled label-distribution-aware margin (LDAM) loss motivated by minimizing a margin-based generalization bound. This loss replaces the standard cross-entropy objective during training and can be applied with prior strategies for training with class-imbalance such as re-weighting or re-sampling. Second, we propose a simple, yet effective, training schedule that defers re-weighting until after the initial stage, allowing the model to learn an initial representation while avoiding some of the complications associated with re-weighting or re-sampling. We test our methods on several benchmark vision tasks including the real-world imbalanced dataset iNaturalist 2018. Our experiments show that either of these methods alone can already improve over existing techniques and their combination achieves even better performance gains.

Motivation & Objective

• Motivate and address poor generalization on minority classes in imbalanced datasets.
• Derive a per-class margin theory to guide loss design for minority classes.
• Propose a practical LDAM loss that increases margins for rare classes.
• Introduce a deferred re-balancing training schedule to stabilize optimization with re-weighting.
• Demonstrate empirical gains on benchmark vision tasks and real-world data.

Proposed method

• Derive per-class margin bounds showing optimal margins scale as n_j^{-1/4} for class j.
• Define LDAM loss that enforces class-dependent margins with Delta_j = C / n_j^{1/4} for each class j.
• Normalize logits and final-layer weights to unit length to stabilize margin adjustments.
• Propose a deferred re-balancing optimization: train with LDAM/ERM in stage one, then apply re-weighted LDAM with reduced learning rate in stage two.
• Combine LDAM with re-weighting or re-sampling schemes (DRW/DRS) for enhanced performance (LDAM-DRW as main pipeline).
• Evaluate on imbalanced versions of IMDB, CIFAR-10/100, Tiny ImageNet, and real-world iNaturalist 2018 datasets.

Experimental results

Research questions

• RQ1 How should margins be allocated across classes to optimize generalization under imbalanced label distributions?
• RQ2 Can a label-distribution-aware margin loss improve minority-class performance without harming frequent classes?
• RQ3 Does deferring re-weighting (DRW) lead to better feature representations and later gains when re-weighting is applied?
• RQ4 What empirical gains do LDAM and DRW yield on standard imbalanced benchmarks and real-world data?
• RQ5 How does LDAM interact with existing re-weighting and re-sampling techniques?

Key findings

• LDAM loss improves minority-class generalization by enforcing larger margins for rare classes (Delta_j ∝ n_j^{-1/4}).
• Deferred re-balancing (DRW) training stabilizes optimization and yields strong gains when followed by re-weighted LDAM.
• LDAM-DRW outperforms ERM, re-weighting, re-sampling, and focal-loss baselines across IMDB, CIFAR-10/100, and iNaturalist 2018.
• Independently, LDAM and DRW provide significant improvements; together they yield the best results (LDAM-DRW).
• Experiments on artificially imbalanced CIFAR variants and real-world iNaturalist 2018 show notable top-1 error reductions compared to prior state-of-the-art methods.