PaperlayerView exact revision on arXiv

Source-linked AI summary

Learning Imbalanced Datasets with Label-Distribution-Aware Margin Loss

Kaidi Cao, Colin Wei, Adrien Gaidon, Nikos Arechiga, Tengyu Ma

arXiv:1906.07413v2cs.LGcs.CVstat.ML

TL;DR

The paper addresses poor minority-class generalization under heavily imbalanced training data when evaluation emphasizes less frequent classes. It introduces LDAM, a theoretically motivated class-dependent margin loss, and deferred re-weighting, then reports improved performance across benchmark vision tasks, with the combination performing best. The methods assume knowledge of the test label distribution and analyze separable training cases.

  • Problem

    Deep networks can perform poorly on underrepresented classes when training labels are imbalanced but evaluation emphasizes minority classes.

  • Method

    The paper combines label-distribution-aware margin loss, which encourages larger minority-class margins, with a deferred re-balancing schedule that postpones re-weighting.

  • Results

    The two methods significantly improve performance across benchmark vision tasks, and their practical combination improves artificially imbalanced CIFAR, Tiny ImageNet, and iNaturalist’18.

  • Takeaways & Limitations

    LDAM provides a theoretically principled uniform-label generalization-bound justification, while deferred re-weighting can avoid drawbacks of re-weighting or re-sampling during initial representation learning.

  • Takeaways & Limitations

    The analysis considers separable training cases, and the precise theoretical explanation for deferred re-weighting’s success remains unclear.

Abstract

from arXiv · show

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.

1 Introduction

Long-tailed datasets can hurt performance on minority classes when evaluation emphasizes them. The paper proposes stronger minority-class regularization through label-distribution-aware margins and combines it with deferred re-balancing.

  • Long-tailed training data can produce poor performance on underrepresented classes, especially when testing uses uniform-label or minimum-class accuracy.
  • Re-weighting and re-sampling can improve class-accuracy trade-offs but may overfit minority classes because those classes provide less information.
  • The paper regularizes minority classes more strongly than frequent classes to improve their generalization without sacrificing fit on frequent classes.
  • LDAM uses per-class margins motivated by a uniform-label generalization bound, while remaining compatible with re-weighting and re-sampling.
  • Deferred re-balancing postpones re-weighting, and the combined methods improve performance on artificially imbalanced CIFAR and Tiny ImageNet and real-world iNaturalist’18.

2 Related Works

Prior work addresses imbalance through re-sampling, re-weighting, domain-adaptation methods, meta-learning, and margin-based objectives. This paper instead adds class-dependent margin regularization and defers re-balancing to improve optimization and minority-class generalization.

  • Re-sampling: Re-sampling either discards frequent-class data or repeatedly exposes minority examples, with under-sampling becoming infeasible under extreme imbalance and over-sampling risking minority overfitting.
  • Re-weighting: Re-weighting can make deep-model optimization difficult under extreme imbalance and may reduce performance on frequent classes when inverse class frequency is used.
  • Re-weighting: Focal loss and related methods adjust weights using individual-example properties such as classification difficulty or gradient magnitude.
  • Novelty: The paper’s approach adds minority-focused regularization orthogonal to re-weighting and proposes deferred re-balancing to improve a generic re-weighting scheme.
  • Margin loss: Unlike class-independent angular-margin methods, the paper uses class-distribution-aware margins to target minority-class generalization.
  • Related settings: The paper assumes the test label distribution is known, unlike distributionally robust formulations that only bound the amount of label shift.
  • Meta-learning: The authors regard loss-based methods as generally more computationally efficient than meta-learning approaches.

3 Main Approach

The approach derives class-specific margin targets from balanced-test generalization bounds, then implements them with LDAM and deferred re-balancing. LDAM enlarges minority-class margins, while deferred re-balancing applies re-weighting only after initial ERM training.

  • Problem setup and notations: The analysis assumes class-conditional distributions are unchanged between training and testing and evaluates error under a balanced test distribution.The balanced distribution samples classes uniformly, then samples inputs from each class-conditional distribution.
  • Theoretical motivations: Per-class margin bounds motivate larger margins for minority classes, whose smaller sample counts make generalization more difficult.The paper studies per-class and balanced-test error bounds rather than only the minimum margin across all examples.
  • Label-distribution-aware margin loss: The proposed class-dependent margins scale as ∆j = C n_j^-1/4, assigning larger margins to classes with fewer training examples.This trade-off is derived from the binary-class analysis and extended to multiple classes.
  • Label-distribution-aware margin loss: LDAM extends soft-margin classification by enforcing label-distribution-dependent margins and is orthogonal to re-weighting and re-sampling.The implementation uses a smooth cross-entropy relaxation of the multiclass hinge loss, with normalized logits to simplify margin tuning.
  • Deferred Re-balancing Optimization Schedule: Deferred re-balancing first trains with vanilla ERM and LDAM, then applies re-weighted LDAM after learning-rate annealing at a smaller learning rate.The first stage is intended to provide a good initialization, while the second stage changes the weights relatively little; this addresses optimization and over-fitting issues associated with re-balancing.

4 Experiments

Experiments across imbalanced text and vision benchmarks evaluate LDAM, deferred re-balancing, and their combination. LDAM-DRW achieves the strongest reported performance, including substantial gains on iNaturalist 2018, while analyses show improved minority-class generalization and benefits from deferred re-balancing.

  • Datasets: Experiments cover artificially imbalanced IMDB, CIFAR-10, CIFAR-100, and Tiny ImageNet, plus real-world iNaturalist 2018.The evaluation uses controllable imbalance settings for the artificial datasets and a large-scale naturally imbalanced dataset.
  • Methods: LDAM-DRW combines label-distribution-aware margins with deferred re-weighting, while DRW and DRS apply re-balancing only after the initial ERM stage.LDAM is also evaluated with other losses and schedules, and LDAM-DRW is the main proposed algorithm.
  • CIFAR results: LDAM-DRW achieves the best top-1 validation-error performance on imbalanced CIFAR-10 and CIFAR-100, with each component beneficial in combinations with other losses or schedules.Table 2 compares ResNet-32 models under long-tailed and step imbalance settings.
  • CIFAR results: LDAM is superior to cross-entropy and focal loss without a data-rebalance schedule, while the full pipeline outperforms previous state-of-the-art methods by a large margin.Uniform-margin alternatives and hinge loss are included in the ablation analysis; hinge loss has optimization issues with 100 classes.
  • Analysis: On CIFAR-10 with step imbalance, the proposed algorithm generalizes better on minority classes while leaving frequent-class performance almost unaffected; under long-tailed imbalance, DRW becomes superior after learning-rate annealing.Figure 2 contrasts RW under-fitting and RS minority over-fitting, while Figure 3 marks annealing at epoch 160.
  • iNaturalist 2018: 10.86% and 6.88% are the reported improvements in iNaturalist 2018 top-1 error over ERM and the previous state of the art, respectively.The experiment uses 437,513 training images, 8,142 classes, balanced validation labels, and a ResNet-50 backbone.

5 Conclusion

The paper introduces LDAM and deferred re-weighting for imbalanced datasets, combines theoretical justification with empirical evaluation, and reports improved performance across benchmark vision tasks.

  • 5 Conclusion: LDAM encourages larger margins for minority classes and is theoretically justified through a uniform-label generalization error bound.The analysis motivates class-dependent margins from per-class generalization bounds.
  • 5 Conclusion: Deferred re-weighting postpones re-weighting until after the model learns an initial representation.The authors state that this may avoid drawbacks of re-weighting or re-sampling during early training, although its precise explanation remains unclear.
  • 5 Conclusion: LDAM and deferred re-weighting achieve significantly improved performance across varied imbalanced-dataset vision tasks.The conclusion reports broad benchmark improvements while identifying the theoretical explanation of deferred re-weighting as an open direction.

A Missing Proofs and Derivations in Section 3.1

The derivation bounds balanced-class generalization using class-dependent margins and then analyzes how margin shifts can optimize the binary-class bound.

  • A Missing Proofs and Derivations in Section 3.1: The proof defines class-specific hard-margin losses and empirical Rademacher complexities to control generalization from each class’s samples.The class-j sample count enters the complexity-based bound through the corresponding margin analysis.
  • A Missing Proofs and Derivations in Section 3.1: Theorem 2 bounds balanced-class generalization for all positive class-dependent margins and hypotheses in the function class.The proof applies a classwise margin bound and combines the class-specific results by a union bound and averaging.
  • A Missing Proofs and Derivations in Section 3.1: For binary classification, adding bias terms can shift the two class margins while preserving their total and reach the bound’s optimal margin allocation.The construction starts from a classifier with positive margins, selects bias terms, and obtains the optimized margins before applying Theorem 2.
  • A Missing Proofs and Derivations in Section 3.1: Substituting γ2 = β − γ1 and setting the derivative to zero yields the margin allocation that optimizes the two-class bound.The resulting optimized margins are then used to obtain the stated generalization guarantee.

B Implementation details

The experiments use standard convolutional architectures, augmentation, optimization settings, and dataset-specific schedules across artificially imbalanced CIFAR, Tiny ImageNet, and iNaturalist 2018.

  • B Implementation details: Artificial imbalance distributions include long-tailed and step configurations for CIFAR datasets.Figure 4 illustrates the number of training examples per class for these generated distributions.
  • B Implementation details: CIFAR experiments train ResNet-32 for 200 epochs with batch size 128, momentum 0.9, and weight decay 2 × 10−4.Training uses four-pixel padding, random crops, and horizontal flips.
  • B Implementation details: Tiny ImageNet experiments use ResNet-18, 64 × 64 crops, batch size 128, and a 120-epoch training schedule.The passage specifies horizontal flips, eight-pixel padding, and an initial learning rate of 0.1.
  • B Implementation details: iNaturalist 2018 experiments train ResNet-50 for 90 epochs and begin rebalancing at epoch 60 in the two-stage schedule.Images are resized to a 256-pixel shorter side, cropped to 224 × 224, and evaluated with margins normalized so the largest enforced margin is 0.3.
  • B Implementation details: Table 4 reports validation error on imbalanced Tiny ImageNet across different loss functions and training schedules.The supplied table caption identifies the comparison structure but does not provide numerical entries.

C.1 Feature visualization

Feature visualizations qualitatively compare methods on imbalanced data, while Tiny ImageNet results provide an additional benchmark showing consistent improvements from LDAM and deferred re-weighting.

  • C.1 Feature visualization: The feature visualization constrains representations to three normalized dimensions, with training distributions in the top row and validation distributions in the second row.The setup enables scattering features on a unit hyper-sphere for visual comparison.
  • C.1 Feature visualization: Tiny ImageNet has 200 classes with 500 training images and 50 validation images per class, and is evaluated under generated long-tailed and step imbalance.The results are presented in Table 4.
  • C.1 Feature visualization: LDAM and deferred re-weighting show consistent improvements over ERM on imbalanced Tiny ImageNet, whereas Class-Balanced Softmax performs worse than ERM.This comparison extends evaluation beyond artificially imbalanced CIFAR.

C.3 Comparing feature extractors trained by different schemes

Feature extractors from ERM-trained backbones are most expressive in the highly imbalanced setting examined. Deferred re-balancing can be combined with either re-weighting or re-sampling, which perform similarly in that schedule.

  • Feature extractor comparison: Under step imbalance with ρ = 100 and µ = 0.5, ERM-trained backbone networks produce the most expressive feature embedding among the three methods.This comparison uses a linear classifier trained on features extracted from differently pretrained backbones.
  • Deferred re-balancing: Re-weighting and re-sampling exhibit similar performance when combined with deferred re-balancing.The authors use re-weighting as the default choice in the main paper.

C.5 Imbalanced Test Label Distributions

The extended algorithm adapts LDAM to imbalanced test distributions and is evaluated when training and validation label distributions differ. Across two random test-distribution seeds, the results are qualitatively similar and the algorithm works well in this setting.

  • LDAM adaptation: The LDAM extension adapts the class-specific margin ∆i using the training and test class counts ni and n′i.The cited passages introduce the relevant counts and state that LDAM is adapted by changing the margin for class i.
  • Test-distribution construction: Test labels are generated with an imbalanced distribution and their frequencies are randomly permuted so the training and test distributions differ substantially.For the long-tailed setting with ρ = 100, the training distribution follows Figure 7a, while test distributions are shown in Figure 7b and Figure 7c.
  • Robustness across test distributions: Across two random seeds for generating test label distributions, the reported settings show qualitatively similar results.The precise generated label distributions are provided in the authors’ code.
  • Experimental setting: The extended algorithm is evaluated with both imbalanced training and imbalanced validation distributions for CIFAR-10.Table 6 reports Top-1 validation error for this scheme.
Loading 1906.07413v2…